Integrate OpenAI (#103)

Co-authored-by: Max Deichmann <[email protected]>

langfuse/__init__.py

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

langfuse/integrations.py

@@ -0,0 +1,127 @@
+import os
+import threading
+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, metadata=None, **kwargs):
+        self.args = {}
+        self.args["name"] = name
+        self.args["metadata"] = metadata
+        self.kwargs = kwargs
+
+    def get_langfuse_args(self):
+        return {**self.args, **self.kwargs}
+
+    def get_openai_args(self):
+        return self.kwargs
+
+
+class OpenAILangfuse:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super(OpenAILangfuse, cls).__new__(cls)
+                    cls._instance.initialize()
+        return cls._instance
+
+    def initialize(self):
+        self.langfuse = Langfuse()
+
+    @classmethod
+    def flush(cls):
+        cls._instance.langfuse.flush()
+
+    def _get_call_details(self, result, **kwargs):
+        name = kwargs.get("name", "OpenAI-generation")
+
+        if not isinstance(name, str):
+            raise TypeError("name must be a string")
+
+        metadata = kwargs.get("metadata", {})
+
+        if not isinstance(metadata, dict):
+            raise TypeError("metadata must be a dictionary")
+
+        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 = None
+        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,
+            "metadata": metadata,
+        }
+        return all_details
+
+    def _log_result(self, result, call_details):
+        generation = InitialGeneration(**call_details)
+        self.langfuse.generation(generation)
+        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))
+
+        setattr(openai, "flush_langfuse", self.flush)
+
+
+modifier = OpenAILangfuse()
+modifier.replace_openai_funcs()

tests/test_integrations.py

@@ -0,0 +1,155 @@
+import os
+import pytest
+from dotenv import load_dotenv
+from langfuse.integrations import openai
+from langfuse.api.client import FintoLangfuse
+
+
+from tests.utils import create_uuid
+
+
+load_dotenv()
+
+api = FintoLangfuse(
+    environment=os.environ["LANGFUSE_HOST"],
+    username=os.environ["LANGFUSE_PUBLIC_KEY"],
+    password=os.environ["LANGFUSE_SECRET_KEY"],
+)
+
+
+@pytest.mark.skip(reason="inference cost")
+def test_openai_chat_completion():
+    generation_name = create_uuid()
+    completion = openai.ChatCompletion.create(
+        name=generation_name,
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "1 + 1 = "}],
+        temperature=0,
+        metadata={"someKey": "someResponse"},
+    )
+
+    openai.flush_langfuse()
+
+    generation = api.observations.get_many(name=generation_name, type="GENERATION")
+
+    assert len(generation.data) != 0
+    assert generation.data[0].name == generation_name
+    assert generation.data[0].metadata == {"someKey": "someResponse"}
+    assert len(completion.choices) != 0
+    assert completion.choices[0].message.content == generation.data[0].output
+    assert generation.data[0].input == "1 + 1 = "
+    assert generation.data[0].type == "GENERATION"
+    assert generation.data[0].model == "gpt-3.5-turbo-0613"
+    assert generation.data[0].start_time is not None
+    assert generation.data[0].end_time is not None
+    assert generation.data[0].start_time < generation.data[0].end_time
+    assert generation.data[0].model_parameters == {
+        "temperature": "0",
+        "top_p": "1",
+        "frequency_penalty": "0",
+        "maxTokens": "inf",
+        "presence_penalty": "0",
+    }
+    assert generation.data[0].prompt_tokens is not None
+    assert generation.data[0].completion_tokens is not None
+    assert generation.data[0].total_tokens is not None
+
+
+@pytest.mark.skip(reason="inference cost")
+def test_openai_chat_completion_two_calls():
+    generation_name = create_uuid()
+    completion = openai.ChatCompletion.create(
+        name=generation_name,
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "1 + 1 = "}],
+        temperature=0,
+        metadata={"someKey": "someResponse"},
+    )
+
+    generation_name_2 = create_uuid()
+
+    completion_2 = openai.ChatCompletion.create(
+        name=generation_name_2,
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "2 + 2 = "}],
+        temperature=0,
+        metadata={"someKey": "someResponse"},
+    )
+
+    openai.flush_langfuse()
+
+    generation = api.observations.get_many(name=generation_name, type="GENERATION")
+
+    assert len(generation.data) != 0
+    assert generation.data[0].name == generation_name
+    assert len(completion.choices) != 0
+    assert completion.choices[0].message.content == generation.data[0].output
+    assert generation.data[0].input == "1 + 1 = "
+
+    generation_2 = api.observations.get_many(name=generation_name_2, type="GENERATION")
+
+    assert len(generation_2.data) != 0
+    assert generation_2.data[0].name == generation_name_2
+    assert len(completion_2.choices) != 0
+    assert completion_2.choices[0].message.content == generation_2.data[0].output
+    assert generation_2.data[0].input == "2 + 2 = "
+
+
+@pytest.mark.skip(reason="inference cost")
+def test_openai_completion():
+    generation_name = create_uuid()
+    completion = openai.Completion.create(
+        name=generation_name,
+        model="gpt-3.5-turbo-instruct",
+        prompt="1 + 1 = ",
+        temperature=0,
+        metadata={"someKey": "someResponse"},
+    )
+
+    openai.flush_langfuse()
+
+    generation = api.observations.get_many(name=generation_name, type="GENERATION")
+
+    assert len(generation.data) != 0
+    assert generation.data[0].name == generation_name
+    assert generation.data[0].metadata == {"someKey": "someResponse"}
+    assert len(completion.choices) != 0
+    assert completion.choices[0].text == generation.data[0].output
+    assert generation.data[0].input == "1 + 1 = "
+    assert generation.data[0].type == "GENERATION"
+    assert generation.data[0].model == "gpt-3.5-turbo-instruct"
+    assert generation.data[0].start_time is not None
+    assert generation.data[0].end_time is not None
+    assert generation.data[0].start_time < generation.data[0].end_time
+    assert generation.data[0].model_parameters == {
+        "temperature": "0",
+        "top_p": "1",
+        "frequency_penalty": "0",
+        "maxTokens": "inf",
+        "presence_penalty": "0",
+    }
+    assert generation.data[0].prompt_tokens is not None
+    assert generation.data[0].completion_tokens is not None
+    assert generation.data[0].total_tokens is not None
+
+
+@pytest.mark.skip(reason="inference cost")
+def test_fails_wrong_name():
+    with pytest.raises(TypeError, match="name must be a string"):
+        openai.Completion.create(
+            name={"key": "generation_name"},
+            model="gpt-3.5-turbo-instruct",
+            prompt="1 + 1 = ",
+            temperature=0,
+        )
+
+
+@pytest.mark.skip(reason="inference cost")
+def test_fails_wrong_metadata():
+    with pytest.raises(TypeError, match="name must be a string"):
+        openai.Completion.create(
+            metadata="metadata",
+            model="gpt-3.5-turbo-instruct",
+            prompt="1 + 1 = ",
+            temperature=0,
+        )