braintrust 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

braintrust/devserver/schemas.py

@@ -1,7 +1,8 @@
 import json
-from typing import Any, Dict, List, Optional, Sequence, Union, get_args, get_origin
+from collections.abc import Sequence
+from typing import Any, Union, get_args, get_origin, get_type_hints
 
-from typing_extensions import TypedDict, get_type_hints
+from typing_extensions import TypedDict
 
 # This is not beautiful code, but it saves us from introducing Pydantic as a dependency, and it is fairly
 # straightforward for an LLM to keep it up to date with runEvalBodySchema in JS.
@@ -16,12 +17,12 @@ class ValidationError(Exception):
 class ParsedFunctionId(TypedDict, total=False):
     """Parsed function identifier."""
 
-    function_id: Optional[str]
-    version: Optional[str]
-    name: Optional[str]
-    prompt_session_id: Optional[str]
-    inline_code: Optional[str]
-    global_function: Optional[str]
+    function_id: str | None
+    version: str | None
+    name: str | None
+    prompt_session_id: str | None
+    inline_code: str | None
+    global_function: str | None
 
 
 class ParsedParent(TypedDict):
@@ -35,16 +36,16 @@ class ParsedEvalBody(TypedDict, total=False):
     """Type for parsed eval request body."""
 
     name: str  # Required
-    parameters: Dict[str, Any]
+    parameters: dict[str, Any]
     data: Any
-    scores: List[ParsedFunctionId]
+    scores: list[ParsedFunctionId]
     experiment_name: str
     project_id: str
-    parent: Union[str, ParsedParent]
+    parent: str | ParsedParent
     stream: bool
 
 
-def validate_typed_dict(data: Any, typed_dict_class: type, path: str = "") -> Dict[str, Any]:
+def validate_typed_dict(data: Any, typed_dict_class: type, path: str = "") -> dict[str, Any]:
     """Validate data against a TypedDict definition."""
     if not isinstance(data, dict):
         raise ValidationError(f"{path or 'Root'} must be a dictionary, got {type(data).__name__}")
@@ -107,7 +108,7 @@ def validate_value(value: Any, expected_type: type, path: str) -> Any:
         return validate_value(value, inner_type, path)
 
     # Handle List/Sequence
-    if origin in (list, List, Sequence):
+    if origin in (list, list, Sequence):
         if not isinstance(value, list):
             raise ValidationError(f"{path} must be a list, got {type(value).__name__}")
 
@@ -115,7 +116,7 @@ def validate_value(value: Any, expected_type: type, path: str) -> Any:
         return [validate_value(item, item_type, f"{path}[{i}]") for i, item in enumerate(value)]
 
     # Handle Dict/Mapping
-    if origin in (dict, Dict):
+    if origin in (dict, dict):
         if not isinstance(value, dict):
             raise ValidationError(f"{path} must be a dict, got {type(value).__name__}")
 
@@ -172,7 +173,7 @@ def parse_function_id(data: Any, path: str = "function") -> ParsedFunctionId:
     raise ValidationError(f"{path} must specify function_id, name, prompt_session_id, or inline_code")
 
 
-def parse_eval_body(request_data: Union[str, bytes, dict]) -> ParsedEvalBody:
+def parse_eval_body(request_data: str | bytes | dict) -> ParsedEvalBody:
     """
     Parse request body for eval execution.
 
@@ -221,10 +222,12 @@ def parse_eval_body(request_data: Union[str, bytes, dict]) -> ParsedEvalBody:
         parsed_scores = []
         for i, score in enumerate(scores_data):
             try:
-                parsed_scores.append({
-                    "name": score["name"],
-                    "function_id": parse_function_id(score["function_id"], f"scores[{i}]"),
-                })
+                parsed_scores.append(
+                    {
+                        "name": score["name"],
+                        "function_id": parse_function_id(score["function_id"], f"scores[{i}]"),
+                    }
+                )
             except ValidationError as e:
                 raise ValidationError(f"Invalid score at index {i}: {e}")
 

braintrust/devserver/server.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 import sys
 import textwrap
-from typing import Any, Optional, Union
+from typing import Any
 
 try:
     import uvicorn
@@ -40,7 +40,7 @@ _all_evaluators: dict[str, Evaluator[Any, Any]] = {}
 
 
 class CheckAuthorizedMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, allowed_org_name: Optional[str] = None):
+    def __init__(self, app, allowed_org_name: str | None = None):
         super().__init__(app)
         self.allowed_org_name = allowed_org_name
         self.protected_paths = ["/list", "/eval"]
@@ -100,7 +100,7 @@ async def list_evaluators(request: Request) -> JSONResponse:
     return JSONResponse(evaluator_list)
 
 
-async def run_eval(request: Request) -> Union[JSONResponse, StreamingResponse]:
+async def run_eval(request: Request) -> JSONResponse | StreamingResponse:
     """Handle eval execution requests."""
     try:
         # Get request body
@@ -157,12 +157,14 @@ async def run_eval(request: Request) -> Union[JSONResponse, StreamingResponse]:
             result = await evaluator.task(input, hooks)
         else:
             result = evaluator.task(input, hooks)
-        hooks.report_progress({
-            "format": "code",
-            "output_type": "completion",
-            "event": "json_delta",
-            "data": json.dumps(result),
-        })
+        hooks.report_progress(
+            {
+                "format": "code",
+                "output_type": "completion",
+                "event": "json_delta",
+                "data": json.dumps(result),
+            }
+        )
         return result
 
     def on_start_fn(summary: ExperimentSummary):
@@ -214,6 +216,7 @@ async def run_eval(request: Request) -> Union[JSONResponse, StreamingResponse]:
 
             async def event_generator():
                 """Generate SSE events from the queue."""
+
                 # Create a task to run the eval and signal completion
                 async def run_and_complete():
                     try:
@@ -255,7 +258,7 @@ async def run_eval(request: Request) -> Union[JSONResponse, StreamingResponse]:
         return JSONResponse({"error": f"Failed to run evaluation: {str(e)}"}, status_code=500)
 
 
-def create_app(evaluators: list[Evaluator[Any, Any]], org_name: Optional[str] = None):
+def create_app(evaluators: list[Evaluator[Any, Any]], org_name: str | None = None):
     """Create and configure the Starlette app for the dev server.
 
     Args:
@@ -283,7 +286,9 @@ def create_app(evaluators: list[Evaluator[Any, Any]], org_name: Optional[str] =
     return app
 
 
-def run_dev_server(evaluators: list[Evaluator[Any, Any]], host: str = "localhost", port: int = 8300, org_name: Optional[str] = None):
+def run_dev_server(
+    evaluators: list[Evaluator[Any, Any]], host: str = "localhost", port: int = 8300, org_name: str | None = None
+):
     """Start the dev server.
 
     Args:
@@ -305,7 +310,9 @@ def snake_to_camel(snake_str: str) -> str:
     return components[0] + "".join(x.title() for x in components[1:]) if components else snake_str
 
 
-def make_scorer(state: BraintrustState, name: str, score: FunctionId, project_id: Optional[str] = None) -> EvalScorer[Any, Any]:
+def make_scorer(
+    state: BraintrustState, name: str, score: FunctionId, project_id: str | None = None
+) -> EvalScorer[Any, Any]:
     def scorer_fn(input, output, expected, metadata):
         request = {
             **score,

braintrust/devserver/test_cached_login.py

@@ -10,7 +10,7 @@ class TestCachedLogin(unittest.TestCase):
         """Clear the cache before each test."""
         cache._login_cache = cache.LRUCache(max_size=32)
 
-    @patch('braintrust.devserver.cache.login_to_state')
+    @patch("braintrust.devserver.cache.login_to_state")
     def test_cached_login_caches_results(self, mock_login):
         """Test that cached_login caches and reuses results."""
         mock_state = MagicMock()
@@ -26,7 +26,7 @@ class TestCachedLogin(unittest.TestCase):
         self.assertEqual(result2, mock_state)
         self.assertEqual(mock_login.call_count, 1)  # Still 1, not called again
 
-    @patch('braintrust.devserver.cache.login_to_state')
+    @patch("braintrust.devserver.cache.login_to_state")
     def test_cached_login_different_keys(self, mock_login):
         """Test that different cache keys create separate entries."""
         mock_state1 = MagicMock()
@@ -48,7 +48,7 @@ class TestCachedLogin(unittest.TestCase):
         self.assertEqual(result3, mock_state3)
         self.assertEqual(mock_login.call_count, 3)
 
-    @patch('braintrust.devserver.cache.login_to_state')
+    @patch("braintrust.devserver.cache.login_to_state")
     def test_cached_login_with_org_name(self, mock_login):
         """Test caching with org_name parameter."""
         mock_state = MagicMock()
@@ -68,7 +68,7 @@ class TestCachedLogin(unittest.TestCase):
         result3 = asyncio.run(cache.cached_login("api_key_1", "https://app.braintrust.com", org_name="other_org"))
         self.assertEqual(mock_login.call_count, 2)
 
-    @patch('braintrust.devserver.cache.login_to_state')
+    @patch("braintrust.devserver.cache.login_to_state")
     def test_cached_login_propagates_exceptions(self, mock_login):
         """Test that exceptions from login_to_state are propagated."""
         mock_login.side_effect = ValueError("Invalid API key")