-
Notifications
You must be signed in to change notification settings - Fork 111
feat: AgentCore tool search plugin for Strands Agents #494
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| """Gateway Strands plugins.""" | ||
|
|
||
| from .agentcore_tool_search import AgentCoreToolSearchPlugin | ||
|
|
||
| __all__ = ["AgentCoreToolSearchPlugin"] | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,138 @@ | ||
| # Strands AgentCore Tool Search Plugin | ||
|
|
||
| A semantic tool discovery plugin for [Strands Agents](https://github.com/strands-agents/sdk-python) that uses the [Amazon Bedrock AgentCore Gateway](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-using-mcp-semantic-search.html) `x_amz_bedrock_agentcore_search` tool. This enables agents to dynamically load only the relevant tools for each invocation by deriving user intent from conversation history, even when hundreds of tools are registered on the gateway. | ||
|
|
||
| ## Features | ||
|
|
||
| - **Semantic tool discovery** — uses AgentCore Gateway's built-in search to find relevant tools | ||
| - **Intent-based loading** — derives user intent via LLM before searching | ||
| - **No list_tools call** — tools are built directly from search results | ||
| - **Pluggable intent provider** — swap the default intent provider with your own | ||
| - **Agent model reuse** — by default, the intent classifier uses the same model as the parent agent | ||
|
|
||
| ## Installation | ||
|
|
||
| ```bash | ||
| pip install agentcore-tool-search-plugin | ||
| ``` | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This package doesn't exist — it's part of bedrock-agentcore. Same issue on line 25 with the import path. The examples further down (line 90+) have the correct import; these first ones will confuse people. |
||
|
|
||
| ## Usage | ||
|
|
||
| ```python | ||
| from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client | ||
| from strands import Agent | ||
| from strands.tools.mcp import MCPClient | ||
| from agentcore_tool_search_plugin import AgentCoreToolSearchPlugin | ||
|
|
||
| mcp_client = MCPClient(lambda: aws_iam_streamablehttp_client( | ||
| endpoint="https://<gateway-id>.gateway.bedrock-agentcore.<region>.amazonaws.com/mcp", | ||
| aws_region="us-east-1", | ||
| aws_service="bedrock-agentcore", | ||
| )) | ||
|
|
||
| mcp_client.start() | ||
|
|
||
| agent = Agent(plugins=[AgentCoreToolSearchPlugin(mcp_client=mcp_client)]) | ||
|
|
||
| agent("Find me afternoon flights to New York") | ||
| ``` | ||
|
|
||
| Or using a context manager: | ||
|
|
||
| ```python | ||
| with mcp_client: | ||
| agent = Agent(plugins=[AgentCoreToolSearchPlugin(mcp_client=mcp_client)]) | ||
| agent("Find me afternoon flights to New York") | ||
| ``` | ||
|
|
||
| ## How It Works | ||
|
|
||
|  | ||
|
|
||
| On each agent invocation: | ||
|
|
||
| 1. **User query** — The user sends a query to Strands agent. | ||
| 2. **Hook** — The agent triggers the `AgentCoreToolSearchPlugin` before model invocation | ||
| 3. **Derive intent** — The `IntentProvider` sends the last N messages from conversation history to the configured LLM to produce a concise intent string | ||
| 4. **Search gateway** — The intent is passed to AgentCore Gateway's `x_amz_bedrock_agentcore_search` tool to obtain most relevant tools. | ||
| 5. **Invoke LLM** — The agent invokes the LLM with the user query along with the matched tools from registered MCP targets (Lambda, API Gateway, MCP Server) | ||
|
|
||
| Previously loaded tools are cleared before each search, so the agent always has the most relevant tools available. | ||
|
|
||
| ## Intent Provider | ||
|
|
||
| An `IntentProvider` is responsible for analyzing conversation messages and producing a concise intent string that drives tool search. The plugin calls `derive_intent(messages, model)` before each invocation to determine what tools to load. | ||
|
|
||
| ### Default Intent Provider | ||
|
|
||
| `StrandsIntentProvider` uses an LLM to classify the last few conversation messages into a concise intent string. By default it uses the agent's model. | ||
|
|
||
| **Basic usage (uses the agent's model automatically):** | ||
|
|
||
| ```python | ||
| from bedrock_agentcore.gateway.integrations.strands.plugins import AgentCoreToolSearchPlugin | ||
|
|
||
| agent = Agent(plugins=[ | ||
| AgentCoreToolSearchPlugin(mcp_client=mcp_client) | ||
| ]) | ||
| ``` | ||
|
|
||
| **With a custom model for intent classification:** | ||
|
|
||
| ```python | ||
| from strands.models.bedrock import BedrockModel | ||
| from bedrock_agentcore.gateway.integrations.strands.plugins import AgentCoreToolSearchPlugin | ||
| from bedrock_agentcore.gateway.integrations.strands.plugins.agentcore_tool_search.intent_providers import StrandsIntentProvider | ||
|
|
||
| intent_model = BedrockModel(model_id="us.anthropic.claude-haiku-4-5-20251001-v1:0") | ||
| agent = Agent(plugins=[ | ||
| AgentCoreToolSearchPlugin( | ||
| mcp_client=mcp_client, | ||
| intent_provider=StrandsIntentProvider(model=intent_model), | ||
| ) | ||
| ]) | ||
| ``` | ||
|
|
||
| **With a custom system prompt:** | ||
|
|
||
| ```python | ||
| from bedrock_agentcore.gateway.integrations.strands.plugins import AgentCoreToolSearchPlugin | ||
| from bedrock_agentcore.gateway.integrations.strands.plugins.agentcore_tool_search.intent_providers import StrandsIntentProvider | ||
|
|
||
| agent = Agent(plugins=[ | ||
| AgentCoreToolSearchPlugin( | ||
| mcp_client=mcp_client, | ||
| intent_provider=StrandsIntentProvider( | ||
| system_prompt="Classify the user's intent in one sentence. Focus on the action, not details." | ||
| ), | ||
| ) | ||
| ]) | ||
| ``` | ||
|
|
||
| ### Custom Intent Provider | ||
|
|
||
| You can provide your own intent derivation strategy by subclassing `IntentProvider`: | ||
|
|
||
| ```python | ||
| from bedrock_agentcore.gateway.integrations.strands.plugins.agentcore_tool_search.intent_providers import IntentProvider | ||
|
|
||
| class MyIntentProvider(IntentProvider): | ||
| def derive_intent(self, messages: list[dict], model=None) -> str: | ||
| # custom logic to derive intent | ||
| return "intent string" | ||
|
|
||
| agent = Agent(plugins=[ | ||
| AgentCoreToolSearchPlugin( | ||
| mcp_client=mcp_client, | ||
| intent_provider=MyIntentProvider(), | ||
| ) | ||
| ]) | ||
| ``` | ||
|
|
||
| ## Prerequisites | ||
|
|
||
| - An AgentCore Gateway with **[semantic search](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-using-mcp-semantic-search.html) enabled** | ||
| - Tools registered on the gateway with descriptions | ||
| - AWS credentials with access to the gateway | ||
|
|
||
| For more details, see the [AgentCore Gateway Documentation](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-building.html). | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| """AgentCore Tool Search plugin for Strands Agents.""" | ||
|
|
||
| from .intent_providers import StrandsIntentProvider, IntentProvider | ||
| from .plugin import AgentCoreToolSearchPlugin | ||
|
|
||
| __all__ = ["AgentCoreToolSearchPlugin", "IntentProvider", "DefaultIntentProvider"] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| """Intent provider interfaces and implementations.""" | ||
|
|
||
| from .intent_provider import IntentProvider | ||
| from .strands_intent_provider import StrandsIntentProvider | ||
|
|
||
| __all__ = ["StrandsIntentProvider", "IntentProvider"] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| """Intent provider abstract interface.""" | ||
|
|
||
| from abc import ABC, abstractmethod | ||
|
|
||
|
|
||
| class IntentProvider(ABC): | ||
| """Abstract interface for deriving user intent from conversation messages. | ||
|
|
||
| Subclasses must implement the `derive_intent` method to analyze conversation | ||
| messages and return a concise intent string. | ||
| """ | ||
|
|
||
| @abstractmethod | ||
| def derive_intent(self, messages: list[dict], model=None) -> str: | ||
| """Analyze conversation messages and return a concise intent string. | ||
|
|
||
| Args: | ||
| messages: List of conversation message dicts in Strands format. | ||
| model: Optional model instance from the parent agent. Implementations | ||
| can use this for LLM-based intent derivation. | ||
|
|
||
| Returns: | ||
| A plain text string describing the user's intent. | ||
| Returns empty string if intent cannot be determined. | ||
| """ | ||
| ... |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,73 @@ | ||
| """Default LLM-based intent provider implementation.""" | ||
|
|
||
| import logging | ||
|
|
||
| from strands import Agent | ||
|
|
||
| from .intent_provider import IntentProvider | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| INTENT_SYSTEM_PROMPT = ( | ||
| "You are an intent classifier. Given the recent conversation messages, " | ||
| "produce a concise one-sentence description of what the user is trying to accomplish. " | ||
| "Focus on the type of task, not the specific details. " | ||
| "Reply with ONLY the intent description, nothing else." | ||
| ) | ||
|
Comment on lines
+11
to
+16
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. could this be prone to user manipulation? like
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. DefaultIntentProvider uses the agent's model (which can be configured with Bedrock Guardrails). Users can also pass a custom BedrockModel with guardrails attached via the constructor. Additionally, the intent classifier has a constrained system prompt, no tools, and only receives user-typed messages — so the worst case of prompt injection is a poor search query, not code execution or data leakage. Open to recommendations if you think we should add additional hardening here. |
||
|
|
||
|
|
||
| class StrandsIntentProvider(IntentProvider): | ||
| """LLM-based intent provider that classifies the last N messages.""" | ||
|
|
||
| def __init__(self, message_window: int = 5, model=None, system_prompt: str = INTENT_SYSTEM_PROMPT): | ||
| """Initialize StrandsIntentProvider. | ||
|
|
||
| Args: | ||
| message_window: Number of recent messages to consider. | ||
| model: Optional explicit model for intent classification. | ||
| system_prompt: System prompt for the intent classifier. Defaults to INTENT_SYSTEM_PROMPT. | ||
| """ | ||
| self._message_window = message_window | ||
| self._explicit_model = model | ||
| self._system_prompt = system_prompt | ||
|
|
||
| def derive_intent(self, messages: list[dict], model=None) -> str: | ||
| """Derive intent using an LLM. Falls back to agent's model if no explicit model set.""" | ||
| try: | ||
| recent_messages = messages[-self._message_window :] if messages else [] | ||
| if not recent_messages: | ||
| return "" | ||
|
|
||
| kwargs = {"system_prompt": self._system_prompt, "tools": []} | ||
| # Priority: explicit model > agent's model > Strands default | ||
| resolved_model = self._explicit_model or model | ||
| if resolved_model: | ||
| kwargs["model"] = resolved_model | ||
|
|
||
| intent_agent = Agent(**kwargs) | ||
| response = intent_agent(self._format_messages_for_prompt(recent_messages)) | ||
| return str(response).strip() | ||
|
senthilkumarmohan marked this conversation as resolved.
|
||
| except Exception as e: | ||
| logger.error("Failed to derive intent: %s", e) | ||
| return "" | ||
|
|
||
| def _format_messages_for_prompt(self, messages: list[dict]) -> str: | ||
| """Format user messages into a text prompt for the intent LLM. | ||
|
|
||
| Only includes user-role messages to avoid leaking PII or sensitive data | ||
| from tool results or assistant responses. | ||
| """ | ||
| parts = [] | ||
| for msg in messages: | ||
| role = msg.get("role", "") | ||
| if role != "user": | ||
| continue | ||
| content = msg.get("content", []) | ||
| text = "" | ||
| if isinstance(content, list): | ||
| text = " ".join( | ||
| block.get("text", "") for block in content if isinstance(block, dict) and "text" in block | ||
| ) | ||
| if text.strip(): | ||
| parts.append(text.strip()) | ||
| return "\n".join(parts) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,123 @@ | ||
| """AgentCore tool search plugin for Strands Agents.""" | ||
|
|
||
| import json | ||
| import logging | ||
|
|
||
| from mcp.types import Tool as MCPTool | ||
| from strands.hooks import BeforeInvocationEvent | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| from strands.plugins import Plugin, hook | ||
| from strands.tools.mcp import MCPClient | ||
| from strands.tools.mcp.mcp_agent_tool import MCPAgentTool | ||
|
|
||
| from .intent_providers import StrandsIntentProvider, IntentProvider | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class AgentCoreToolSearchPlugin(Plugin): | ||
| """Plugin that dynamically loads tools from AgentCore Gateway based on semantic intent. | ||
|
|
||
| Args: | ||
| mcp_client: MCPClient connected to an AgentCore Gateway. | ||
| intent_provider: Strategy for deriving intent. Defaults to DefaultIntentProvider. | ||
| """ | ||
|
|
||
| name = "agentcore-tool-search-plugin" | ||
|
|
||
| def __init__( | ||
| self, | ||
| mcp_client: MCPClient, | ||
| intent_provider: IntentProvider | None = None, | ||
| ): | ||
| """Initialize the plugin. | ||
|
|
||
| Args: | ||
| mcp_client: MCPClient connected to an AgentCore Gateway. | ||
| intent_provider: Strategy for deriving intent. Defaults to DefaultIntentProvider. | ||
| """ | ||
| super().__init__() | ||
| self._intent_provider = intent_provider or StrandsIntentProvider() | ||
| self._mcp_client = mcp_client | ||
| self._loaded_tool_names: set[str] = set() | ||
|
|
||
| @property | ||
| def tools(self): | ||
| """Return empty list; tools are loaded dynamically via the hook.""" | ||
| return [] | ||
|
|
||
| @hook | ||
| def on_before_invocation(self, event: BeforeInvocationEvent) -> None: | ||
| """Derive intent, search gateway, and load matching tools.""" | ||
| messages = event.messages or [] | ||
|
|
||
| # Pass the agent's model to the intent provider | ||
| intent = self._intent_provider.derive_intent(messages, model=event.agent.model) | ||
| logger.info("Derived intent: %s", intent) | ||
|
|
||
| # Clear all previously loaded dynamic tools | ||
| for name in list(self._loaded_tool_names): | ||
| event.agent.tool_registry.registry.pop(name, None) | ||
| self._loaded_tool_names.clear() | ||
|
|
||
| if not intent: | ||
| return | ||
|
|
||
| try: | ||
| result = self._mcp_client.call_tool_sync( | ||
| tool_use_id="intent-search", | ||
| name="x_amz_bedrock_agentcore_search", | ||
| arguments={"query": intent}, | ||
| ) | ||
| agent_tools = self._build_tools_from_search_result(result) | ||
| except Exception as e: | ||
| logger.error("AgentCore Gateway search failed: %s", e) | ||
| return | ||
|
|
||
| for agent_tool in agent_tools: | ||
| try: | ||
| # Skip if a non-dynamic tool with this name already exists | ||
| if ( | ||
| agent_tool.tool_name in event.agent.tool_registry.registry | ||
| and agent_tool.tool_name not in self._loaded_tool_names | ||
| ): | ||
| logger.debug("Skipping tool %s: already registered as a static tool", agent_tool.tool_name) | ||
| continue | ||
| event.agent.tool_registry.register_tool(agent_tool) | ||
| self._loaded_tool_names.add(agent_tool.tool_name) | ||
| except Exception as e: | ||
| logger.error("Failed to register tool %s: %s", agent_tool.tool_name, e) | ||
|
|
||
| logger.info("Loaded tools: %s", self._loaded_tool_names) | ||
|
senthilkumarmohan marked this conversation as resolved.
|
||
|
|
||
| def _build_tools_from_search_result(self, result) -> list[MCPAgentTool]: | ||
| """Build MCPAgentTool objects from the gateway search response.""" | ||
| tools = [] | ||
| if not result or not isinstance(result, dict): | ||
| return tools | ||
|
|
||
| tool_defs = [] | ||
| structured = result.get("structuredContent") | ||
| if isinstance(structured, dict) and "tools" in structured: | ||
| tool_defs = structured["tools"] | ||
| else: | ||
| for block in result.get("content", []): | ||
| if isinstance(block, dict) and "text" in block: | ||
| try: | ||
| data = json.loads(block["text"]) | ||
| if isinstance(data, dict) and "tools" in data: | ||
| tool_defs = data["tools"] | ||
| break | ||
| except (json.JSONDecodeError, TypeError): | ||
| continue | ||
|
|
||
| for tool_def in tool_defs: | ||
| if not isinstance(tool_def, dict) or "name" not in tool_def: | ||
| continue | ||
| mcp_tool = MCPTool( | ||
| name=tool_def["name"], | ||
| description=tool_def.get("description", ""), | ||
| inputSchema=tool_def.get("inputSchema", {"type": "object", "properties": {}}), | ||
| ) | ||
| tools.append(MCPAgentTool(mcp_tool=mcp_tool, mcp_client=self._mcp_client)) | ||
|
|
||
| return tools | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Missing
__init__.pyatgateway/integrations/andgateway/integrations/strands/. Without those, this won't be importable from an installed wheel since pyproject.toml usespackages = ["src/bedrock_agentcore"].