feat: add shared dataclass for calls so they can be handled by same handler

Author: andrewklatzke

packages/optimization/src/ldai_optimization/__init__.py

@@ -10,6 +10,7 @@
     AIJudgeCallConfig,
     GroundTruthOptimizationOptions,
     GroundTruthSample,
+    LLMCallConfig,
     OptimizationContext,
     OptimizationFromConfigOptions,
     OptimizationJudge,
@@ -28,6 +29,7 @@
     'GroundTruthOptimizationOptions',
     'GroundTruthSample',
     'LDApiError',
+    'LLMCallConfig',
     'OptimizationClient',
     'OptimizationContext',
     'OptimizationFromConfigOptions',

packages/optimization/src/ldai_optimization/dataclasses.py

@@ -14,6 +14,7 @@
     Sequence,
     Union,
 )
+from typing_extensions import Protocol
 
 from ldai import AIAgentConfig
 from ldai.models import LDMessage, ModelConfig
@@ -108,6 +109,33 @@ def from_dict(cls, data: Dict[str, Any]) -> "ToolDefinition":
         )
 
 
+class LLMCallConfig(Protocol):
+    """Structural protocol satisfied by both ``AIAgentConfig`` and ``AIJudgeCallConfig``.
+
+    Use this as the config parameter type when you want a single handler function
+    that can be passed to both ``handle_agent_call`` and ``handle_judge_call``::
+
+        async def handle_llm_call(
+            key: str,
+            config: LLMCallConfig,
+            context: Union[OptimizationContext, OptimizationJudgeContext],
+        ) -> OptimizationResponse:
+            model_name = config.model.name if config.model else "gpt-4o"
+            instructions = config.instructions or ""
+            ...
+
+        OptimizationOptions(
+            handle_agent_call=handle_llm_call,
+            handle_judge_call=handle_llm_call,
+            ...
+        )
+    """
+
+    key: str
+    model: Optional[ModelConfig]
+    instructions: Optional[str]
+
+
 @dataclass
 class AIJudgeCallConfig:
     """

packages/optimization/tests/test_client.py

@@ -3818,6 +3818,16 @@ async def test_commit_not_called_when_run_fails(self):
 
         mock_commit.assert_not_called()
 
+    async def test_succeeds_without_api_key_when_auto_commit_false(self):
+        client = self._make_client_without_key()
+        options = _make_options()  # auto_commit defaults to False
+
+        with patch("ldai_optimization.client.LDApiClient") as mock_api_cls:
+            result = await client.optimize_from_options("test-agent", options)
+
+        mock_api_cls.assert_not_called()
+        assert result is not None
+
     async def test_raises_when_auto_commit_true_and_no_api_key(self):
         client = self._make_client_without_key()
         options = _make_options(auto_commit=True, project_key="my-project")