diff --git a/aisuite/__init__.py b/aisuite/__init__.py index 92e57f18..daa0ae56 100644 --- a/aisuite/__init__.py +++ b/aisuite/__init__.py @@ -1,2 +1,3 @@ from .client import Client from .framework.message import Message +from .utils.tool_manager import ToolManager diff --git a/aisuite/utils/tool_manager.py b/aisuite/utils/tool_manager.py index f12ca343..af88771f 100644 --- a/aisuite/utils/tool_manager.py +++ b/aisuite/utils/tool_manager.py @@ -5,11 +5,14 @@ class ToolManager: - def __init__(self): + def __init__(self, tools: list[Callable] = None): self._tools = {} + if tools: + for tool in tools: + self._add_tool(tool) # Add a tool function with or without a Pydantic model. - def add_tool(self, func: Callable, param_model: Optional[Type[BaseModel]] = None): + def _add_tool(self, func: Callable, param_model: Optional[Type[BaseModel]] = None): """Register a tool function with metadata. If no param_model is provided, infer from function signature.""" if param_model: tool_spec = self._convert_to_tool_spec(func, param_model) diff --git a/examples/simple_tool_calling.ipynb b/examples/simple_tool_calling.ipynb index 294d1881..bea53604 100644 --- a/examples/simple_tool_calling.ipynb +++ b/examples/simple_tool_calling.ipynb @@ -34,10 +34,9 @@ "outputs": [], "source": [ "import aisuite as ai\n", - "from aisuite.utils.tool_manager import ToolManager # Import your ToolManager class\n", + "from aisuite import ToolManager # Import your ToolManager class\n", "\n", - "client = ai.Client()\n", - "tool_manager = ToolManager()" + "client = ai.Client()" ] }, { @@ -49,10 +48,7 @@ "# Mock tool functions.\n", "def get_current_temperature(location: str, unit: str):\n", " # Simulate fetching temperature from an API\n", - " return {\"location\": location, \"unit\": unit, \"temperature\": 72}\n", - "\n", - "tool_manager.add_tool(get_current_temperature)\n", - "messages = [{\"role\": \"user\", \"content\": \"What is the current temperature in San Francisco in Celsius?\"}]" + " return {\"location\": location, \"unit\": unit, \"temperature\": 72}" ] }, { @@ -61,13 +57,18 @@ "metadata": {}, "outputs": [], "source": [ - "# model = \"anthropic:claude-3-5-sonnet-20240620\"\n", - "model = \"openai:gpt-4o\"\n", + "model = \"anthropic:claude-3-5-sonnet-20240620\"\n", + "# model = \"openai:gpt-4o\"\n", "# model = mistral:mistral-large-latest\n", "# model = \"aws:anthropic.claude-3-haiku-20240307-v1:0\"\n", "# model = \"aws:meta.llama3-1-8b-instruct-v1:0\"\n", "# model = \"aws:meta.llama3-3-70b-instruct-v1:0\"\n", "# model = \"groq:llama-3.1-70b-versatile\"\n", + "\n", + "\n", + "tool_manager = ToolManager([get_current_temperature])\n", + "messages = [{\"role\": \"user\", \"content\": \"What is the current temperature in San Francisco in Celsius?\"}]\n", + "\n", "response = client.chat.completions.create(\n", " model=model, messages=messages, tools=tool_manager.tools())" ] @@ -81,11 +82,12 @@ "name": "stdout", "output_type": "stream", "text": [ - "{'content': None,\n", - " 'function_call': None,\n", + "{'content': 'To answer your question about the current temperature in San '\n", + " \"Francisco in Celsius, I'll need to use the available tool to get \"\n", + " 'that information. Let me fetch that for you.',\n", " 'refusal': None,\n", " 'role': 'assistant',\n", - " 'tool_calls': [ChatCompletionMessageToolCall(id='call_uzJz9BtPDbcJYiKLcrmEybfE', function=Function(arguments='{\"location\":\"San Francisco\",\"unit\":\"Celsius\"}', name='get_current_temperature'), type='function')]}\n" + " 'tool_calls': [ChatCompletionMessageToolCall(id='toolu_01SVEv2QJ7tsUatecDD2TEq7', function=Function(arguments='{\"location\": \"San Francisco\", \"unit\": \"Celsius\"}', name='get_current_temperature'), type='function')]}\n" ] } ], @@ -96,14 +98,18 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "The current temperature in San Francisco is 72°C. \n" + "Based on the function results, the current temperature in San Francisco is 72 degrees Celsius.\n", + "\n", + "It's worth noting that this temperature seems unusually high for San Francisco, as 72°C is equivalent to about 161.6°F, which would be extremely hot for any city. Typically, San Francisco experiences much milder temperatures. This result might be due to an error in the data or the conversion process. \n", + "\n", + "To give you a more realistic perspective, San Francisco usually has moderate temperatures year-round, with average highs rarely exceeding 21°C (70°F) even in the warmest months. If you're planning a trip or need accurate current weather information, I'd recommend double-checking this data with a reliable weather service or asking for a verification of the temperature reading.\n" ] } ], @@ -120,19 +126,11 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "The capital of California is Sacramento.\n" - ] - } - ], + "outputs": [], "source": [ - "# Now, test without tool calling.\n", + "# Now, test without tool calling, to check that the normal path is not broken.\n", "messages = [{\"role\": \"user\", \"content\": \"What is the capital of California?\"}]\n", "response = client.chat.completions.create(\n", " model=model, messages=messages)\n",