feat: tool calling and Agent primitive#2059
Closed
ValbuenaVC wants to merge 1 commit into
Closed
Conversation
- S3.1: Add Tool + ToolCallDispatch protocol (pyrit/agent/tools.py) Tool is a Pydantic model with name/description/json_schema/callable fields. ToolCallDispatch is a Protocol satisfied by any dispatcher backend. ToolCall/ToolResult are type aliases mirroring the function_call shape. - S3.2: Add InProcessRuntime (pyrit/agent/runtime.py) Default ToolCallDispatch backend executing Python callables in-process. Unknown tools, exceptions, and malformed args return structured error dicts. - S3.3: Extend TargetCapabilities with tool-usage fields Add supports_tool_usage: bool = False and tool_usage_schema: ToolUsageSchema | None = None. Add TOOL_USAGE = 'supports_tool_usage' to CapabilityName enum. Update _permissive_configuration to include supports_tool_usage=True. Export ToolUsageSchema from pyrit.models. - S3.4: Add Agent(PromptTarget) (pyrit/agent/agent.py) Wraps any PromptTarget and executes tool-call/result loop. Calls inner target's _send_prompt_to_target_async directly to keep interim tool turns in-context without memory duplication. max_tool_iterations guard prevents infinite loops. - S3.5: Attack interface parity tests Agent IS-A PromptTarget; attacks accept it with no signature changes. Tests verify PromptSendingAttack runs end-to-end against Agent. - S3.6: Exports, linting, coverage pyrit/agent/__init__.py exports Agent, InProcessRuntime, Tool, ToolCall, ToolCallDispatch, ToolResult. All pre-commit hooks pass; coverage 96% for new code (gate: 78%). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements Story 3 — Tool Calling and Agent Primitive. Adds a new
pyrit.agentpackage that gives anyPromptTargetuniversal tool-calling capability without modifying existing targets or attacks.Changes
New package:
pyrit/agent/tools.py(S3.1) — Core tool primitives:Tool: Pydantic model withname,description,json_schema,callablefields; hashable by name so tools can live in sets.ToolCall/ToolResult: Type aliases mirroring thefunction_callshape fromOpenAIResponseTarget.ToolCallDispatch: Protocol satisfied by any dispatcher backend (structural subtyping).runtime.py(S3.2) —InProcessRuntimeimplementingToolCallDispatch:OpenAIResponseTarget._execute_call_section_async.agent.py(S3.4) —Agent(PromptTarget):PromptTargetand executes the tool-call/result loop._send_prompt_to_target_asyncdirectly to keep interim tool turns in-context without memory duplication (same pattern asOpenAIResponseTarget).max_tool_iterationsguard (default 10) prevents infinite loops.AgentIS-APromptTarget, attacks and scenarios accept it with no changes.Modified:
pyrit/models/target_capabilities.py(S3.3)ToolUsageSchema: New frozen Pydantic model specifying which data types signal tool calls/results.CapabilityName.TOOL_USAGE = "supports_tool_usage": Points at the bool field soincludes()works correctly.TargetCapabilities: Two new fields —supports_tool_usage: bool = Falseandtool_usage_schema: ToolUsageSchema | None = None. Defaults keep all existing targets unchanged.Modified:
pyrit/prompt_target/common/discover_target_capabilities.pysupports_tool_usage=Trueto_permissive_configurationso it declares every boolean capability as supported.Tests
tests/unit/agent/test_tools.pytests/unit/agent/test_runtime.pytests/unit/models/test_target_capabilities_tool_usage.pytests/unit/agent/test_agent.pytests/unit/agent/test_attack_parity.pyTotal: 50 new tests; 9858 total unit tests pass.
Quality gate
make pre-commit: all hooks pass (ruff, ty, _async suffix, etc.)make ty: zero new type errors introduced (494 pre-existing)