comments

Author: xzrderek

eval_protocol/proxy/__init__.py

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

eval_protocol/proxy/proxy_core/__init__.py

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

eval_protocol/proxy/proxy_core/app.py

@@ -3,8 +3,8 @@
 A FastAPI service that sits in front of LiteLLM and extracts metadata from URL paths.
 """
 
-from fastapi import FastAPI, Depends, HTTPException, Request, Query
-from typing import Optional, Callable, Dict, Any, List
+from fastapi import FastAPI, Depends, Request, Query
+from typing import Optional, List
 import os
 import redis
 import logging
@@ -166,8 +166,8 @@ def get_traces_params(
         )
 
     async def require_auth(request: Request) -> None:
-        account_id = auth_provider.validate_and_return_account_id(request)
-        request.state.account_id = account_id
+        account_info = auth_provider.validate_and_return_account_info(request)
+        request.state.account_id = account_info.account_id if account_info else None
         return None
 
     # =====================

eval_protocol/proxy/proxy_core/auth.py

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

eval_protocol/proxy/proxy_core/litellm.py

@@ -8,7 +8,6 @@
 import logging
 from uuid6 import uuid7
 from fastapi import Request, Response, HTTPException
-from typing import Optional
 import redis
 from .redis_utils import register_insertion_id
 from .models import ProxyConfig, ChatParams

eval_protocol/proxy/proxy_core/models.py

@@ -3,12 +3,24 @@
 """
 
 from pydantic import BaseModel
-from typing import Optional, List, Any, Dict, Callable, TypeAlias
-from fastapi import Request, Query
+from typing import Optional, List, Any, Dict, Protocol
+from fastapi import Request
 
 
-ChatRequestHook: TypeAlias = Callable[[Dict[str, Any], Request, "ChatParams"], tuple[Dict[str, Any], "ChatParams"]]
-TracesRequestHook: TypeAlias = Callable[[Request, "TracesParams"], "TracesParams"]
+class AccountInfo(BaseModel):
+    """Account information returned from authentication."""
+
+    account_id: str
+
+
+class ChatRequestHook(Protocol):
+    def __call__(
+        self, data: Dict[str, Any], request: Request, params: "ChatParams"
+    ) -> tuple[Dict[str, Any], "ChatParams"]: ...
+
+
+class TracesRequestHook(Protocol):
+    def __call__(self, request: Request, params: "TracesParams") -> "TracesParams": ...
 
 
 class ChatParams(BaseModel):