Skip to content
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions src/strands/agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
from pydantic import BaseModel

from .. import _identifier
from ..event_loop.event_loop import event_loop_cycle, run_tool
from ..event_loop.event_loop import event_loop_cycle
from ..handlers.callback_handler import PrintingCallbackHandler, null_callback_handler
from ..hooks import (
AfterInvocationEvent,
Expand All @@ -35,6 +35,8 @@
from ..session.session_manager import SessionManager
from ..telemetry.metrics import EventLoopMetrics
from ..telemetry.tracer import get_tracer, serialize
from ..tools.executors import ConcurrentToolExecutor
from ..tools.executors._executor import Executor as ToolExecutor
from ..tools.registry import ToolRegistry
from ..tools.watcher import ToolWatcher
from ..types.content import ContentBlock, Message, Messages
Expand Down Expand Up @@ -136,13 +138,14 @@ def caller(
"name": normalized_name,
"input": kwargs.copy(),
}
tool_results: list[ToolResult] = []
invocation_state = kwargs

async def acall() -> ToolResult:
# Pass kwargs as invocation_state
async for event in run_tool(self._agent, tool_use, kwargs):
async for event in ToolExecutor._stream(self._agent, tool_use, tool_results, invocation_state):
_ = event

return cast(ToolResult, event)
return tool_results[0]

def tcall() -> ToolResult:
return asyncio.run(acall())
Expand Down Expand Up @@ -208,6 +211,7 @@ def __init__(
state: Optional[Union[AgentState, dict]] = None,
hooks: Optional[list[HookProvider]] = None,
session_manager: Optional[SessionManager] = None,
tool_executor: Optional[ToolExecutor] = None,
):
"""Initialize the Agent with the specified configuration.

Expand Down Expand Up @@ -250,6 +254,7 @@ def __init__(
Defaults to None.
session_manager: Manager for handling agent sessions including conversation history and state.
If provided, enables session-based persistence and state management.
tool_executor: Definition of tool execution stragety (e.g., sequential, concurrent, etc.).

Raises:
ValueError: If agent id contains path separators.
Expand Down Expand Up @@ -324,6 +329,8 @@ def __init__(
if self._session_manager:
self.hooks.add_hook(self._session_manager)

self.tool_executor = tool_executor or ConcurrentToolExecutor()

if hooks:
for hook in hooks:
self.hooks.add_hook(hook)
Expand Down
155 changes: 13 additions & 142 deletions src/strands/event_loop/event_loop.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,22 +11,20 @@
import logging
import time
import uuid
from typing import TYPE_CHECKING, Any, AsyncGenerator, cast
from typing import TYPE_CHECKING, Any, AsyncGenerator

from opentelemetry import trace as trace_api

from ..experimental.hooks import (
AfterModelInvocationEvent,
AfterToolInvocationEvent,
BeforeModelInvocationEvent,
BeforeToolInvocationEvent,
)
from ..hooks import (
MessageAddedEvent,
)
from ..telemetry.metrics import Trace
from ..telemetry.tracer import get_tracer
from ..tools.executor import run_tools, validate_and_prepare_tools
from ..tools._validator import validate_and_prepare_tools
from ..types.content import Message
from ..types.exceptions import (
ContextWindowOverflowException,
Expand All @@ -35,7 +33,7 @@
ModelThrottledException,
)
from ..types.streaming import Metrics, StopReason
from ..types.tools import ToolChoice, ToolChoiceAuto, ToolConfig, ToolGenerator, ToolResult, ToolUse
from ..types.tools import ToolResult, ToolUse
from ._recover_message_on_max_tokens_reached import recover_message_on_max_tokens_reached
from .streaming import stream_messages

Expand Down Expand Up @@ -212,7 +210,7 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
if stop_reason == "max_tokens":
"""
Handle max_tokens limit reached by the model.

When the model reaches its maximum token limit, this represents a potentially unrecoverable
state where the model's response was truncated. By default, Strands fails hard with an
MaxTokensReachedException to maintain consistency with other failure types.
Expand Down Expand Up @@ -306,122 +304,6 @@ async def recurse_event_loop(agent: "Agent", invocation_state: dict[str, Any]) -
recursive_trace.end()


async def run_tool(agent: "Agent", tool_use: ToolUse, invocation_state: dict[str, Any]) -> ToolGenerator:
"""Process a tool invocation.

Looks up the tool in the registry and streams it with the provided parameters.

Args:
agent: The agent for which the tool is being executed.
tool_use: The tool object to process, containing name and parameters.
invocation_state: Context for the tool invocation, including agent state.

Yields:
Tool events with the last being the tool result.
"""
logger.debug("tool_use=<%s> | streaming", tool_use)
tool_name = tool_use["name"]

# Get the tool info
tool_info = agent.tool_registry.dynamic_tools.get(tool_name)
tool_func = tool_info if tool_info is not None else agent.tool_registry.registry.get(tool_name)

# Add standard arguments to invocation_state for Python tools
invocation_state.update(
{
"model": agent.model,
"system_prompt": agent.system_prompt,
"messages": agent.messages,
"tool_config": ToolConfig( # for backwards compatability
tools=[{"toolSpec": tool_spec} for tool_spec in agent.tool_registry.get_all_tool_specs()],
toolChoice=cast(ToolChoice, {"auto": ToolChoiceAuto()}),
),
}
)

before_event = agent.hooks.invoke_callbacks(
BeforeToolInvocationEvent(
agent=agent,
selected_tool=tool_func,
tool_use=tool_use,
invocation_state=invocation_state,
)
)

try:
selected_tool = before_event.selected_tool
tool_use = before_event.tool_use
invocation_state = before_event.invocation_state # Get potentially modified invocation_state from hook

# Check if tool exists
if not selected_tool:
if tool_func == selected_tool:
logger.error(
"tool_name=<%s>, available_tools=<%s> | tool not found in registry",
tool_name,
list(agent.tool_registry.registry.keys()),
)
else:
logger.debug(
"tool_name=<%s>, tool_use_id=<%s> | a hook resulted in a non-existing tool call",
tool_name,
str(tool_use.get("toolUseId")),
)

result: ToolResult = {
"toolUseId": str(tool_use.get("toolUseId")),
"status": "error",
"content": [{"text": f"Unknown tool: {tool_name}"}],
}
# for every Before event call, we need to have an AfterEvent call
after_event = agent.hooks.invoke_callbacks(
AfterToolInvocationEvent(
agent=agent,
selected_tool=selected_tool,
tool_use=tool_use,
invocation_state=invocation_state, # Keep as invocation_state for backward compatibility with hooks
result=result,
)
)
yield after_event.result
return

async for event in selected_tool.stream(tool_use, invocation_state):
yield event

result = event

after_event = agent.hooks.invoke_callbacks(
AfterToolInvocationEvent(
agent=agent,
selected_tool=selected_tool,
tool_use=tool_use,
invocation_state=invocation_state, # Keep as invocation_state for backward compatibility with hooks
result=result,
)
)
yield after_event.result

except Exception as e:
logger.exception("tool_name=<%s> | failed to process tool", tool_name)
error_result: ToolResult = {
"toolUseId": str(tool_use.get("toolUseId")),
"status": "error",
"content": [{"text": f"Error: {str(e)}"}],
}
after_event = agent.hooks.invoke_callbacks(
AfterToolInvocationEvent(
agent=agent,
selected_tool=selected_tool,
tool_use=tool_use,
invocation_state=invocation_state, # Keep as invocation_state for backward compatibility with hooks
result=error_result,
exception=e,
)
)
yield after_event.result


async def _handle_tool_execution(
stop_reason: StopReason,
message: Message,
Expand All @@ -431,18 +313,12 @@ async def _handle_tool_execution(
cycle_start_time: float,
invocation_state: dict[str, Any],
) -> AsyncGenerator[dict[str, Any], None]:
tool_uses: list[ToolUse] = []
tool_results: list[ToolResult] = []
invalid_tool_use_ids: list[str] = []

"""
Handles the execution of tools requested by the model during an event loop cycle.
"""Handles the execution of tools requested by the model during an event loop cycle.

Args:
stop_reason: The reason the model stopped generating.
message: The message from the model that may contain tool use requests.
event_loop_metrics: Metrics tracking object for the event loop.
event_loop_parent_span: Span for the parent of this event loop.
agent: Agent for which tools are being executed.
cycle_trace: Trace object for the current event loop cycle.
cycle_span: Span object for tracing the cycle (type may vary).
cycle_start_time: Start time of the current cycle.
Expand All @@ -456,23 +332,18 @@ async def _handle_tool_execution(
- The updated event loop metrics,
- The updated request state.
"""
validate_and_prepare_tools(message, tool_uses, tool_results, invalid_tool_use_ids)
tool_uses: list[ToolUse] = []
tool_results: list[ToolResult] = []
invalid_tool_use_ids: list[str] = []

validate_and_prepare_tools(message, tool_uses, tool_results, invalid_tool_use_ids)
tool_uses = [tool_use for tool_use in tool_uses if tool_use.get("toolUseId") not in invalid_tool_use_ids]
if not tool_uses:
yield {"stop": (stop_reason, message, agent.event_loop_metrics, invocation_state["request_state"])}
return

def tool_handler(tool_use: ToolUse) -> ToolGenerator:
return run_tool(agent, tool_use, invocation_state)

tool_events = run_tools(
handler=tool_handler,
tool_uses=tool_uses,
event_loop_metrics=agent.event_loop_metrics,
invalid_tool_use_ids=invalid_tool_use_ids,
tool_results=tool_results,
cycle_trace=cycle_trace,
parent_span=cycle_span,
tool_events = agent.tool_executor._execute(
agent, tool_uses, tool_results, cycle_trace, cycle_span, invocation_state
)
async for tool_event in tool_events:
yield tool_event
Expand Down
45 changes: 45 additions & 0 deletions src/strands/tools/_validator.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
"""Tool validation utilities."""

from ..tools.tools import InvalidToolUseNameException, validate_tool_use
from ..types.content import Message
from ..types.tools import ToolResult, ToolUse


def validate_and_prepare_tools(
message: Message,
tool_uses: list[ToolUse],
tool_results: list[ToolResult],
invalid_tool_use_ids: list[str],
) -> None:
"""Validate tool uses and prepare them for execution.

Args:
message: Current message.
tool_uses: List to populate with tool uses.
tool_results: List to populate with tool results for invalid tools.
invalid_tool_use_ids: List to populate with invalid tool use IDs.
"""
# Extract tool uses from message
for content in message["content"]:
if isinstance(content, dict) and "toolUse" in content:
tool_uses.append(content["toolUse"])

# Validate tool uses
# Avoid modifying original `tool_uses` variable during iteration
tool_uses_copy = tool_uses.copy()
for tool in tool_uses_copy:
try:
validate_tool_use(tool)
except InvalidToolUseNameException as e:
# Replace the invalid toolUse name and return invalid name error as ToolResult to the LLM as context
tool_uses.remove(tool)
tool["name"] = "INVALID_TOOL_NAME"
invalid_tool_use_ids.append(tool["toolUseId"])
tool_uses.append(tool)
tool_results.append(
{
"toolUseId": tool["toolUseId"],
"status": "error",
"content": [{"text": f"Error: {str(e)}"}],
}
)
35 changes: 30 additions & 5 deletions src/strands/tools/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ def my_tool(param1: str, param2: int = 42) -> dict:
import functools
import inspect
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import (
Any,
Callable,
Expand Down Expand Up @@ -450,11 +451,7 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw
# Inject special framework-provided parameters
self._metadata.inject_special_parameters(validated_input, tool_use, invocation_state)

# "Too few arguments" expected, hence the type ignore
if inspect.iscoroutinefunction(self._tool_func):
result = await self._tool_func(**validated_input) # type: ignore
else:
result = await asyncio.to_thread(self._tool_func, **validated_input) # type: ignore
result = await self._stream(validated_input, **kwargs)

# FORMAT THE RESULT for Strands Agent
if isinstance(result, dict) and "status" in result and "content" in result:
Expand Down Expand Up @@ -488,6 +485,34 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw
"content": [{"text": f"Error: {error_type} - {error_msg}"}],
}

async def _stream(self, validated_input: dict[str, Any], **kwargs: Any) -> Any:
"""Execute the tool function based on type.

Args:
validated_input: Validated input parameters for the tool function.
**kwargs: Additional keyword arguments for future extensibility.

Returns:
The result of the tool function execution.
"""
# "Too few arguments" expected, hence the type ignores

if inspect.iscoroutinefunction(self._tool_func):
return await self._tool_func(**validated_input) # type: ignore

thread_pool = kwargs.get("thread_pool")

if isinstance(thread_pool, ThreadPoolExecutor):
return await asyncio.get_event_loop().run_in_executor(
thread_pool,
lambda: self._tool_func(**validated_input), # type: ignore
)

if thread_pool == "asyncio":
return await asyncio.to_thread(self._tool_func, **validated_input) # type: ignore

return self._tool_func(**validated_input) # type:ignore

@property
def supports_hot_reload(self) -> bool:
"""Check if this tool supports automatic reloading when modified.
Expand Down
Loading
Loading