diff --git a/tests/gemini/test_generate_content_with_tool_calls.py b/tests/gemini/test_generate_content_with_tool_calls.py new file mode 100644 index 0000000..4e61ed6 --- /dev/null +++ b/tests/gemini/test_generate_content_with_tool_calls.py @@ -0,0 +1,191 @@ +"""Test the chat completions gateway calls with tool calling and processing response.""" + +import os +import sys +import uuid +import pytest + +from google import genai +from google.genai import types + +# add tests folder (parent) to sys.path +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from util import * # needed for pytest fixtures + +pytest_plugins = ("pytest_asyncio",) + + +def set_light_values(brightness: int, color_temp: str) -> dict[str, int | str]: + """Set the brightness and color temperature of a room light. (mock API). + + Args: + brightness: Light level from 0 to 100. Zero is off and 100 is full brightness + color_temp: Color temperature of the light fixture, which can be `daylight`, + `cool` or `warm`. + + Returns: + A dictionary containing the set brightness and color temperature. + """ + return { + "brightness": brightness, + "colorTemperature": color_temp, + } + + +SYSTEM_INSTRUCTION = "This the system instruction. Use the function call." +USER_PROMPT = ( + "Turn the light to 50% brightness and set the color temperature to daylight. \ + Once you are done, respond with 'DONE'." +) +SET_LIGHT_VALUES_TOOL_CALL_ARGS = {"brightness": 50, "color_temp": "daylight"} +SET_LIGHT_VALUES_TOOL_CALL = { + "type": "function", + "function": { + "name": "set_light_values", + "arguments": SET_LIGHT_VALUES_TOOL_CALL_ARGS, + }, +} + + +async def _verify_trace_from_explorer( + context, explorer_api_url, dataset_name, expected_final_assistant_message +) -> None: + # Fetch the trace ids for the dataset. + # There will be 2 traces - the first will contain the system instruction, user prompt + # and the assistant tool call. + # The second will contain the system instruction, user prompt, the assistant tool call, + # the tool response and the assistant response. + 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) == 2 + trace_id_1 = traces[0]["id"] + trace_id_2 = traces[1]["id"] + + # Fetch the trace + trace_response_1 = await context.request.get( + f"{explorer_api_url}/api/v1/trace/{trace_id_1}" + ) + trace_1 = await trace_response_1.json() + + trace_response_2 = await context.request.get( + f"{explorer_api_url}/api/v1/trace/{trace_id_2}" + ) + trace_2 = await trace_response_2.json() + + # Verify the trace messages + assert trace_1["messages"] == [ + { + "role": "system", + "content": SYSTEM_INSTRUCTION, + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": USER_PROMPT, + } + ], + }, + { + "role": "assistant", + "tool_calls": [SET_LIGHT_VALUES_TOOL_CALL], + }, + ] + + assert trace_2["messages"] == [ + { + "role": "system", + "content": SYSTEM_INSTRUCTION, + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": USER_PROMPT, + } + ], + }, + { + "role": "assistant", + "tool_calls": [SET_LIGHT_VALUES_TOOL_CALL], + }, + { + "role": "tool", + "tool_name": "set_light_values", + "content": { + "brightness": 50, + "colorTemperature": "daylight", + }, + }, + {"role": "assistant", "content": expected_final_assistant_message}, + ] + + +@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set") +@pytest.mark.parametrize( + "do_stream, push_to_explorer", + [(True, True), (True, False), (False, True), (False, False)], +) +async def test_generate_content_with_tool_call( + context, explorer_api_url, gateway_url, push_to_explorer, do_stream +): + """ + Test the generate content gateway calls with tool calling and response processing + without streaming. + """ + dataset_name = "test-dataset-gemini-tool-call-" + str(uuid.uuid4()) + + client = genai.Client( + api_key=os.getenv("GEMINI_API_KEY"), + http_options={ + "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini" + if push_to_explorer + else f"{gateway_url}/api/v1/gateway/gemini", + "headers": { + "invariant-authorization": "Bearer " + }, # This key is not used for local tests + }, + ) + + request = { + "model": "gemini-2.0-flash", + "contents": USER_PROMPT, + "config": types.GenerateContentConfig( + tools=[set_light_values], + system_instruction=SYSTEM_INSTRUCTION, + ), + } + + chat_response = ( + client.models.generate_content(**request) + if not do_stream + else client.models.generate_content_stream(**request) + ) + + if not do_stream: + assert "DONE" in chat_response.candidates[0].content.parts[0].text + expected_final_assistant_message = ( + chat_response.candidates[0].content.parts[0].text + ) + else: + full_response = "" + for chunk in chat_response: + if ( + chunk.candidates + and chunk.candidates[0].content + and chunk.candidates[0].content.parts + ): + for text_part in chunk.candidates[0].content.parts: + full_response += text_part.text + assert "DONE" in full_response.upper() + expected_final_assistant_message = full_response + + if push_to_explorer: + await _verify_trace_from_explorer( + context, explorer_api_url, dataset_name, expected_final_assistant_message + ) diff --git a/tests/gemini/test_generate_content_without_tool_calls.py b/tests/gemini/test_generate_content_without_tool_calls.py index f92cd69..28a0e68 100644 --- a/tests/gemini/test_generate_content_without_tool_calls.py +++ b/tests/gemini/test_generate_content_without_tool_calls.py @@ -1,6 +1,5 @@ """Test the generate content gateway calls without tool calling.""" -import base64 import os import sys import uuid @@ -8,7 +7,6 @@ from pathlib import Path from unittest.mock import patch import pytest -from httpx import Client from google import genai import PIL.Image @@ -52,10 +50,11 @@ async def test_generate_content( }, } - if not do_stream: - chat_response = client.models.generate_content(**request) - else: - chat_response = client.models.generate_content_stream(**request) + chat_response = ( + client.models.generate_content(**request) + if not do_stream + else client.models.generate_content_stream(**request) + ) # Verify the chat response if not do_stream: