diff --git a/tests/open_ai/test_chat_with_tool_call.py b/tests/open_ai/test_chat_with_tool_call.py index fbd15db..422e5ff 100644 --- a/tests/open_ai/test_chat_with_tool_call.py +++ b/tests/open_ai/test_chat_with_tool_call.py @@ -22,7 +22,10 @@ pytest_plugins = ("pytest_asyncio",) 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.""" + """ + Test the chat completions proxy calls with tool calling and response processing + without streaming. + """ dataset_name = "test-dataset-open-ai-tool-call-" + str(uuid.uuid4()) client = OpenAI( @@ -109,32 +112,120 @@ async def test_chat_completion_with_tool_call_without_streaming( trace = await trace_response.json() # Verify the trace messages - assert trace["messages"] == [ + expected_messages = history + [ { - "role": "user", - "content": "What is the weather in New York?", - }, + "role": "assistant", + "content": chat_response_final.choices[0].message.content, + } + ] + expected_messages[1]["tool_calls"][0]["function"]["arguments"] = json.loads( + expected_messages[1]["tool_calls"][0]["function"]["arguments"] + ) + assert trace["messages"] == expected_messages + + +@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set") +async def test_chat_completion_with_tool_call_with_streaming( + context, explorer_api_url, proxy_url +): + """ + Test the chat completions proxy calls with tool calling and response processing + while streaming. + """ + dataset_name = "test-dataset-open-ai-tool-call-" + str(uuid.uuid4()) + + client = OpenAI( + http_client=Client( + headers={ + "Invariant-Authorization": "Bearer " + }, # 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"], + }, + }, + } + ], + stream=True, + ) + + tool_call = {"function": {}} + for chunk in chat_response: + if chunk.choices and chunk.choices[0].delta.tool_calls: + partial_tool_call = chunk.choices[0].delta.tool_calls[0] + tool_call.setdefault("id", partial_tool_call.id) + tool_call.setdefault("type", partial_tool_call.type) + tool_call["function"].setdefault("name", partial_tool_call.function.name) + tool_call["function"].setdefault("arguments", "") + tool_call["function"]["arguments"] += partial_tool_call.function.arguments + + assert tool_call["function"]["name"] == "get_weather" + assert tool_call["function"]["arguments"] == '{"location":"New York"}' + + # 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": json.loads(tool_call.function.arguments), - "name": tool_call.function.name, + "arguments": tool_call["function"]["arguments"], + "name": tool_call["function"]["name"], }, - "id": tool_call.id, - "type": tool_call.type, + "id": tool_call["id"], + "type": tool_call["type"], } ], }, { "role": "tool", - "tool_call_id": tool_call.id, + "tool_call_id": tool_call["id"], "tool_name": "get_weather", "content": tool_result, }, - { - "role": "assistant", - "content": chat_response_final.choices[0].message.content, - }, ] + + # Send mock response back to OpenAI with history of chat + chat_response_final = client.chat.completions.create( + model="gpt-4o", messages=history, stream=True + ) + final_response = {"role": "assistant", "content": ""} + for chunk in chat_response_final: + if chunk.choices and chunk.choices[0].delta.content: + final_response["content"] += chunk.choices[0].delta.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 + expected_messages = history + [final_response] + expected_messages[1]["tool_calls"][0]["function"]["arguments"] = json.loads( + expected_messages[1]["tool_calls"][0]["function"]["arguments"] + ) + assert trace["messages"] == expected_messages