|
5 | 5 | import pytest |
6 | 6 | from pydantic import BaseModel, Field |
7 | 7 |
|
8 | | -from copilot import define_tool |
| 8 | +from copilot import ToolSet, define_tool |
9 | 9 | from copilot.rpc import ( |
10 | 10 | PermissionDecisionApproveOnce, |
11 | 11 | PermissionDecisionReject, |
@@ -48,6 +48,49 @@ def encrypt_string(params: EncryptParams, invocation: ToolInvocation) -> str: |
48 | 48 | assistant_message = await get_final_assistant_message(session) |
49 | 49 | assert "HELLO" in assistant_message.data.content |
50 | 50 |
|
| 51 | + async def test_low_level_tool_definition(self, ctx: E2ETestContext): |
| 52 | + class PhaseArgs(BaseModel): |
| 53 | + phase: str = Field( |
| 54 | + description="Current phase", |
| 55 | + pattern="^(searching|analyzing|done)$", |
| 56 | + ) |
| 57 | + |
| 58 | + class SearchArgs(BaseModel): |
| 59 | + keyword: str |
| 60 | + |
| 61 | + current_phase = "" |
| 62 | + |
| 63 | + @define_tool("set_current_phase", description="Sets the current phase of the agent") |
| 64 | + def set_current_phase(params: PhaseArgs, invocation: ToolInvocation) -> str: |
| 65 | + nonlocal current_phase |
| 66 | + current_phase = params.phase |
| 67 | + return f"Phase set to {params.phase}" |
| 68 | + |
| 69 | + @define_tool("search_items", description="Search for items by keyword") |
| 70 | + def search_items(params: SearchArgs, invocation: ToolInvocation) -> str: |
| 71 | + args = invocation.arguments or {} |
| 72 | + keyword = str(args.get("keyword", "")) |
| 73 | + assert keyword == "copilot" |
| 74 | + return "Found: item_alpha, item_beta" |
| 75 | + |
| 76 | + session = await ctx.client.create_session( |
| 77 | + on_permission_request=PermissionHandler.approve_all, |
| 78 | + available_tools=ToolSet().add_custom("*").add_builtin("web_fetch"), |
| 79 | + tools=[set_current_phase, search_items], |
| 80 | + ) |
| 81 | + |
| 82 | + prompt = ( |
| 83 | + "First, set the current phase to 'analyzing'. Then search for items with " |
| 84 | + "keyword 'copilot'. Report the phase and search results." |
| 85 | + ) |
| 86 | + await session.send(prompt) |
| 87 | + assistant_message = await get_final_assistant_message(session) |
| 88 | + content = assistant_message.data.content or "" |
| 89 | + assert content != "" |
| 90 | + assert "analyzing" in content.lower() |
| 91 | + assert "item_alpha" in content.lower() or "item_beta" in content.lower() |
| 92 | + assert current_phase == "analyzing" |
| 93 | + |
51 | 94 | async def test_handles_tool_calling_errors(self, ctx: E2ETestContext): |
52 | 95 | @define_tool("get_user_location", description="Gets the user's location") |
53 | 96 | def get_user_location() -> str: |
|
0 commit comments