wehub-resource-sync
274 lines
10 KiB
Python
274 lines
10 KiB
Python
import unittest
|
|
from pydantic import BaseModel
|
|
from typing import Dict, Literal, Optional
|
|
from aisuite.utils.tools import Tools # Import your ToolManager class
|
|
from enum import Enum
|
|
|
|
|
|
# Define a sample tool function and Pydantic model for testing
|
|
class TemperatureUnit(str, Enum):
|
|
CELSIUS = "Celsius"
|
|
FAHRENHEIT = "Fahrenheit"
|
|
|
|
|
|
class TemperatureParamsV2(BaseModel):
|
|
location: str
|
|
unit: TemperatureUnit = TemperatureUnit.CELSIUS
|
|
|
|
|
|
class TemperatureParams(BaseModel):
|
|
location: str
|
|
unit: str = "Celsius"
|
|
|
|
|
|
def get_current_temperature(location: str, unit: str = "Celsius") -> Dict[str, str]:
|
|
"""Gets the current temperature for a specific location and unit."""
|
|
return {"location": location, "unit": unit, "temperature": "72"}
|
|
|
|
|
|
def missing_annotation_tool(location, unit="Celsius"):
|
|
"""Tool function without type annotations."""
|
|
return {"location": location, "unit": unit, "temperature": "72"}
|
|
|
|
|
|
def get_current_temperature_v2(
|
|
location: str, unit: TemperatureUnit = TemperatureUnit.CELSIUS
|
|
) -> Dict[str, str]:
|
|
"""Gets the current temperature for a specific location and unit (with enum support)."""
|
|
return {"location": location, "unit": unit, "temperature": "72"}
|
|
|
|
|
|
def analyze_items(
|
|
tags: list[str],
|
|
scores: dict[str, int],
|
|
mode: Literal["fast", "thorough"] = "fast",
|
|
note: Optional[str] = None,
|
|
limit: int | None = None,
|
|
) -> Dict[str, object]:
|
|
"""Analyze items using complex argument types."""
|
|
return {
|
|
"tags": tags,
|
|
"scores": scores,
|
|
"mode": mode,
|
|
"note": note,
|
|
"limit": limit,
|
|
}
|
|
|
|
|
|
class TestToolManager(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tool_manager = Tools()
|
|
|
|
def test_add_tool_with_pydantic_model(self):
|
|
"""Test adding a tool with an explicit Pydantic model."""
|
|
self.tool_manager._add_tool(get_current_temperature, TemperatureParams)
|
|
|
|
expected_tool_spec = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_temperature",
|
|
"description": "Gets the current temperature for a specific location and unit.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "",
|
|
},
|
|
"unit": {
|
|
"type": "string",
|
|
"description": "",
|
|
"default": "Celsius",
|
|
},
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
tools = self.tool_manager.tools()
|
|
self.assertIn(
|
|
"get_current_temperature", [tool["function"]["name"] for tool in tools]
|
|
)
|
|
assert (
|
|
tools == expected_tool_spec
|
|
), f"Expected {expected_tool_spec}, but got {tools}"
|
|
|
|
def test_add_tool_with_signature_inference(self):
|
|
"""Test adding a tool and inferring parameters from the function signature."""
|
|
self.tool_manager._add_tool(get_current_temperature)
|
|
# Expected output from tool_manager.tools() when called with OpenAI format
|
|
expected_tool_spec = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_temperature",
|
|
"description": "Gets the current temperature for a specific location and unit.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "", # No description provided in function signature
|
|
},
|
|
"unit": {
|
|
"type": "string",
|
|
"description": "",
|
|
"default": "Celsius",
|
|
},
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
tools = self.tool_manager.tools()
|
|
print(tools)
|
|
self.assertIn(
|
|
"get_current_temperature", [tool["function"]["name"] for tool in tools]
|
|
)
|
|
assert (
|
|
tools == expected_tool_spec
|
|
), f"Expected {expected_tool_spec}, but got {tools}"
|
|
|
|
def test_add_tool_missing_annotation_raises_exception(self):
|
|
"""Test that adding a tool with missing type annotations raises a TypeError."""
|
|
with self.assertRaises(TypeError):
|
|
self.tool_manager._add_tool(missing_annotation_tool)
|
|
|
|
def test_execute_tool_valid_parameters(self):
|
|
"""Test executing a registered tool with valid parameters."""
|
|
self.tool_manager._add_tool(get_current_temperature, TemperatureParams)
|
|
tool_call = {
|
|
"id": "call_1",
|
|
"function": {
|
|
"name": "get_current_temperature",
|
|
"arguments": {"location": "San Francisco", "unit": "Celsius"},
|
|
},
|
|
}
|
|
result, result_message = self.tool_manager.execute_tool(tool_call)
|
|
|
|
# Assuming result is returned as a list with a single dictionary
|
|
result_dict = result[0] if isinstance(result, list) else result
|
|
|
|
# Check that the result matches expected output
|
|
self.assertEqual(result_dict["location"], "San Francisco")
|
|
self.assertEqual(result_dict["unit"], "Celsius")
|
|
self.assertEqual(result_dict["temperature"], "72")
|
|
|
|
def test_execute_tool_invalid_parameters(self):
|
|
"""Test that executing a tool with invalid parameters raises a ValueError."""
|
|
self.tool_manager._add_tool(get_current_temperature, TemperatureParams)
|
|
tool_call = {
|
|
"id": "call_1",
|
|
"function": {
|
|
"name": "get_current_temperature",
|
|
"arguments": {"location": 123}, # Invalid type for location
|
|
},
|
|
}
|
|
|
|
with self.assertRaises(ValueError) as context:
|
|
self.tool_manager.execute_tool(tool_call)
|
|
|
|
# Verify the error message contains information about the validation error
|
|
self.assertIn(
|
|
"Error in tool 'get_current_temperature' parameters", str(context.exception)
|
|
)
|
|
|
|
def test_add_tool_with_enum(self):
|
|
"""Test adding a tool with an enum parameter."""
|
|
self.tool_manager._add_tool(get_current_temperature_v2, TemperatureParamsV2)
|
|
|
|
expected_tool_spec = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_temperature_v2",
|
|
"description": "Gets the current temperature for a specific location and unit (with enum support).",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "",
|
|
},
|
|
"unit": {
|
|
"type": "string",
|
|
"enum": ["Celsius", "Fahrenheit"],
|
|
"description": "",
|
|
"default": "Celsius",
|
|
},
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
tools = self.tool_manager.tools()
|
|
assert (
|
|
tools == expected_tool_spec
|
|
), f"Expected {expected_tool_spec}, but got {tools}"
|
|
|
|
def test_add_tool_with_complex_annotations_uses_json_schema_types(self):
|
|
"""Test complex annotations generate valid JSON Schema types."""
|
|
self.tool_manager._add_tool(analyze_items)
|
|
|
|
parameters = self.tool_manager.tools()[0]["function"]["parameters"]
|
|
properties = parameters["properties"]
|
|
|
|
self.assertEqual(properties["tags"]["type"], "array")
|
|
self.assertEqual(properties["tags"]["items"], {"type": "string"})
|
|
self.assertEqual(properties["scores"]["type"], "object")
|
|
self.assertEqual(
|
|
properties["scores"]["additionalProperties"], {"type": "integer"}
|
|
)
|
|
self.assertEqual(properties["mode"]["type"], "string")
|
|
self.assertEqual(properties["mode"]["enum"], ["fast", "thorough"])
|
|
self.assertEqual(properties["mode"]["default"], "fast")
|
|
self.assertEqual(properties["note"]["type"], "string")
|
|
self.assertIsNone(properties["note"]["default"])
|
|
self.assertEqual(properties["limit"]["type"], "integer")
|
|
self.assertIsNone(properties["limit"]["default"])
|
|
self.assertEqual(parameters["required"], ["tags", "scores"])
|
|
|
|
for schema in properties.values():
|
|
self.assertNotIn("list[", str(schema))
|
|
self.assertNotIn("dict[", str(schema))
|
|
self.assertNotIn("| None", str(schema))
|
|
|
|
def test_execute_tool_with_complex_annotations(self):
|
|
"""Test Pydantic validation still executes complex-typed tools."""
|
|
self.tool_manager._add_tool(analyze_items)
|
|
tool_call = {
|
|
"id": "call_1",
|
|
"function": {
|
|
"name": "analyze_items",
|
|
"arguments": {
|
|
"tags": ["urgent", "finance"],
|
|
"scores": {"urgent": 10, "finance": 7},
|
|
"mode": "thorough",
|
|
"limit": 3,
|
|
},
|
|
},
|
|
}
|
|
|
|
result, _ = self.tool_manager.execute_tool(tool_call)
|
|
|
|
self.assertEqual(
|
|
result[0],
|
|
{
|
|
"tags": ["urgent", "finance"],
|
|
"scores": {"urgent": 10, "finance": 7},
|
|
"mode": "thorough",
|
|
"note": None,
|
|
"limit": 3,
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|