Tool Execution

1
from smallestai.atoms.agent.tools import ToolRegistry, function_tool
2
3
class MyAgent(OutputAgentNode):
4
    def __init__(self):
5
        super().__init__(name="my-agent")
6
        self.llm = OpenAIClient(model="gpt-4o-mini")
7
        
8
        # Create registry and discover tools
9
        self.tool_registry = ToolRegistry()
10
        self.tool_registry.discover(self)
11
    
12
    @function_tool()
13
    def get_weather(self, location: str):
14
        """Get weather for a location."""
15
        return {"temp": 72, "conditions": "Sunny"}

1
response = await self.llm.chat(
2
    messages=self.context.messages,
3
    stream=True,
4
    tools=self.tool_registry.get_schemas()
5
)

1
from typing import List
2
from smallestai.atoms.agent.clients.types import ToolCall
3
4
async def generate_response(self):
5
    response = await self.llm.chat(
6
        messages=self.context.messages,
7
        stream=True,
8
        tools=self.tool_registry.get_schemas()
9
    )
10
    
11
    tool_calls: List[ToolCall] = []
12
    
13
    # Stream response and collect tool calls
14
    async for chunk in response:
15
        if chunk.content:
16
            yield chunk.content
17
        if chunk.tool_calls:
18
            tool_calls.extend(chunk.tool_calls)
19
    
20
    # Execute tools if any were called
21
    if tool_calls:
22
        results = await self.tool_registry.execute(
23
            tool_calls=tool_calls,
24
            parallel=True
25
        )
26
        
27
        # Add to context and get final response
28
        # (see below)

1
if tool_calls:
2
    results = await self.tool_registry.execute(tool_calls=tool_calls, parallel=True)
3
    
4
    # Add tool calls and results to context
5
    self.context.add_messages([
6
        {
7
            "role": "assistant",
8
            "content": "",
9
            "tool_calls": [
10
                {
11
                    "id": tc.id,
12
                    "type": "function",
13
                    "function": {
14
                        "name": tc.name,
15
                        "arguments": str(tc.arguments),
16
                    },
17
                }
18
                for tc in tool_calls
19
            ],
20
        },
21
        *[
22
            {"role": "tool", "tool_call_id": tc.id, "content": result.content}
23
            for tc, result in zip(tool_calls, results)
24
        ],
25
    ])
26
    
27
    # Call LLM again with tool results
28
    final_response = await self.llm.chat(
29
        messages=self.context.messages,
30
        stream=True
31
    )
32
    
33
    async for chunk in final_response:
34
        if chunk.content:
35
            yield chunk.content

1
# Parallel (faster)
2
results = await self.tool_registry.execute(tool_calls=tool_calls, parallel=True)
3
4
# Sequential (if dependencies exist)
5
results = await self.tool_registry.execute(tool_calls=tool_calls, parallel=False)