chore: update types for lint

Author: andrewklatzke

packages/optimization/src/ldai_optimization/dataclasses.py

@@ -119,6 +119,7 @@ async def handle_llm_call(
             key: str,
             config: LLMCallConfig,
             context: LLMCallContext,
+            is_evaluation: bool,
         ) -> OptimizationResponse:
             model_name = config.model.name if config.model else "gpt-4o"
             instructions = config.instructions or ""
@@ -132,9 +133,12 @@ async def handle_llm_call(
         )
     """
 
-    key: str
-    model: Optional[ModelConfig]
-    instructions: Optional[str]
+    @property
+    def key(self) -> str: ...
+    @property
+    def model(self) -> Optional[ModelConfig]: ...
+    @property
+    def instructions(self) -> Optional[str]: ...
 
 
 class LLMCallContext(Protocol):
@@ -144,8 +148,10 @@ class LLMCallContext(Protocol):
     ``handle_agent_call`` and ``handle_judge_call``.
     """
 
-    user_input: Optional[str]
-    current_variables: Dict[str, Any]
+    @property
+    def user_input(self) -> Optional[str]: ...
+    @property
+    def current_variables(self) -> Dict[str, Any]: ...
 
 
 @dataclass

packages/optimization/src/ldai_optimization/util.py

@@ -4,7 +4,7 @@
 import json
 import logging
 import re
-from typing import Any, Awaitable, Dict, List, Optional, Tuple, Union
+from typing import Any, Awaitable, Dict, List, Optional, Tuple, TypeVar, Union
 
 from ldai_optimization.dataclasses import ToolDefinition
 
@@ -156,18 +156,19 @@ def restore_variable_placeholders(
     return text, warnings
 
 
-async def await_if_needed(
-    result: Union[str, Awaitable[str]]
-) -> str:
+_T = TypeVar("_T")
+
+
+async def await_if_needed(result: Union[_T, Awaitable[_T]]) -> _T:
     """
     Handle both sync and async callable results.
 
-    :param result: Either a string or an awaitable that returns a string
-    :return: The string result
+    :param result: Either a value or an awaitable that returns a value
+    :return: The resolved value
     """
-    if isinstance(result, str):
-        return result
-    return await result
+    if inspect.isawaitable(result):
+        return await result  # type: ignore[return-value]
+    return result  # type: ignore[return-value]
 
 
 def create_evaluation_tool() -> ToolDefinition:

packages/optimization/tests/test_ld_api_client.py

@@ -11,9 +11,9 @@
 
 from ldai_optimization.ld_api_client import (
     AgentOptimizationConfig,
+    AgentOptimizationResultPost as OptimizationResultPayload,
     LDApiClient,
     LDApiError,
-    OptimizationResultPayload,
     _parse_agent_optimization,
 )