fix(openai): correctly parse openai function call responses (#186)

Author: maxdeichmann

langfuse/openai.py

@@ -15,6 +15,10 @@
 
 from langfuse.model import UpdateGeneration
 
+import logging
+
+log = logging.getLogger("langfuse")
+
 
 class OpenAiDefinition:
     module: str
@@ -236,7 +240,7 @@ def _get_langfuse_data_from_default_response(resource: OpenAiDefinition, respons
         choices = response.get("choices", [])
         if len(choices) > 0:
             choice = choices[-1]
-            completion = choice.message.content if _is_openai_v1() else choice.get("message", None).get("content", None)
+            completion = choice.message.json() if _is_openai_v1() else choice.get("message", None)
 
     usage = response.get("usage", None)
 
@@ -269,8 +273,10 @@ def _wrap(open_ai_resource: OpenAiDefinition, initialize, wrapped, args, kwargs)
         else:
             model, completion, usage = _get_langfuse_data_from_default_response(open_ai_resource, openai_response.__dict__ if _is_openai_v1() else openai_response)
             generation.update(UpdateGeneration(model=model, completion=completion, end_time=datetime.now(), usage=usage))
+
         return openai_response
     except Exception as ex:
+        log.warn(ex)
         model = kwargs.get("model", None)
         generation.update(UpdateGeneration(endTime=datetime.now(), statusMessage=str(ex), level="ERROR", model=model))
         raise ex

tests/test_openai.py

@@ -32,7 +32,6 @@ def test_openai_chat_completion():
     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 == [{"content": "1 + 1 = ", "role": "user"}]
     assert generation.data[0].type == "GENERATION"
     assert generation.data[0].model == "gpt-3.5-turbo-0613"
@@ -49,7 +48,7 @@ def test_openai_chat_completion():
     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
-    assert generation.data[0].output == "2"
+    assert "2" in generation.data[0].output
 
 
 def test_openai_chat_completion_stream():
@@ -249,15 +248,15 @@ def test_openai_chat_completion_two_calls():
     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 == [{"content": "1 + 1 = ", "role": "user"}]
 
     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 == [{"content": "2 + 2 = ", "role": "user"}]
 
 
@@ -478,7 +477,7 @@ async def test_async_chat():
     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 == [{"content": "1 + 1 = ", "role": "user"}]
     assert generation.data[0].type == "GENERATION"
     assert generation.data[0].model == "gpt-3.5-turbo-0613"
@@ -495,7 +494,7 @@ async def test_async_chat():
     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
-    assert generation.data[0].output == "2"
+    assert "2" in generation.data[0].output
 
 
 @pytest.mark.asyncio
@@ -534,3 +533,39 @@ async def test_async_chat_stream():
     assert generation.data[0].completion_tokens is not None
     assert generation.data[0].total_tokens is not None
     assert generation.data[0].output == "2"
+
+
+def test_openai_function_call():
+    from typing import List
+
+    from pydantic import BaseModel
+
+    api = get_api()
+    generation_name = create_uuid()
+
+    class StepByStepAIResponse(BaseModel):
+        title: str
+        steps: List[str]
+
+    import json
+
+    response = openai.chat.completions.create(
+        name=generation_name,
+        model="gpt-3.5-turbo-0613",
+        messages=[{"role": "user", "content": "Explain how to assemble a PC"}],
+        functions=[{"name": "get_answer_for_user_query", "description": "Get user answer in series of steps", "parameters": StepByStepAIResponse.schema()}],
+        function_call={"name": "get_answer_for_user_query"},
+    )
+
+    output = json.loads(response.choices[0].message.function_call.arguments)
+
+    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].output is not None
+    assert "function_call" in generation.data[0].output
+
+    assert output["title"] is not None