Cache airline flight DB and harden adapters

Author: benjibc

eval_protocol/adapters/langfuse.py

@@ -9,7 +9,7 @@
 import random
 import time
 from datetime import datetime, timedelta
-from typing import Any, Dict, List, Optional, Protocol
+from typing import Any, Callable, Dict, List, Optional, Protocol, cast
 
 from eval_protocol.models import EvaluationRow, InputMetadata, Message
 from .base import BaseAdapter
@@ -44,14 +44,19 @@ def __call__(
         ...
 
 
+LangfuseClient = Any
+
+_get_langfuse_client: Callable[[], Any] | None
+
 try:
-    from langfuse import get_client  # pyright: ignore[reportPrivateImportUsage]
+    from langfuse import get_client as _get_langfuse_client  # type: ignore[attr-defined, reportPrivateImportUsage]
     from langfuse.api.resources.trace.types.traces import Traces
     from langfuse.api.resources.commons.types.trace import Trace
     from langfuse.api.resources.commons.types.trace_with_full_details import TraceWithFullDetails
 
     LANGFUSE_AVAILABLE = True
-except ImportError:
+except ImportError:  # pragma: no cover - optional dependency
+    _get_langfuse_client = None
     LANGFUSE_AVAILABLE = False
 
 
@@ -219,7 +224,11 @@ def __init__(self):
         if not LANGFUSE_AVAILABLE:
             raise ImportError("Langfuse not installed. Install with: pip install 'eval-protocol[langfuse]'")
 
-        self.client = get_client()
+        if _get_langfuse_client is None:
+            raise ImportError("Langfuse not installed. Install with: pip install 'eval-protocol[langfuse]'")
+
+        client_factory = cast(Callable[[], LangfuseClient], _get_langfuse_client)
+        self.client: LangfuseClient = client_factory()
 
     def get_evaluation_rows(
         self,

eval_protocol/adapters/langsmith.py

@@ -10,18 +10,23 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List, Optional, Iterable
+from typing import Any, Callable, Dict, Iterable, List, Optional, cast
 
 from eval_protocol.models import EvaluationRow, InputMetadata, Message
 from .base import BaseAdapter
 
 logger = logging.getLogger(__name__)
 
+LangSmithClient = Any
+
+_LANGSMITH_CLIENT_CTOR: Callable[..., LangSmithClient] | None
+
 try:
-    from langsmith import Client  # type: ignore
+    from langsmith import Client as _LANGSMITH_CLIENT_CTOR  # type: ignore[attr-defined]
 
     LANGSMITH_AVAILABLE = True
-except ImportError:
+except ImportError:  # pragma: no cover - optional dependency
+    _LANGSMITH_CLIENT_CTOR = None
     LANGSMITH_AVAILABLE = False
 
 
@@ -35,10 +40,17 @@ class LangSmithAdapter(BaseAdapter):
     - outputs: { messages: [...] } | { content } | { result } | { answer } | { output } | str | list[dict]
     """
 
-    def __init__(self, client: Optional[Client] = None) -> None:
+    def __init__(self, client: Optional[LangSmithClient] = None) -> None:
         if not LANGSMITH_AVAILABLE:
             raise ImportError("LangSmith not installed. Install with: pip install 'eval-protocol[langsmith]'")
-        self.client = client or Client()
+        if client is not None:
+            self.client = client
+            return
+
+        if _LANGSMITH_CLIENT_CTOR is None:
+            raise ImportError("LangSmith client constructor unavailable despite successful import check")
+
+        self.client: LangSmithClient = cast(LangSmithClient, _LANGSMITH_CLIENT_CTOR())
 
     def get_evaluation_rows(
         self,

eval_protocol/mcp_servers/tau2/airplane_environment/airline_environment.py

@@ -9,8 +9,7 @@
 import json
 import logging
 import os
-import time
-from copy import deepcopy
+from functools import lru_cache
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
@@ -24,6 +23,16 @@
 from vendor.tau2.domains.airline.utils import AIRLINE_DB_PATH
 
 
+@lru_cache(maxsize=1)
+def _load_flight_db(path: str) -> FlightDB:
+    """Load and cache the flight database for reuse across resets."""
+
+    logger.info("🗂️ Loading airline database from disk (cached)")
+    db_loaded = FlightDB.load(path)
+    assert isinstance(db_loaded, FlightDB)
+    return db_loaded
+
+
 class AirlineEnvironment:
     """
     Airline environment that integrates τ²-Bench simulation pattern
@@ -37,13 +46,10 @@ def __init__(self, config: Optional[Dict[str, Any]] = None):
 
     def reset(self, seed: Optional[int] = None) -> Tuple[Dict[str, Any], Dict[str, Any]]:
         """Reset the environment to initial state"""
-        logger.info("🔄 Resetting airline environment - reloading database from disk")
-        # FlightDB.load expects a str path
-        # Ensure type matches expected FlightDB
-        # FlightDB.load returns vendor.tau2.domains.airline.data_model.FlightDB which is compatible
-        db_loaded = FlightDB.load(str(AIRLINE_DB_PATH))
-        assert isinstance(db_loaded, FlightDB)
-        self.db = db_loaded
+        logger.info("🔄 Resetting airline environment - using cached airline database")
+        cached_db = _load_flight_db(str(AIRLINE_DB_PATH))
+        # Provide a fresh copy for each environment reset without re-reading from disk.
+        self.db = cached_db.model_copy(deep=True)
         self.airline_tools = AirlineTools(self.db)
 
         return {}, {}

tests/pytest/test_mcp_session_autocreate.py

@@ -31,18 +31,24 @@ async def test_tool_call_returns_json_without_prior_initial_state():
     try:
         base_url = "http://127.0.0.1:9780/mcp"
         client = httpx.Client(timeout=1.0)
-        deadline = time.time() + 20
+        start_time = time.time()
+        deadline = start_time + 20
+        ready_time = None
         while time.time() < deadline:
             try:
                 r = client.get(base_url)
                 if r.status_code in (200, 307, 406):
+                    ready_time = time.time()
                     break
             except Exception:
                 pass
             time.sleep(0.2)
         else:
             pytest.fail("Server did not start on port 9780 in time")
 
+        assert ready_time is not None, "Server did not return a successful status before exiting loop"
+        assert ready_time - start_time < 20, f"Server took too long to respond: {ready_time - start_time:.2f}s"
+
         session = MCPSession(base_url=base_url, session_id="test-autocreate", seed=None, model_id="test-model")
 
         mgr = MCPConnectionManager()