Add tests with tool_calls for Gemini.

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 <some-key>"
+            },  # 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
+        )

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: