Add test for openai chat completions API via proxy with tool call.

2026-07-06 10:57:50 +02:00 · 2025-02-10 15:21:46 +01:00
parent ba26d7d0c7
commit 64ddd9f96d
1 changed files with 140 additions and 0 deletions
@@ -0,0 +1,140 @@
+"""Test the chat completions proxy calls with tool calling and processing response."""
+
+import json
+import os
+import sys
+import uuid
+
+import pytest
+from httpx import Client
+
+# add tests folder (parent) to sys.path
+from openai import OpenAI
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from util import *  # needed for pytest fixtures
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+async def test_chat_completion_with_tool_call_without_streaming(
+    context, explorer_api_url, proxy_url
+):
+    """Test the chat completions proxy calls with tool calling and response processing."""
+    dataset_name = "test-dataset-open-ai-tool-call-" + str(uuid.uuid4())
+
+    client = OpenAI(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": "Bearer <some-key>"
+            },  # This key is not used for local tests
+        ),
+        base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/openai",
+    )
+
+    chat_response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "What is the weather in New York?"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                    },
+                },
+            }
+        ],
+    )
+
+    assert chat_response.choices[0].message.role == "assistant"
+    # Extract tool call
+    assert len(chat_response.choices[0].message.tool_calls) == 1
+    assert chat_response.choices[0].message.tool_calls[0].function.name == "get_weather"
+    assert (
+        chat_response.choices[0].message.tool_calls[0].function.arguments
+        == '{"location":"New York"}'
+    )
+    tool_call = chat_response.choices[0].message.tool_calls[0]
+
+    # Mock response of tool call
+    tool_result = "The temperature in New York is 15°C and it is raining."
+    history = [
+        {"role": "user", "content": "What is the weather in New York?"},
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "function": {
+                        "arguments": tool_call.function.arguments,
+                        "name": tool_call.function.name,
+                    },
+                    "id": tool_call.id,
+                    "type": tool_call.type,
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "tool_call_id": tool_call.id,
+            "tool_name": "get_weather",
+            "content": tool_result,
+        },
+    ]
+
+    # Send mock response back to OpenAI with history of chat
+    chat_response_final = client.chat.completions.create(
+        model="gpt-4o",
+        messages=history,
+    )
+    assert "15°C" in chat_response_final.choices[0].message.content
+
+    # Fetch the trace ids for the dataset
+    traces_response = await context.request.get(
+        f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
+    )
+    traces = await traces_response.json()
+    assert len(traces) == 1
+    trace_id = traces[0]["id"]
+
+    # Fetch the trace
+    trace_response = await context.request.get(
+        f"{explorer_api_url}/api/v1/trace/{trace_id}"
+    )
+    trace = await trace_response.json()
+
+    # Verify the trace messages
+    assert trace["messages"] == [
+        {
+            "role": "user",
+            "content": "What is the weather in New York?",
+        },
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "function": {
+                        "arguments": json.loads(tool_call.function.arguments),
+                        "name": tool_call.function.name,
+                    },
+                    "id": tool_call.id,
+                    "type": tool_call.type,
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "tool_call_id": tool_call.id,
+            "tool_name": "get_weather",
+            "content": tool_result,
+        },
+        {
+            "role": "assistant",
+            "content": chat_response_final.choices[0].message.content,
+        },
+    ]