Improve type annotations for overloaded client methods (#70)

* Fix type annotations in overload_client to preserve method return types

* Further refine types in overload.py

* refactor: simplify method overloading by using direct assignment

Replace getattr/setattr with direct attribute assignment in overload_client function. Add type ignore comments to handle type checking. This change improves code readability while maintaining the same functionality.

* chore: cleaned up imports in sync_client.py

Author: AlePouroullis

src/humanloop/client.py

@@ -1,36 +1,34 @@
+import logging
 import os
 import typing
 from typing import Any, List, Optional, Sequence, Tuple
-import logging
 
 import httpx
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.trace import Tracer
 
+from humanloop.base_client import AsyncBaseHumanloop, BaseHumanloop
 from humanloop.core.client_wrapper import SyncClientWrapper
-
+from humanloop.decorators.flow import flow as flow_decorator_factory
+from humanloop.decorators.prompt import prompt_decorator_factory
+from humanloop.decorators.tool import tool_decorator_factory as tool_decorator_factory
+from humanloop.environment import HumanloopEnvironment
 from humanloop.evals import run_eval
 from humanloop.evals.types import (
     DatasetEvalConfig,
-    EvaluatorEvalConfig,
     EvaluatorCheck,
+    EvaluatorEvalConfig,
     FileEvalConfig,
 )
-
-from humanloop.base_client import AsyncBaseHumanloop, BaseHumanloop
-from humanloop.overload import overload_client
-from humanloop.decorators.flow import flow as flow_decorator_factory
-from humanloop.decorators.prompt import prompt_decorator_factory
-from humanloop.decorators.tool import tool_decorator_factory as tool_decorator_factory
-from humanloop.environment import HumanloopEnvironment
 from humanloop.evaluations.client import EvaluationsClient
 from humanloop.otel import instrument_provider
 from humanloop.otel.exporter import HumanloopSpanExporter
 from humanloop.otel.processor import HumanloopSpanProcessor
+from humanloop.overload import overload_client
 from humanloop.prompt_utils import populate_template
 from humanloop.prompts.client import PromptsClient
-from humanloop.sync.sync_client import SyncClient, DEFAULT_CACHE_SIZE
+from humanloop.sync.sync_client import DEFAULT_CACHE_SIZE, SyncClient
 
 logger = logging.getLogger("humanloop.sdk")
 
@@ -168,6 +166,7 @@ def __init__(
 
         # Overload the .log method of the clients to be aware of Evaluation Context
         # and the @flow decorator providing the trace_id
+        # Additionally, call and log methods are overloaded in the prompts and agents client to support the use of local files
         self.prompts = overload_client(
             client=self.prompts, sync_client=self._sync_client, use_local_files=self.use_local_files
         )

src/humanloop/overload.py

@@ -1,31 +1,32 @@
 import inspect
 import logging
 import types
-from typing import Any, Dict, Optional, Union, Callable
+from typing import Any, Callable, Dict, Optional, TypeVar, Union
 
+from humanloop.agents.client import AgentsClient
 from humanloop.context import (
     get_decorator_context,
     get_evaluation_context,
     get_trace_id,
 )
+from humanloop.datasets.client import DatasetsClient
 from humanloop.error import HumanloopRuntimeError
-from humanloop.sync.sync_client import SyncClient
-from humanloop.prompts.client import PromptsClient
+from humanloop.evaluators.client import EvaluatorsClient
 from humanloop.flows.client import FlowsClient
-from humanloop.datasets.client import DatasetsClient
-from humanloop.agents.client import AgentsClient
+from humanloop.prompts.client import PromptsClient
+from humanloop.sync.sync_client import SyncClient
 from humanloop.tools.client import ToolsClient
-from humanloop.evaluators.client import EvaluatorsClient
 from humanloop.types import FileType
+from humanloop.types.agent_call_response import AgentCallResponse
 from humanloop.types.create_evaluator_log_response import CreateEvaluatorLogResponse
 from humanloop.types.create_flow_log_response import CreateFlowLogResponse
 from humanloop.types.create_prompt_log_response import CreatePromptLogResponse
 from humanloop.types.create_tool_log_response import CreateToolLogResponse
 from humanloop.types.prompt_call_response import PromptCallResponse
-from humanloop.types.agent_call_response import AgentCallResponse
 
 logger = logging.getLogger("humanloop.sdk")
 
+
 LogResponseType = Union[
     CreatePromptLogResponse,
     CreateToolLogResponse,
@@ -39,6 +40,9 @@
 ]
 
 
+T = TypeVar("T", bound=Union[PromptsClient, AgentsClient, ToolsClient, FlowsClient, DatasetsClient, EvaluatorsClient])
+
+
 def _get_file_type_from_client(
     client: Union[PromptsClient, AgentsClient, ToolsClient, FlowsClient, DatasetsClient, EvaluatorsClient],
 ) -> FileType:
@@ -55,13 +59,13 @@ def _get_file_type_from_client(
         return "dataset"
     elif isinstance(client, EvaluatorsClient):
         return "evaluator"
-
-    raise ValueError(f"Unsupported client type: {type(client)}")
+    else:
+        raise ValueError(f"Unsupported client type: {type(client)}")
 
 
-def _handle_tracing_context(kwargs: Dict[str, Any], client: Any) -> Dict[str, Any]:
+def _handle_tracing_context(kwargs: Dict[str, Any], client: T) -> Dict[str, Any]:
     """Handle tracing context for both log and call methods."""
-    trace_id = get_trace_id()
+    trace_id = get_trace_id()   
     if trace_id is not None:
         if "flow" in str(type(client).__name__).lower():
             context = get_decorator_context()
@@ -86,7 +90,7 @@ def _handle_tracing_context(kwargs: Dict[str, Any], client: Any) -> Dict[str, An
 
 def _handle_local_files(
     kwargs: Dict[str, Any],
-    client: Any,
+    client: T,
     sync_client: Optional[SyncClient],
     use_local_files: bool,
 ) -> Dict[str, Any]:
@@ -136,7 +140,7 @@ def _handle_evaluation_context(kwargs: Dict[str, Any]) -> tuple[Dict[str, Any],
     return kwargs, None
 
 
-def _overload_log(self: Any, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> LogResponseType:
+def _overload_log(self: T, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> LogResponseType:
     try:
         # Special handling for flows - prevent direct log usage
         if type(self) is FlowsClient and get_trace_id() is not None:
@@ -158,7 +162,7 @@ def _overload_log(self: Any, sync_client: Optional[SyncClient], use_local_files:
             kwargs = _handle_local_files(kwargs, self, sync_client, use_local_files)
 
         kwargs, eval_callback = _handle_evaluation_context(kwargs)
-        response = self._log(**kwargs)  # Use stored original method
+        response = self._log(**kwargs)  # type: ignore[union-attr] # Use stored original method
         if eval_callback is not None:
             eval_callback(response.id)
         return response
@@ -170,11 +174,11 @@ def _overload_log(self: Any, sync_client: Optional[SyncClient], use_local_files:
         raise HumanloopRuntimeError from e
 
 
-def _overload_call(self: Any, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> CallResponseType:
+def _overload_call(self: T, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> CallResponseType:
     try:
         kwargs = _handle_tracing_context(kwargs, self)
         kwargs = _handle_local_files(kwargs, self, sync_client, use_local_files)
-        return self._call(**kwargs)  # Use stored original method
+        return self._call(**kwargs)  # type: ignore[union-attr] # Use stored original method
     except HumanloopRuntimeError:
         # Re-raise HumanloopRuntimeError without wrapping to preserve the message
         raise
@@ -184,33 +188,37 @@ def _overload_call(self: Any, sync_client: Optional[SyncClient], use_local_files
 
 
 def overload_client(
-    client: Any,
+    client: T,
     sync_client: Optional[SyncClient] = None,
     use_local_files: bool = False,
-) -> Any:
+) -> T:
     """Overloads client methods to add tracing, local file handling, and evaluation context."""
     # Store original log method as _log for all clients. Used in flow decorator
     if hasattr(client, "log") and not hasattr(client, "_log"):
-        client._log = client.log  # type: ignore[attr-defined]
+        # Store original method with type ignore
+        client._log = client.log  # type: ignore
 
         # Create a closure to capture sync_client and use_local_files
-        def log_wrapper(self: Any, **kwargs) -> LogResponseType:
+        def log_wrapper(self: T, **kwargs) -> LogResponseType:
             return _overload_log(self, sync_client, use_local_files, **kwargs)
 
-        client.log = types.MethodType(log_wrapper, client)
+        # Replace the log method with type ignore
+        client.log = types.MethodType(log_wrapper, client)  # type: ignore
 
     # Overload call method for Prompt and Agent clients
     if _get_file_type_from_client(client) in ["prompt", "agent"]:
         if sync_client is None and use_local_files:
             logger.error("sync_client is None but client has call method and use_local_files=%s", use_local_files)
             raise HumanloopRuntimeError("sync_client is required for clients that support call operations")
         if hasattr(client, "call") and not hasattr(client, "_call"):
-            client._call = client.call  # type: ignore[attr-defined]
+            # Store original method with type ignore
+            client._call = client.call  # type: ignore
 
             # Create a closure to capture sync_client and use_local_files
-            def call_wrapper(self: Any, **kwargs) -> CallResponseType:
+            def call_wrapper(self: T, **kwargs) -> CallResponseType:
                 return _overload_call(self, sync_client, use_local_files, **kwargs)
 
-            client.call = types.MethodType(call_wrapper, client)
+            # Replace the call method with type ignore
+            client.call = types.MethodType(call_wrapper, client)  # type: ignore
 
     return client

src/humanloop/sync/sync_client.py

@@ -1,11 +1,12 @@
+import json
 import logging
-from pathlib import Path
-from typing import List, Optional, Tuple, TYPE_CHECKING, Union
-from functools import lru_cache
-import typing
 import time
+import typing
+from functools import lru_cache
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
 from humanloop.error import HumanloopRuntimeError
-import json
 
 if TYPE_CHECKING:
     from humanloop.base_client import BaseHumanloop