Resolve Typing Issues on Wrapping OpenAI Client (#1156)

* added TypeVar so that output type matches input type * updated var name to match according to mypy standards
comet-ml · Jan 28, 2025 · 4e06ae3 · 4e06ae3
1 parent 0e6f946
commit 4e06ae3
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/sdks/python/src/opik/integrations/openai/opik_tracker.py b/sdks/python/src/opik/integrations/openai/opik_tracker.py
@@ -1,14 +1,16 @@
-from typing import Optional, Union
+from typing import Optional, TypeVar
 
 import openai
 
 from . import chat_completion_chunks_aggregator, openai_decorator
 
+OpenAIClient = TypeVar("OpenAIClient", openai.OpenAI, openai.AsyncOpenAI)
+
 
 def track_openai(
-    openai_client: Union[openai.OpenAI, openai.AsyncOpenAI],
+    openai_client: OpenAIClient,
     project_name: Optional[str] = None,
-) -> Union[openai.OpenAI, openai.AsyncOpenAI]:
+) -> OpenAIClient:
     """Adds Opik tracking to an OpenAI client.
 
     Tracks calls to: