diff --git a/langfuse/openai.py b/langfuse/openai.py index b8243eb97..f7b3e3e48 100644 --- a/langfuse/openai.py +++ b/langfuse/openai.py @@ -10,6 +10,7 @@ from langfuse.client import InitialGeneration, CreateTrace, StatefulGenerationClient import openai +from openai import AsyncOpenAI, OpenAI from wrapt import wrap_function_wrapper from langfuse.model import UpdateGeneration @@ -336,15 +337,15 @@ def register_tracing(self): setattr(openai, "langfuse_host", None) setattr(openai, "flush_langfuse", self.flush) - setattr(openai.AsyncOpenAI, "langfuse_public_key", None) - setattr(openai.AsyncOpenAI, "langfuse_secret_key", None) - setattr(openai.AsyncOpenAI, "langfuse_host", None) - setattr(openai.AsyncOpenAI, "flush_langfuse", self.flush) + setattr(AsyncOpenAI, "langfuse_public_key", None) + setattr(AsyncOpenAI, "langfuse_secret_key", None) + setattr(AsyncOpenAI, "langfuse_host", None) + setattr(AsyncOpenAI, "flush_langfuse", self.flush) - setattr(openai.OpenAI, "langfuse_public_key", None) - setattr(openai.OpenAI, "langfuse_secret_key", None) - setattr(openai.OpenAI, "langfuse_host", None) - setattr(openai.OpenAI, "flush_langfuse", self.flush) + setattr(OpenAI, "langfuse_public_key", None) + setattr(OpenAI, "langfuse_secret_key", None) + setattr(OpenAI, "langfuse_host", None) + setattr(OpenAI, "flush_langfuse", self.flush) modifier = OpenAILangfuse() diff --git a/tests/test_openai.py b/tests/test_openai.py index 0adafbb3d..4638b45b4 100644 --- a/tests/test_openai.py +++ b/tests/test_openai.py @@ -2,10 +2,9 @@ import pytest from langfuse.client import Langfuse from langfuse.model import CreateTrace -from langfuse.openai import _is_openai_v1, _is_streaming_response, openai +from langfuse.openai import _is_openai_v1, _is_streaming_response, openai, AsyncOpenAI from tests.utils import create_uuid, get_api -from openai import AsyncOpenAI chat_func = openai.chat.completions.create if _is_openai_v1() else openai.ChatCompletion.create @@ -503,6 +502,7 @@ async def test_async_chat(): async def test_async_chat_stream(): api = get_api() client = AsyncOpenAI() + generation_name = create_uuid() completion = await client.chat.completions.create(messages=[{"role": "user", "content": "1 + 1 = "}], model="gpt-3.5-turbo", name=generation_name, stream=True)