From bde9b5d5484d332eebec4b0382f356e6ea6f9cb6 Mon Sep 17 00:00:00 2001 From: Jeremy Mumford Date: Mon, 27 Jan 2025 12:01:42 -0700 Subject: [PATCH 1/2] added TypeVar so that output type matches input type --- sdks/python/src/opik/integrations/openai/opik_tracker.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sdks/python/src/opik/integrations/openai/opik_tracker.py b/sdks/python/src/opik/integrations/openai/opik_tracker.py index 109c42f00b..8038135228 100644 --- 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 +T = TypeVar("OpenAI Client", openai.OpenAI, openai.AsyncOpenAI) + def track_openai( - openai_client: Union[openai.OpenAI, openai.AsyncOpenAI], + openai_client: T, project_name: Optional[str] = None, -) -> Union[openai.OpenAI, openai.AsyncOpenAI]: +) -> T: """Adds Opik tracking to an OpenAI client. Tracks calls to: From f8735ddce00cdc12f74fde0225f7216074b26616 Mon Sep 17 00:00:00 2001 From: Jeremy Mumford Date: Mon, 27 Jan 2025 15:24:36 -0700 Subject: [PATCH 2/2] updated var name to match according to mypy standards --- sdks/python/src/opik/integrations/openai/opik_tracker.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/python/src/opik/integrations/openai/opik_tracker.py b/sdks/python/src/opik/integrations/openai/opik_tracker.py index 8038135228..630678779c 100644 --- a/sdks/python/src/opik/integrations/openai/opik_tracker.py +++ b/sdks/python/src/opik/integrations/openai/opik_tracker.py @@ -4,13 +4,13 @@ from . import chat_completion_chunks_aggregator, openai_decorator -T = TypeVar("OpenAI Client", openai.OpenAI, openai.AsyncOpenAI) +OpenAIClient = TypeVar("OpenAIClient", openai.OpenAI, openai.AsyncOpenAI) def track_openai( - openai_client: T, + openai_client: OpenAIClient, project_name: Optional[str] = None, -) -> T: +) -> OpenAIClient: """Adds Opik tracking to an OpenAI client. Tracks calls to: