fix(openai): export openai client directly from langchain (#183)

* export openai client directly from langchain
langfuse · Nov 20, 2023 · bc67380 · bc67380
1 parent e5c37b8
commit bc67380
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git 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
@@ -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)