Integrate OpenAI

langfuse/__init__.py

@@ -1 +1,2 @@
 from .client import Langfuse
+from .integrations import openai

langfuse/integrations.py

@@ -0,0 +1,99 @@
+import os
+import functools
+from datetime import datetime
+from dotenv import load_dotenv
+
+import openai
+from openai.api_resources import ChatCompletion, Completion
+
+from langfuse import Langfuse
+from langfuse.client import InitialGeneration
+from langfuse.api.resources.commons.types.llm_usage import LlmUsage
+
+
+load_dotenv()
+
+
+class CreateArgsExtractor:
+    def __init__(self, name=None, **kwargs):
+        self.args = {}
+        self.args["name"] = name
+        self.kwargs = kwargs
+
+    def get_langfuse_args(self):
+        return {**self.args, **self.kwargs}
+
+    def get_openai_args(self):
+        return self.kwargs
+
+
+class OpenAILangfuse:
+    def __init__(self):
+        self.langfuse = Langfuse(os.environ["LF_PK"], os.environ["LF_SK"], os.environ["HOST"])
+
+    def _get_call_details(self, result, **kwargs):
+        name = kwargs.get("name", "OpenAI-generation")
+        if result.object == "chat.completion":
+            prompt = kwargs.get("messages", [{}])[-1].get("content", "")
+            completion = result.choices[-1].message.content
+        elif result.object == "text_completion":
+            prompt = kwargs.get("prompt", "")
+            completion = result.choices[-1].text
+        else:
+            completion = ""
+        model = result.model
+        usage = None if result.usage is None else LlmUsage(**result.usage)
+        endTime = datetime.now()
+        modelParameters = {
+            "temperature": kwargs.get("temperature", 1),
+            "maxTokens": kwargs.get("max_tokens", float("inf")),
+            "top_p": kwargs.get("top_p", 1),
+            "frequency_penalty": kwargs.get("frequency_penalty", 0),
+            "presence_penalty": kwargs.get("presence_penalty", 0),
+        }
+        all_details = {
+            "name": name,
+            "prompt": prompt,
+            "completion": completion,
+            "endTime": endTime,
+            "model": model,
+            "modelParameters": modelParameters,
+            "usage": usage,
+        }
+        return all_details
+
+    def _log_result(self, result, call_details):
+        generation = InitialGeneration(**call_details)
+        self.langfuse.generation(generation)
+        self.langfuse.flush()
+        return result
+
+    def langfuse_modified(self, func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                startTime = datetime.now()
+                arg_extractor = CreateArgsExtractor(*args, **kwargs)
+                result = func(**arg_extractor.get_openai_args())
+                call_details = self._get_call_details(result, **arg_extractor.get_langfuse_args())
+                call_details["startTime"] = startTime
+            except Exception as ex:
+                raise ex
+
+            return self._log_result(result, call_details)
+
+        return wrapper
+
+    def replace_openai_funcs(self):
+        api_resources_classes = [
+            (ChatCompletion, "create"),
+            (Completion, "create"),
+        ]
+
+        for api_resource_class, method in api_resources_classes:
+            create_method = getattr(api_resource_class, method)
+            setattr(api_resource_class, method, self.langfuse_modified(create_method))
+
+
+modifier = OpenAILangfuse()
+modifier.replace_openai_funcs()

tests/test_integrations.py

@@ -0,0 +1,40 @@
+import os
+from dotenv import load_dotenv
+from langfuse.integrations import openai
+from langfuse.api.client import FintoLangfuse
+from langfuse.version import __version__ as version
+
+from tests.utils import create_uuid
+
+
+load_dotenv()
+
+api = FintoLangfuse(
+    environment=os.environ["HOST"],
+    username=os.environ["LF_PK"],
+    password=os.environ["LF_SK"],
+)
+
+
+def test_openai_chat_completion():
+    trace_id = create_uuid()
+    completion = openai.ChatCompletion.create(
+        name=trace_id, model="gpt-3.5-turbo", messages=[{"role": "user", "content": "1 + 1 = "}], temperature=0
+    )
+
+    generation = api.generations.get(name=trace_id)
+
+    assert len(generation.data) != 0
+    assert len(completion.choices) != 0
+
+
+def test_openai_completion():
+    trace_id = create_uuid()
+    completion = openai.Completion.create(
+        name=trace_id, model="gpt-3.5-turbo-instruct", prompt="1 + 1 = ", temperature=0
+    )
+
+    generation = api.generations.get(name=trace_id)
+
+    assert len(generation.data) != 0
+    assert len(completion.choices) != 0