Add preprocess hooks

Author: xzrderek

eval_protocol/proxy/README.md

@@ -332,32 +332,6 @@ eval_protocol/proxy/
 └── README.md               # This file
 ```
 
-### Adding Custom Authentication
-
-Extend `AuthProvider` in `auth.py`:
-```python
-from .auth import AuthProvider
-from fastapi import HTTPException, Request
-
-class MyAuthProvider(AuthProvider):
-    def validate(self, request: Request) -> Optional[str]:
-        api_key = None
-        auth_header = request.headers.get("authorization", "")
-        if auth_header.startswith("Bearer "):
-            api_key = auth_header.replace("Bearer ", "").strip()
-        if not api_key:
-            raise HTTPException(status_code=401, detail="Invalid API key")
-        return api_key
-```
-
-Then pass it to `create_app`:
-```python
-from proxy_core import create_app
-from my_auth import MyAuthProvider
-
-app = create_app(auth_provider=MyAuthProvider())
-```
-
 ### Testing
 
 #### Test chat completion:

eval_protocol/proxy/__init__.py

@@ -5,11 +5,13 @@
 Langfuse tracing for distributed evaluation workflows.
 """
 
-from .proxy_core import create_app, AuthProvider, NoAuthProvider, ProxyConfig
+from .proxy_core import create_app, AuthProvider, NoAuthProvider, ProxyConfig, ChatParams, TracesParams
 
 __all__ = [
     "create_app",
     "AuthProvider",
     "NoAuthProvider",
     "ProxyConfig",
+    "ChatParams",
+    "TracesParams",
 ]

eval_protocol/proxy/docker-compose.yml

@@ -49,6 +49,7 @@ services:
       - LOG_LEVEL=INFO
       # Langfuse and secrets
       - SECRETS_PATH=/app/proxy_core/secrets.yaml
+      - LANGFUSE_HOST=${LANGFUSE_HOST:-https://cloud.langfuse.com}
     ports:
       - "4000:4000"  # Main public-facing port
     networks:

eval_protocol/proxy/proxy_core/__init__.py

@@ -1,9 +1,11 @@
-from .models import ProxyConfig
+from .models import ProxyConfig, ChatParams, TracesParams
 from .app import create_app
 from .auth import AuthProvider, NoAuthProvider
 
 __all__ = [
     "ProxyConfig",
+    "ChatParams",
+    "TracesParams",
     "create_app",
     "AuthProvider",
     "NoAuthProvider",

eval_protocol/proxy/proxy_core/app.py

@@ -4,7 +4,7 @@
 """
 
 from fastapi import FastAPI, Depends, HTTPException, Request, Query
-from typing import Optional, List
+from typing import Optional, Callable, Dict, Any, List
 import os
 import redis
 import logging
@@ -13,24 +13,28 @@
 import sys
 from contextlib import asynccontextmanager
 
-from .models import ProxyConfig, LangfuseTracesResponse
+from .models import ProxyConfig, LangfuseTracesResponse, TracesParams, ChatParams, ChatRequestHook, TracesRequestHook
 from .auth import AuthProvider, NoAuthProvider
 from .litellm import handle_chat_completion, proxy_to_litellm
 from .langfuse import fetch_langfuse_traces
 
 # Configure logging before any other imports (so all modules inherit this config)
 log_level = os.getenv("LOG_LEVEL", "INFO").upper()
-logging.basicConfig(
-    level=getattr(logging, log_level),
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    handlers=[logging.StreamHandler(sys.stdout)],
-)
+if not logging.getLogger().hasHandlers():
+    logging.basicConfig(
+        level=getattr(logging, log_level),
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
 
 logger = logging.getLogger(__name__)
 
 
-def build_proxy_config() -> ProxyConfig:
-    """Load environment and secrets, and build ProxyConfig (no Redis)."""
+def build_proxy_config(
+    preprocess_chat_request: Optional[ChatRequestHook] = None,
+    preprocess_traces_request: Optional[TracesRequestHook] = None,
+) -> ProxyConfig:
+    """Load environment and secrets, and build ProxyConfig"""
     # Env
     litellm_url = os.getenv("LITELLM_URL")
     if not litellm_url:
@@ -67,6 +71,8 @@ def build_proxy_config() -> ProxyConfig:
         langfuse_host=langfuse_host,
         langfuse_keys=langfuse_keys,
         default_project_id=default_project_id,
+        preprocess_chat_request=preprocess_chat_request,
+        preprocess_traces_request=preprocess_traces_request,
     )
 
 
@@ -97,12 +103,15 @@ def init_redis() -> redis.Redis:
 
 def create_app(
     auth_provider: AuthProvider = NoAuthProvider(),
+    preprocess_chat_request: Optional[ChatRequestHook] = None,
+    preprocess_traces_request: Optional[TracesRequestHook] = None,
 ) -> FastAPI:
     @asynccontextmanager
     async def lifespan(app: FastAPI):
         # Build runtime on startup
-        app.state.config = build_proxy_config()
+        app.state.config = build_proxy_config(preprocess_chat_request, preprocess_traces_request)
         app.state.redis = init_redis()
+
         try:
             yield
         finally:
@@ -119,8 +128,46 @@ def get_config(request: Request) -> ProxyConfig:
     def get_redis(request: Request) -> redis.Redis:
         return request.app.state.redis
 
+    def get_traces_params(
+        tags: Optional[List[str]] = Query(default=None),
+        project_id: Optional[str] = None,
+        limit: int = 100,
+        sample_size: Optional[int] = None,
+        user_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        name: Optional[str] = None,
+        environment: Optional[str] = None,
+        version: Optional[str] = None,
+        release: Optional[str] = None,
+        fields: Optional[str] = None,
+        hours_back: Optional[int] = None,
+        from_timestamp: Optional[str] = None,
+        to_timestamp: Optional[str] = None,
+        sleep_between_gets: float = 2.5,
+        max_retries: int = 3,
+    ) -> TracesParams:
+        return TracesParams(
+            tags=tags,
+            project_id=project_id,
+            limit=limit,
+            sample_size=sample_size,
+            user_id=user_id,
+            session_id=session_id,
+            name=name,
+            environment=environment,
+            version=version,
+            release=release,
+            fields=fields,
+            hours_back=hours_back,
+            from_timestamp=from_timestamp,
+            to_timestamp=to_timestamp,
+            sleep_between_gets=sleep_between_gets,
+            max_retries=max_retries,
+        )
+
     async def require_auth(request: Request) -> None:
-        auth_provider.validate(request)
+        account_id = auth_provider.validate_and_return_account_id(request)
+        request.state.account_id = account_id
         return None
 
     # =====================
@@ -161,11 +208,9 @@ async def chat_completion_with_full_metadata(
         encoded_base_url: Optional[str] = None,
         config: ProxyConfig = Depends(get_config),
         redis_client: redis.Redis = Depends(get_redis),
+        _: None = Depends(require_auth),
     ):
-        return await handle_chat_completion(
-            config=config,
-            redis_client=redis_client,
-            request=request,
+        params = ChatParams(
             project_id=project_id,
             rollout_id=rollout_id,
             invocation_id=invocation_id,
@@ -174,6 +219,12 @@ async def chat_completion_with_full_metadata(
             row_id=row_id,
             encoded_base_url=encoded_base_url,
         )
+        return await handle_chat_completion(
+            config=config,
+            redis_client=redis_client,
+            request=request,
+            params=params,
+        )
 
     @app.post("/project_id/{project_id}/chat/completions")
     @app.post("/v1/project_id/{project_id}/chat/completions")
@@ -182,12 +233,14 @@ async def chat_completion_with_project_only(
         request: Request,
         config: ProxyConfig = Depends(get_config),
         redis_client: redis.Redis = Depends(get_redis),
+        _: None = Depends(require_auth),
     ):
+        params = ChatParams(project_id=project_id)
         return await handle_chat_completion(
             config=config,
             redis_client=redis_client,
             request=request,
-            project_id=project_id,
+            params=params,
         )
 
     # ===============
@@ -198,45 +251,20 @@ async def chat_completion_with_project_only(
     @app.get("/project_id/{project_id}/traces", response_model=LangfuseTracesResponse)
     @app.get("/v1/project_id/{project_id}/traces", response_model=LangfuseTracesResponse)
     async def get_langfuse_traces(
-        tags: List[str] = Query(...),  # REQUIRED query param
+        request: Request,
+        params: TracesParams = Depends(get_traces_params),
         project_id: Optional[str] = None,
-        limit: int = 100,
-        sample_size: Optional[int] = None,
-        user_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-        name: Optional[str] = None,
-        environment: Optional[str] = None,
-        version: Optional[str] = None,
-        release: Optional[str] = None,
-        fields: Optional[str] = None,
-        hours_back: Optional[int] = None,
-        from_timestamp: Optional[str] = None,
-        to_timestamp: Optional[str] = None,
-        sleep_between_gets: float = 2.5,
-        max_retries: int = 3,
         config: ProxyConfig = Depends(get_config),
         redis_client: redis.Redis = Depends(get_redis),
         _: None = Depends(require_auth),
     ) -> LangfuseTracesResponse:
+        if project_id is not None:
+            params.project_id = project_id
         return await fetch_langfuse_traces(
             config=config,
             redis_client=redis_client,
-            tags=tags,
-            project_id=project_id,
-            limit=limit,
-            sample_size=sample_size,
-            user_id=user_id,
-            session_id=session_id,
-            name=name,
-            environment=environment,
-            version=version,
-            release=release,
-            fields=fields,
-            hours_back=hours_back,
-            from_timestamp=from_timestamp,
-            to_timestamp=to_timestamp,
-            sleep_between_gets=sleep_between_gets,
-            max_retries=max_retries,
+            request=request,
+            params=params,
         )
 
     # Health

eval_protocol/proxy/proxy_core/auth.py

@@ -1,13 +1,18 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+import logging
 from fastapi import Request
+from fastapi import HTTPException
+import httpx
+from typing import Optional
+
+logger = logging.getLogger(__name__)
 
 
 class AuthProvider(ABC):
     @abstractmethod
-    def validate(self, request: Request) -> Optional[str]: ...
+    def validate_and_return_account_id(self, request: Request) -> Optional[str]: ...
 
 
 class NoAuthProvider(AuthProvider):
-    def validate(self, request: Request) -> Optional[str]:
+    def validate_and_return_account_id(self, request: Request) -> Optional[str]:
         return None

eval_protocol/proxy/proxy_core/langfuse.py

@@ -8,16 +8,18 @@
 import asyncio
 from typing import List, Optional, Dict, Any, Set
 from datetime import datetime, timedelta
-from fastapi import HTTPException
+from fastapi import HTTPException, Request
 import redis
 from .redis_utils import get_insertion_ids
-from .models import ProxyConfig, LangfuseTracesResponse, TraceResponse
+from .models import ProxyConfig, LangfuseTracesResponse, TraceResponse, TracesParams
 
 logger = logging.getLogger(__name__)
 
 
-def _extract_tag_value(tags: List[str], prefix: str) -> Optional[str]:
+def _extract_tag_value(tags: Optional[List[str]], prefix: str) -> Optional[str]:
     """Extract value from a tag with the given prefix (e.g., 'rollout_id:' or 'insertion_id:')."""
+    if not tags:
+        return None
     for tag in tags:
         if tag.startswith(prefix):
             return tag.split(":", 1)[1]
@@ -60,7 +62,7 @@ async def _fetch_trace_list_with_retry(
     langfuse_client: Any,
     page: int,
     limit: int,
-    tags: List[str],
+    tags: Optional[List[str]],
     user_id: Optional[str],
     session_id: Optional[str],
     name: Optional[str],
@@ -152,22 +154,8 @@ async def _fetch_trace_detail_with_retry(
 async def fetch_langfuse_traces(
     config: ProxyConfig,
     redis_client: redis.Redis,
-    tags: List[str],
-    project_id: Optional[str] = None,
-    limit: int = 100,
-    sample_size: Optional[int] = None,
-    user_id: Optional[str] = None,
-    session_id: Optional[str] = None,
-    name: Optional[str] = None,
-    environment: Optional[str] = None,
-    version: Optional[str] = None,
-    release: Optional[str] = None,
-    fields: Optional[str] = None,
-    hours_back: Optional[int] = None,
-    from_timestamp: Optional[str] = None,
-    to_timestamp: Optional[str] = None,
-    sleep_between_gets: float = 2.5,
-    max_retries: int = 3,
+    request: Request,
+    params: TracesParams,
 ):
     """
     Fetch full traces from Langfuse for the specified project.
@@ -184,9 +172,27 @@ async def fetch_langfuse_traces(
 
     Returns a list of full trace objects (including observations) in JSON format.
     """
-    # Validate tags
-    if not tags or not any(tag.startswith("rollout_id:") for tag in tags):
-        raise HTTPException(status_code=422, detail="Tags must include at least one 'rollout_id:*' tag")
+
+    # Preprocess traces request
+    if config.preprocess_traces_request:
+        params = config.preprocess_traces_request(request, params)
+
+    tags = params.tags
+    project_id = params.project_id
+    limit = params.limit
+    sample_size = params.sample_size
+    user_id = params.user_id
+    session_id = params.session_id
+    name = params.name
+    environment = params.environment
+    version = params.version
+    release = params.release
+    fields = params.fields
+    hours_back = params.hours_back
+    from_timestamp = params.from_timestamp
+    to_timestamp = params.to_timestamp
+    sleep_between_gets = params.sleep_between_gets
+    max_retries = params.max_retries
 
     # Use default project if not specified
     if project_id is None: