sglang/sgl-router/py_test/e2e_response_api/mixins/function_call.py

"""
Base test class for function calling tests.

This module provides test cases for function calling functionality
across different backends.
"""

import json
import sys
from pathlib import Path

# Add current directory for local imports
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR))

from util import CustomTestCase


class ResponseAPIBaseTest(CustomTestCase):
    """Base class for Response API tests with common utilities."""

    # To be set by subclasses
    base_url: str = None
    api_key: str = None
    model: str = None

    def make_request(
        self,
        endpoint: str,
        method: str = "POST",
        json_data: dict = None,
        params: dict = None,
    ):
        """
        Make HTTP request to router.

        This is a minimal implementation - subclasses should import from basic_crud.
        """
        import requests

        url = f"{self.base_url}{endpoint}"
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        if method == "POST":
            resp = requests.post(url, json=json_data, headers=headers, params=params)
        elif method == "GET":
            resp = requests.get(url, headers=headers, params=params)
        elif method == "DELETE":
            resp = requests.delete(url, headers=headers, params=params)
        else:
            raise ValueError(f"Unsupported method: {method}")
        return resp


class FunctionCallingBaseTest(ResponseAPIBaseTest):
    """Base class for function calling tests."""

    def test_basic_function_call(self):
        """
        Test basic function calling workflow.

        This test follows the pattern from function_call_test.py:
        1. Define a function tool (get_horoscope)
        2. Send user message asking for horoscope
        3. Model should return function_call
        4. Execute function locally and provide output
        5. Model should generate final response using the function output
        """
        # 1. Define a list of callable tools for the model
        tools = [
            {
                "type": "function",
                "name": "get_horoscope",
                "description": "Get today's horoscope for an astrological sign.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "sign": {
                            "type": "string",
                            "description": "An astrological sign like Taurus or Aquarius",
                        },
                    },
                    "required": ["sign"],
                },
            },
        ]
        system_prompt = (
            "You are a helpful assistant that can call functions. "
            "When a user asks for horoscope information, call the function. "
            "IMPORTANT: Don't reply directly to the user, only call the function. "
        )

        # Create a running input list we will add to over time
        input_list = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "What is my horoscope? I am an Aquarius."},
        ]

        # 2. Prompt the model with tools defined
        resp = self.make_request(
            "/v1/responses",
            "POST",
            {
                "model": self.model,
                "tools": tools,
                "input": input_list,
            },
        )

        # Should successfully make the request
        self.assertEqual(resp.status_code, 200)

        data = resp.json()

        # Basic response structure
        self.assertIn("id", data)
        self.assertIn("status", data)
        self.assertEqual(data["status"], "completed")
        self.assertIn("output", data)

        # Verify output array is not empty
        output = data["output"]
        self.assertIsInstance(output, list)
        self.assertGreater(len(output), 0)

        # Check for function_call in output
        function_calls = [
            item for item in output if item.get("type") == "function_call"
        ]
        self.assertGreater(
            len(function_calls), 0, "Response should contain at least one function_call"
        )

        # Verify function_call structure
        function_call = function_calls[0]
        self.assertIn("call_id", function_call)
        self.assertIn("name", function_call)
        self.assertEqual(function_call["name"], "get_horoscope")
        self.assertIn("arguments", function_call)

        # Parse arguments
        args = json.loads(function_call["arguments"])
        self.assertIn("sign", args)
        self.assertEqual(args["sign"].lower(), "aquarius")

        # 3. Save function call outputs for subsequent requests
        input_list += output

        # 4. Execute the function logic for get_horoscope
        horoscope = f"{args['sign']}: Next Tuesday you will befriend a baby otter."

        # 5. Provide function call results to the model
        input_list.append(
            {
                "type": "function_call_output",
                "call_id": function_call["call_id"],
                "output": json.dumps({"horoscope": horoscope}),
            }
        )

        # 6. Make second request with function output
        resp2 = self.make_request(
            "/v1/responses",
            "POST",
            {
                "model": self.model,
                "instructions": "Respond only with a horoscope generated by a tool.",
                "tools": tools,
                "input": input_list,
            },
        )
        data2 = resp2.json()
        self.assertEqual(data2["status"], "completed")

        # The model should be able to give a response using the function output
        output2 = data2["output"]
        self.assertGreater(len(output2), 0)

        # Find message output
        messages = [item for item in output2 if item.get("type") == "message"]
        self.assertGreater(
            len(messages), 0, "Response should contain at least one message"
        )

        # Verify message contains the horoscope
        message = messages[0]
        self.assertIn("content", message)
        content_parts = message["content"]
        self.assertGreater(len(content_parts), 0)

        # Get text from content
        text_parts = [
            part.get("text", "")
            for part in content_parts
            if part.get("type") == "output_text"
        ]
        full_text = " ".join(text_parts).lower()

        # Should mention the horoscope or baby otter
        self.assertTrue(
            "baby otter" in full_text or "aquarius" in full_text,
            "Response should reference the horoscope content",
        )