agentprobe-testing 1.7.0__tar.gz → 1.8.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/PKG-INFO +1 -1
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/__init__.py +8 -1
- agentprobe_testing-1.8.0/agentprobe/agents/openai_chat.py +92 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/cli.py +36 -0
- agentprobe_testing-1.8.0/agentprobe/connect.py +62 -0
- agentprobe_testing-1.8.0/agentprobe/local_runner.py +188 -0
- agentprobe_testing-1.8.0/agentprobe/mcp_server.py +105 -0
- agentprobe_testing-1.8.0/agentprobe/sandbox.py +251 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe_testing.egg-info/PKG-INFO +1 -1
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe_testing.egg-info/SOURCES.txt +10 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/pyproject.toml +1 -1
- agentprobe_testing-1.8.0/tests/test_connect.py +86 -0
- agentprobe_testing-1.8.0/tests/test_local_runner.py +86 -0
- agentprobe_testing-1.8.0/tests/test_mcp_server.py +76 -0
- agentprobe_testing-1.8.0/tests/test_runner_connection.py +53 -0
- agentprobe_testing-1.8.0/tests/test_sandbox.py +91 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/LICENSE +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/README.md +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agents/__init__.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agents/base.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agents/browser_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agents/rule_based.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agents/scripted.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agents/target_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/agreement.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/classifier.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/diff.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/__init__.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/__init__.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/clean.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/complex_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/decoy.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/entities.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/injector_prompt.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/rule_based_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/scenarios.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/split.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/tools.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/access_control/trap.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/__init__.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/chat_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/injector_prompt.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/loader.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/openai_compat_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/oracle.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/tools.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/user_sim.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/tau_retail/world.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/__init__.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/clean.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/entities.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/injector_prompt.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/rule_based_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/scenarios.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/tools.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/domains/web_browsing/trap.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/feedback.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/fixes.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/generic_world.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/injection.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/injector.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/llm.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/playbook.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/priority.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/quickstart.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/reachability.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/registry.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/report.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/runner.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/savings.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenario.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenarios/__init__.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenarios/clean.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenarios/decoy.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenarios/registry.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenarios/split.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/scenarios/trap.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/summarize.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/termui.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/toolcheck.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/toolimport.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/tools.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/trajectory.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/triage.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/validate_scenarios.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe/world.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe_testing.egg-info/dependency_links.txt +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe_testing.egg-info/entry_points.txt +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe_testing.egg-info/requires.txt +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/agentprobe_testing.egg-info/top_level.txt +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/setup.cfg +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_access_control_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_access_control_domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_access_control_rule_based_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_access_control_scenarios.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_access_control_tools.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_agreement.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_browser_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_classifier.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_cli.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_complex_access_control_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_diff.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_failures.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_feedback.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_fixes.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_generic_world.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_injection.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_injection_mode.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_injector.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_llm.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_package_api.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_playbook.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_priority.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_quickstart.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_reachability.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_registry.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_report.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_rule_based_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_runner.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_runner_timeouts.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_savings.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_scenarios.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_target_agent.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_tau_retail.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_termui.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_toolcheck.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_toolimport.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_tools.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_trajectory.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_triage.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_validate_scenarios.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_web_browsing_domain.py +0 -0
- {agentprobe_testing-1.7.0 → agentprobe_testing-1.8.0}/tests/test_world.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: agentprobe-testing
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.8.0
|
|
4
4
|
Summary: Adaptive chaos-testing for LLM agents: a live model-driven Injector that reads an agent's real trajectory and decides where to break something, instead of scripting perturbations in advance.
|
|
5
5
|
License: Business Source License 1.1
|
|
6
6
|
|
|
@@ -57,15 +57,22 @@ from agentprobe.registry import (
|
|
|
57
57
|
from agentprobe.report import InjectionRecord, Report
|
|
58
58
|
from agentprobe.runner import run_recovery, run_robustness_pair
|
|
59
59
|
from agentprobe.savings import FixedFaultInjector, estimate_cost, fixed_fault_for, model_injector_for, retest_failures
|
|
60
|
+
from agentprobe.connect import MODES as CONNECTION_MODES, connect
|
|
61
|
+
from agentprobe.sandbox import DriverAgent, ToolBox, ToolCallError
|
|
60
62
|
from agentprobe.toolcheck import check_agent_tools
|
|
61
63
|
from agentprobe.toolimport import normalize_tools
|
|
62
64
|
from agentprobe.scenario import CommitPattern, FactPattern, GoalSpec, Scenario
|
|
63
65
|
from agentprobe.scenarios.registry import ALL_SCENARIOS as TICKET_SCENARIOS
|
|
64
66
|
from agentprobe.scenarios.registry import BY_ID as TICKET_SCENARIOS_BY_ID
|
|
65
67
|
|
|
66
|
-
__version__ = "1.
|
|
68
|
+
__version__ = "1.8.0"
|
|
67
69
|
|
|
68
70
|
__all__ = [
|
|
71
|
+
"CONNECTION_MODES",
|
|
72
|
+
"DriverAgent",
|
|
73
|
+
"ToolBox",
|
|
74
|
+
"ToolCallError",
|
|
75
|
+
"connect",
|
|
69
76
|
"check_agent_tools",
|
|
70
77
|
"normalize_tools",
|
|
71
78
|
"FixedFaultInjector",
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
"""A tool-calling agent over any OpenAI-compatible chat API (OpenAI, OpenRouter, Moonshot, vLLM, ...).
|
|
2
|
+
|
|
3
|
+
The customer gives a system prompt and a model; the tools come from the test's domain. A plain
|
|
4
|
+
text reply, with no tool call, is the agent's final answer.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import json
|
|
10
|
+
import os
|
|
11
|
+
import time
|
|
12
|
+
import urllib.request
|
|
13
|
+
from typing import Any, Callable, Optional
|
|
14
|
+
|
|
15
|
+
from agentprobe.agents.base import Agent, AgentAction
|
|
16
|
+
|
|
17
|
+
DEFAULT_BASE_URL = "https://api.openai.com/v1"
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class OpenAIChatAgent(Agent):
|
|
21
|
+
def __init__(
|
|
22
|
+
self,
|
|
23
|
+
model: str,
|
|
24
|
+
system_prompt: str,
|
|
25
|
+
api_key: Optional[str] = None,
|
|
26
|
+
base_url: str = DEFAULT_BASE_URL,
|
|
27
|
+
price_per_mtok: tuple[float, float] = (0.0, 0.0),
|
|
28
|
+
post: Optional[Callable[[str, dict, dict], dict]] = None,
|
|
29
|
+
max_tokens: int = 2048,
|
|
30
|
+
temperature: float = 0.0,
|
|
31
|
+
):
|
|
32
|
+
self.model, self.system_prompt = model, system_prompt
|
|
33
|
+
self.url = base_url.rstrip("/") + "/chat/completions"
|
|
34
|
+
self.api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
|
|
35
|
+
self.price = price_per_mtok
|
|
36
|
+
self.max_tokens, self.temperature = max_tokens, temperature
|
|
37
|
+
self._post = post or self._urllib_post
|
|
38
|
+
self.messages: list[dict] = []
|
|
39
|
+
self.tools: list[dict] = []
|
|
40
|
+
self._pending_call_id: Optional[str] = None
|
|
41
|
+
|
|
42
|
+
@staticmethod
|
|
43
|
+
def _urllib_post(url: str, body: dict, headers: dict) -> dict:
|
|
44
|
+
req = urllib.request.Request(url, data=json.dumps(body).encode(), method="POST", headers=headers)
|
|
45
|
+
with urllib.request.urlopen(req, timeout=180) as r: # noqa: S310 -- https/http chosen by the caller
|
|
46
|
+
return json.load(r)
|
|
47
|
+
|
|
48
|
+
def start(self, task: str, tool_schemas: list[dict]) -> None:
|
|
49
|
+
self.tools = [
|
|
50
|
+
{"type": "function", "function": {"name": s["name"], "description": s.get("description", ""), "parameters": s["input_schema"]}}
|
|
51
|
+
for s in tool_schemas
|
|
52
|
+
]
|
|
53
|
+
self.messages = [{"role": "system", "content": self.system_prompt}, {"role": "user", "content": task}]
|
|
54
|
+
self._pending_call_id = None
|
|
55
|
+
|
|
56
|
+
def next_action(self) -> AgentAction:
|
|
57
|
+
t0 = time.monotonic()
|
|
58
|
+
body = {"model": self.model, "messages": self.messages, "max_tokens": self.max_tokens, "temperature": self.temperature}
|
|
59
|
+
if self.tools:
|
|
60
|
+
body["tools"] = self.tools
|
|
61
|
+
body["parallel_tool_calls"] = False
|
|
62
|
+
headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
|
|
63
|
+
data = self._post(self.url, body, headers)
|
|
64
|
+
msg = data["choices"][0]["message"]
|
|
65
|
+
usage = data.get("usage") or {}
|
|
66
|
+
cost = usage.get("cost")
|
|
67
|
+
if cost is None:
|
|
68
|
+
cost = (usage.get("prompt_tokens", 0) * self.price[0] + usage.get("completion_tokens", 0) * self.price[1]) / 1e6
|
|
69
|
+
latency = time.monotonic() - t0
|
|
70
|
+
calls = msg.get("tool_calls") or []
|
|
71
|
+
if calls:
|
|
72
|
+
call = calls[0]
|
|
73
|
+
self.messages.append({"role": "assistant", "content": msg.get("content"), "tool_calls": [call]})
|
|
74
|
+
self._pending_call_id = call["id"]
|
|
75
|
+
try:
|
|
76
|
+
args = json.loads(call["function"].get("arguments") or "{}")
|
|
77
|
+
except json.JSONDecodeError:
|
|
78
|
+
args = {}
|
|
79
|
+
return AgentAction(kind="tool_call", tool_name=call["function"]["name"], tool_args=args, latency_s=latency, cost_usd=cost)
|
|
80
|
+
text = msg.get("content") or ""
|
|
81
|
+
self.messages.append({"role": "assistant", "content": text})
|
|
82
|
+
return AgentAction(kind="final_answer", text=text, latency_s=latency, cost_usd=cost)
|
|
83
|
+
|
|
84
|
+
def observe(self, tool_name: str, tool_args: dict, result: Any, ok: bool) -> None:
|
|
85
|
+
text = result if isinstance(result, str) else json.dumps(result, default=str)
|
|
86
|
+
if not ok:
|
|
87
|
+
text = "Error: " + (result.get("error", text) if isinstance(result, dict) else text)
|
|
88
|
+
self.messages.append({"role": "tool", "tool_call_id": self._pending_call_id, "content": text})
|
|
89
|
+
self._pending_call_id = None
|
|
90
|
+
|
|
91
|
+
def receive_message(self, text: str) -> None:
|
|
92
|
+
self.messages.append({"role": "user", "content": text})
|
|
@@ -4,6 +4,7 @@ from __future__ import annotations
|
|
|
4
4
|
|
|
5
5
|
import argparse
|
|
6
6
|
import json
|
|
7
|
+
import os
|
|
7
8
|
import sys
|
|
8
9
|
|
|
9
10
|
from dotenv import load_dotenv
|
|
@@ -343,9 +344,44 @@ def main(argv: list[str] | None = None) -> int:
|
|
|
343
344
|
help="Exit 1 if any scenario that passed in `before` fails in `after` -- for a CI gate.",
|
|
344
345
|
)
|
|
345
346
|
|
|
347
|
+
def _local_opts(p, scenario_required: bool) -> None:
|
|
348
|
+
p.add_argument("--domain-key", help="A domain generated on sampact.com. Omit to use a built-in demo domain (--domain).")
|
|
349
|
+
p.add_argument("--domain", choices=["ticket", "access_control"], default="ticket", help="The built-in demo domain, when no --domain-key is given.")
|
|
350
|
+
p.add_argument("--api-key", default=os.environ.get("AGENTPROBE_API_KEY"), help="Your sampact.com API key (or set AGENTPROBE_API_KEY): fetches a private domain and uploads the report.")
|
|
351
|
+
p.add_argument("--base-url", default=None, help="Another AgentProbe server (default: https://sampact.com).")
|
|
352
|
+
p.add_argument("--fault", default="false_success", help="none, model (adaptive injector, needs ANTHROPIC_API_KEY), or one of: false_success, tool_error, phantom_success, unverified_outcome.")
|
|
353
|
+
p.add_argument("--no-upload", action="store_true", help="Don't upload the report to your account.")
|
|
354
|
+
p.add_argument("--out", default=None, help="Write the HTML report here.")
|
|
355
|
+
if scenario_required:
|
|
356
|
+
p.add_argument("--scenario", required=True, help="The scenario id to run.")
|
|
357
|
+
|
|
358
|
+
mcp_p = sub.add_parser("mcp", help="Test an MCP-capable agent: this process is its tool server (stdio).")
|
|
359
|
+
_local_opts(mcp_p, True)
|
|
360
|
+
|
|
361
|
+
rd_p = sub.add_parser("run-domain", help="Run a domain's scenarios from your machine against an agent on your network, then upload the report.")
|
|
362
|
+
_local_opts(rd_p, False)
|
|
363
|
+
rd_p.add_argument("--scenario", default=None, help="Comma-separated scenario ids (default: all).")
|
|
364
|
+
rd_p.add_argument("--connection", choices=["http", "model", "openai"], default="http", help="How your agent connects.")
|
|
365
|
+
rd_p.add_argument("--agent-url", help="--connection http: your agent's decide endpoint.")
|
|
366
|
+
rd_p.add_argument("--header", action="append", help="--connection http: an extra header, 'Name: value' (repeatable).")
|
|
367
|
+
rd_p.add_argument("--system-prompt", help="--connection model/openai: your system prompt, or a path to a file holding it.")
|
|
368
|
+
rd_p.add_argument("--model", help="--connection model/openai: the model name.")
|
|
369
|
+
rd_p.add_argument("--agent-api-key", help="--connection openai: the key for your OpenAI-compatible provider.")
|
|
370
|
+
rd_p.add_argument("--skip-if-clean-fails", action="store_true", help="Don't inject into scenarios your agent already fails without faults.")
|
|
371
|
+
rd_p.add_argument("--budget-usd", type=float, default=None, help="Stop starting scenarios once spend reaches this.")
|
|
372
|
+
|
|
346
373
|
args = parser.parse_args(argv)
|
|
347
374
|
load_dotenv()
|
|
348
375
|
|
|
376
|
+
if args.command in ("mcp", "run-domain"):
|
|
377
|
+
from agentprobe import local_runner
|
|
378
|
+
|
|
379
|
+
try:
|
|
380
|
+
return local_runner.cmd_mcp(args) if args.command == "mcp" else local_runner.cmd_run_domain(args)
|
|
381
|
+
except (ValueError, KeyError) as e:
|
|
382
|
+
print(f"agentprobe: {e}", file=sys.stderr)
|
|
383
|
+
return 1
|
|
384
|
+
|
|
349
385
|
if args.command == "diff":
|
|
350
386
|
from agentprobe.diff import diff_runs, render_diff
|
|
351
387
|
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
"""One switch for every way of connecting an agent.
|
|
2
|
+
|
|
3
|
+
agent = ap.connect("http", url="https://my-agent.example.com/decide")
|
|
4
|
+
agent = ap.connect("function", fn=my_decide_fn)
|
|
5
|
+
agent = ap.connect("model", system_prompt="You are ...", model="claude-sonnet-5")
|
|
6
|
+
agent = ap.connect("openai", system_prompt="You are ...", model="gpt-4o", api_key="sk-...")
|
|
7
|
+
agent = ap.connect("tools", driver=my_driver) # an agent that runs its own loop
|
|
8
|
+
|
|
9
|
+
report = ap.quick_test_all(scenarios, agent, ...) # the same call whichever you picked
|
|
10
|
+
|
|
11
|
+
MODES lists what's available and what each needs. (An MCP server, `agentprobe mcp`, is the other
|
|
12
|
+
way in, for agents that are MCP clients; it is started from the command line, not from here.)
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
from typing import Any, Callable
|
|
18
|
+
|
|
19
|
+
from agentprobe.agents.base import Agent
|
|
20
|
+
|
|
21
|
+
MODES: dict[str, dict] = {
|
|
22
|
+
"http": {"needs": ["url"], "optional": ["headers", "timeout"], "what": "Your agent is a web endpoint that returns its next step."},
|
|
23
|
+
"function": {"needs": ["fn"], "optional": [], "what": "Your agent is a Python function (task, history) -> AgentAction."},
|
|
24
|
+
"model": {"needs": ["system_prompt"], "optional": ["model", "max_tokens"], "what": "AgentProbe runs a Claude tool loop with your system prompt (the tools come from the test)."},
|
|
25
|
+
"openai": {"needs": ["system_prompt", "model"], "optional": ["api_key", "base_url", "max_tokens"], "what": "The same, over any OpenAI-compatible API (OpenAI, OpenRouter, Moonshot, vLLM...)."},
|
|
26
|
+
"tools": {"needs": ["driver"], "optional": ["timeout"], "what": "Your agent runs its own loop, calling tools you take from the test (LangChain, OpenAI Agents SDK, anything)."},
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def connect(mode: str, **options: Any) -> Callable[[], Agent]:
|
|
31
|
+
"""A target factory for quick_test/quick_test_all, for the connection `mode` you choose."""
|
|
32
|
+
spec = MODES.get(mode)
|
|
33
|
+
if spec is None:
|
|
34
|
+
raise ValueError(f"unknown connection mode {mode!r}; choose one of {', '.join(MODES)}")
|
|
35
|
+
missing = [k for k in spec["needs"] if k not in options or options[k] in (None, "")]
|
|
36
|
+
if missing:
|
|
37
|
+
raise ValueError(f"connection mode {mode!r} needs: {', '.join(missing)}")
|
|
38
|
+
unknown = [k for k in options if k not in spec["needs"] and k not in spec["optional"]]
|
|
39
|
+
if unknown:
|
|
40
|
+
raise ValueError(f"connection mode {mode!r} does not take: {', '.join(unknown)}")
|
|
41
|
+
|
|
42
|
+
if mode == "http":
|
|
43
|
+
from agentprobe.quickstart import http_agent
|
|
44
|
+
|
|
45
|
+
return http_agent(options["url"], headers=options.get("headers"), **({"timeout": options["timeout"]} if "timeout" in options else {}))
|
|
46
|
+
if mode == "function":
|
|
47
|
+
from agentprobe.quickstart import wrap_agent
|
|
48
|
+
|
|
49
|
+
return wrap_agent(options["fn"])
|
|
50
|
+
if mode == "model":
|
|
51
|
+
from agentprobe.agents.target_agent import TargetAgent
|
|
52
|
+
|
|
53
|
+
kwargs = {k: options[k] for k in ("model", "max_tokens") if k in options}
|
|
54
|
+
return lambda: TargetAgent(system_prompt=options["system_prompt"], **kwargs)
|
|
55
|
+
if mode == "openai":
|
|
56
|
+
from agentprobe.agents.openai_chat import OpenAIChatAgent
|
|
57
|
+
|
|
58
|
+
kwargs = {k: options[k] for k in ("api_key", "base_url", "max_tokens") if k in options}
|
|
59
|
+
return lambda: OpenAIChatAgent(model=options["model"], system_prompt=options["system_prompt"], **kwargs)
|
|
60
|
+
from agentprobe.sandbox import DriverAgent
|
|
61
|
+
|
|
62
|
+
return DriverAgent.factory(options["driver"], **({"timeout": options["timeout"]} if "timeout" in options else {}))
|
|
@@ -0,0 +1,188 @@
|
|
|
1
|
+
"""Run a test on YOUR machine, inside YOUR network, and send only the report to your account.
|
|
2
|
+
|
|
3
|
+
This is the "customer-hosted runner": an agent that can't be reached from the internet (a staging
|
|
4
|
+
service on a private network, a laptop, a browser agent) is tested locally, against a domain you
|
|
5
|
+
generated on sampact.com, and the result shows up under Your Runs like any other run. Nothing
|
|
6
|
+
about your agent has to be exposed.
|
|
7
|
+
|
|
8
|
+
agentprobe run-domain --domain-key K --api-key ap_live_... --connection http --agent-url http://localhost:8080/decide
|
|
9
|
+
agentprobe mcp --domain-key K --api-key ap_live_... --scenario refund-1 # an MCP client is the agent
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from __future__ import annotations
|
|
13
|
+
|
|
14
|
+
import os
|
|
15
|
+
import sys
|
|
16
|
+
from dataclasses import dataclass
|
|
17
|
+
from typing import Any, Callable, Optional
|
|
18
|
+
|
|
19
|
+
from agentprobe.connect import connect
|
|
20
|
+
|
|
21
|
+
FAULTS = {
|
|
22
|
+
"false_success": "FALSE_SUCCESS",
|
|
23
|
+
"tool_error": "TOOL_ERROR",
|
|
24
|
+
"phantom_success": "PHANTOM_SUCCESS",
|
|
25
|
+
"unverified_outcome": "UNVERIFIED_OUTCOME",
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@dataclass
|
|
30
|
+
class ResolvedDomain:
|
|
31
|
+
domain: Any
|
|
32
|
+
scenarios: dict
|
|
33
|
+
key: Optional[str]
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def resolve_domain(domain_key: Optional[str], builtin: str = "ticket", api_key: Optional[str] = None, base_url: Optional[str] = None) -> ResolvedDomain:
|
|
37
|
+
"""A generated domain from sampact.com (by key), or one of the built-in demo domains."""
|
|
38
|
+
if domain_key:
|
|
39
|
+
from agentprobe.registry import DEFAULT_BASE_URL, fetch_domain
|
|
40
|
+
|
|
41
|
+
fetched = fetch_domain(domain_key, base_url=base_url or DEFAULT_BASE_URL, api_key=api_key)
|
|
42
|
+
return ResolvedDomain(fetched.domain, fetched.scenarios, domain_key)
|
|
43
|
+
if builtin == "access_control":
|
|
44
|
+
from agentprobe.domains.access_control.domain import ACCESS_CONTROL_DOMAIN
|
|
45
|
+
from agentprobe.domains.access_control.scenarios import SCENARIOS
|
|
46
|
+
|
|
47
|
+
return ResolvedDomain(ACCESS_CONTROL_DOMAIN, dict(SCENARIOS), None)
|
|
48
|
+
from agentprobe.domain import TICKET_DOMAIN
|
|
49
|
+
from agentprobe.scenarios.registry import BY_ID
|
|
50
|
+
|
|
51
|
+
return ResolvedDomain(TICKET_DOMAIN, dict(BY_ID), None)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def injector_factory_for(fault: str, domain, model: Optional[str] = None) -> Callable:
|
|
55
|
+
"""`fault` is none, model (the adaptive injector), or one of FAULTS (a fixed fault on the
|
|
56
|
+
scenario's first required commit, no model call)."""
|
|
57
|
+
from agentprobe.injection import InjectionKind
|
|
58
|
+
from agentprobe.injector import NullInjector
|
|
59
|
+
from agentprobe.savings import fixed_fault_for, model_injector_for
|
|
60
|
+
|
|
61
|
+
if fault == "none":
|
|
62
|
+
return lambda sc: NullInjector()
|
|
63
|
+
if fault == "model":
|
|
64
|
+
return lambda sc: model_injector_for(domain, **({"model": model} if model else {}))
|
|
65
|
+
if fault in FAULTS:
|
|
66
|
+
kind = InjectionKind[FAULTS[fault]]
|
|
67
|
+
return lambda sc: fixed_fault_for(sc, kind)
|
|
68
|
+
raise ValueError(f"unknown fault {fault!r}; choose none, model, or one of {', '.join(FAULTS)}")
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def _read(path_or_text: Optional[str]) -> Optional[str]:
|
|
72
|
+
if path_or_text and os.path.isfile(path_or_text):
|
|
73
|
+
return open(path_or_text).read()
|
|
74
|
+
return path_or_text
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def build_connection(connection: str, opts: dict) -> Callable:
|
|
78
|
+
"""Turn command-line style options into a connection (see agentprobe.connect)."""
|
|
79
|
+
if connection == "http":
|
|
80
|
+
headers = {}
|
|
81
|
+
for h in opts.get("header") or []:
|
|
82
|
+
k, _, v = h.partition(":")
|
|
83
|
+
headers[k.strip()] = v.strip()
|
|
84
|
+
return connect("http", url=opts["agent_url"], **({"headers": headers} if headers else {}))
|
|
85
|
+
if connection == "model":
|
|
86
|
+
return connect("model", system_prompt=_read(opts["system_prompt"]), **({"model": opts["model"]} if opts.get("model") else {}))
|
|
87
|
+
if connection == "openai":
|
|
88
|
+
extra = {k: opts[k] for k in ("base_url",) if opts.get(k)}
|
|
89
|
+
if opts.get("agent_api_key"):
|
|
90
|
+
extra["api_key"] = opts["agent_api_key"]
|
|
91
|
+
return connect("openai", system_prompt=_read(opts["system_prompt"]), model=opts["model"], **extra)
|
|
92
|
+
raise ValueError(f"--connection must be http, model or openai here (got {connection!r}); use `agentprobe mcp` for an MCP client")
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def run_domain(
|
|
96
|
+
resolved: ResolvedDomain,
|
|
97
|
+
agent_factory: Callable,
|
|
98
|
+
scenario_ids: Optional[list[str]] = None,
|
|
99
|
+
fault: str = "false_success",
|
|
100
|
+
api_key: Optional[str] = None,
|
|
101
|
+
base_url: Optional[str] = None,
|
|
102
|
+
upload: bool = True,
|
|
103
|
+
out_path: Optional[str] = None,
|
|
104
|
+
**quick_test_all_kwargs: Any,
|
|
105
|
+
):
|
|
106
|
+
from agentprobe.quickstart import quick_test_all
|
|
107
|
+
|
|
108
|
+
chosen = {sid: resolved.scenarios[sid] for sid in (scenario_ids or resolved.scenarios)}
|
|
109
|
+
kwargs = dict(quick_test_all_kwargs)
|
|
110
|
+
if upload and api_key:
|
|
111
|
+
kwargs["upload_api_key"] = api_key
|
|
112
|
+
kwargs["domain_key"] = resolved.key
|
|
113
|
+
report = quick_test_all(
|
|
114
|
+
chosen,
|
|
115
|
+
agent_factory,
|
|
116
|
+
injector_factory=injector_factory_for(fault, resolved.domain),
|
|
117
|
+
domain=resolved.domain,
|
|
118
|
+
injector_model=None if fault == "model" else (f"fixed fault: {fault}" if fault != "none" else None),
|
|
119
|
+
record_injections=bool(upload and api_key),
|
|
120
|
+
**kwargs,
|
|
121
|
+
)
|
|
122
|
+
if out_path:
|
|
123
|
+
with open(out_path, "w") as f:
|
|
124
|
+
f.write(report.render_html())
|
|
125
|
+
return report
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def mcp_report(trajectory, scenario, domain, fault: str, target_label: str = "an MCP client"):
|
|
129
|
+
"""A Report for the one run an MCP client just did (chaos only: there is no clean baseline)."""
|
|
130
|
+
from agentprobe.injection import injection_was_triggered
|
|
131
|
+
from agentprobe.report import InjectionRecord, Report
|
|
132
|
+
from agentprobe.toolcheck import required_tools
|
|
133
|
+
|
|
134
|
+
records = [InjectionRecord(scenario_id=trajectory.scenario_id, applied=a, triggered=injection_was_triggered(a, trajectory.steps)) for a in trajectory.injections]
|
|
135
|
+
return Report(
|
|
136
|
+
mode="recovery",
|
|
137
|
+
injector_model="none" if fault == "none" else f"fixed fault: {fault}" if fault in FAULTS else "model",
|
|
138
|
+
target_model=target_label,
|
|
139
|
+
clean_trajectories=[],
|
|
140
|
+
chaos_trajectories=[trajectory],
|
|
141
|
+
injection_records=records,
|
|
142
|
+
injection_mode="pinned" if fault in FAULTS else ("exploration" if fault == "model" else None),
|
|
143
|
+
tool_expectations={"known": sorted(domain.all_tools), "required": {scenario.id: sorted(required_tools(scenario))}},
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def cmd_mcp(args) -> int:
|
|
148
|
+
from agentprobe.mcp_server import serve
|
|
149
|
+
from agentprobe.registry import upload_run
|
|
150
|
+
|
|
151
|
+
resolved = resolve_domain(args.domain_key, args.domain, api_key=args.api_key, base_url=args.base_url)
|
|
152
|
+
if args.scenario not in resolved.scenarios:
|
|
153
|
+
print(f"agentprobe: no scenario {args.scenario!r}; choose one of {', '.join(sorted(resolved.scenarios))}", file=sys.stderr)
|
|
154
|
+
return 1
|
|
155
|
+
scenario = resolved.scenarios[args.scenario]
|
|
156
|
+
injector = injector_factory_for(args.fault, resolved.domain)(scenario)
|
|
157
|
+
|
|
158
|
+
def on_finish(traj):
|
|
159
|
+
report = mcp_report(traj, scenario, resolved.domain, args.fault)
|
|
160
|
+
out = args.out or f"agentprobe_mcp_{scenario.id}.html"
|
|
161
|
+
with open(out, "w") as f:
|
|
162
|
+
f.write(report.render_html())
|
|
163
|
+
print(report.render(), file=sys.stderr)
|
|
164
|
+
print(f"agentprobe: report written to {out}", file=sys.stderr)
|
|
165
|
+
if args.api_key and not args.no_upload:
|
|
166
|
+
upload_run(report, api_key=args.api_key, domain_key=resolved.key, **({"base_url": args.base_url} if args.base_url else {}))
|
|
167
|
+
print("agentprobe: uploaded to your account (Your Runs)", file=sys.stderr)
|
|
168
|
+
|
|
169
|
+
serve(scenario, resolved.domain, injector, on_finish=on_finish)
|
|
170
|
+
return 0
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
def cmd_run_domain(args) -> int:
|
|
174
|
+
resolved = resolve_domain(args.domain_key, args.domain, api_key=args.api_key, base_url=args.base_url)
|
|
175
|
+
try:
|
|
176
|
+
agent = build_connection(args.connection, vars(args))
|
|
177
|
+
except (ValueError, KeyError) as e:
|
|
178
|
+
print(f"agentprobe: {e}", file=sys.stderr)
|
|
179
|
+
return 1
|
|
180
|
+
ids = [s.strip() for s in args.scenario.split(",")] if args.scenario else None
|
|
181
|
+
report = run_domain(
|
|
182
|
+
resolved, agent, ids, fault=args.fault, api_key=args.api_key, base_url=args.base_url,
|
|
183
|
+
upload=not args.no_upload, out_path=args.out,
|
|
184
|
+
**({"skip_injection_if_clean_fails": True} if args.skip_if_clean_fails else {}),
|
|
185
|
+
**({"budget_usd": args.budget_usd} if args.budget_usd is not None else {}),
|
|
186
|
+
)
|
|
187
|
+
print(report.render())
|
|
188
|
+
return 0
|
|
@@ -0,0 +1,105 @@
|
|
|
1
|
+
"""Test any MCP-capable agent (Claude Desktop, Cursor, an MCP client you wrote) by being its tools.
|
|
2
|
+
|
|
3
|
+
Run `agentprobe mcp ...`, point the client at it as an MCP server (stdio), and give it the task.
|
|
4
|
+
Every tool the client calls goes into the test world, where failures are injected; the answer goes
|
|
5
|
+
back. When the agent is done it calls the `end_conversation` tool (or the client disconnects), and
|
|
6
|
+
the report is written.
|
|
7
|
+
|
|
8
|
+
JSON-RPC over stdin/stdout, one message per line, the subset of MCP a tool server needs:
|
|
9
|
+
initialize, tools/list, tools/call (plus ping and the initialized notification).
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from __future__ import annotations
|
|
13
|
+
|
|
14
|
+
import json
|
|
15
|
+
import sys
|
|
16
|
+
from typing import IO, Callable, Optional
|
|
17
|
+
|
|
18
|
+
from agentprobe.sandbox import ExternalSession
|
|
19
|
+
|
|
20
|
+
END_TOOL = "end_conversation"
|
|
21
|
+
PROTOCOL = "2024-11-05"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def _end_tool_schema() -> dict:
|
|
25
|
+
return {
|
|
26
|
+
"name": END_TOOL,
|
|
27
|
+
"description": "Call this when you are finished with the task. Put your final message to the customer in `message`.",
|
|
28
|
+
"inputSchema": {"type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"]},
|
|
29
|
+
}
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def serve(
|
|
33
|
+
scenario,
|
|
34
|
+
domain,
|
|
35
|
+
injector,
|
|
36
|
+
stdin: IO[str] = None,
|
|
37
|
+
stdout: IO[str] = None,
|
|
38
|
+
on_finish: Optional[Callable] = None,
|
|
39
|
+
) -> object:
|
|
40
|
+
"""Run one test over MCP. Returns the finished Trajectory. `on_finish(trajectory)` is called
|
|
41
|
+
once, before the final reply, so a caller can write the report."""
|
|
42
|
+
stdin = stdin or sys.stdin
|
|
43
|
+
stdout = stdout or sys.stdout
|
|
44
|
+
session = ExternalSession(scenario, domain, injector, variant="recovery")
|
|
45
|
+
session.start()
|
|
46
|
+
|
|
47
|
+
def send(msg: dict) -> None:
|
|
48
|
+
stdout.write(json.dumps(msg) + "\n")
|
|
49
|
+
stdout.flush()
|
|
50
|
+
|
|
51
|
+
def result(i, payload: dict) -> None:
|
|
52
|
+
send({"jsonrpc": "2.0", "id": i, "result": payload})
|
|
53
|
+
|
|
54
|
+
instructions = (
|
|
55
|
+
f"You are working a task. Use ONLY the tools of this server. When you are done, call {END_TOOL} with your final message.\n\n"
|
|
56
|
+
f"TASK: {scenario.task}"
|
|
57
|
+
)
|
|
58
|
+
done = False
|
|
59
|
+
|
|
60
|
+
def finish(text: str) -> None:
|
|
61
|
+
nonlocal done
|
|
62
|
+
if done:
|
|
63
|
+
return
|
|
64
|
+
done = True
|
|
65
|
+
traj = session.finish(text)
|
|
66
|
+
if on_finish:
|
|
67
|
+
on_finish(traj)
|
|
68
|
+
|
|
69
|
+
for line in stdin:
|
|
70
|
+
line = line.strip()
|
|
71
|
+
if not line:
|
|
72
|
+
continue
|
|
73
|
+
try:
|
|
74
|
+
req = json.loads(line)
|
|
75
|
+
except json.JSONDecodeError:
|
|
76
|
+
continue
|
|
77
|
+
method, i = req.get("method"), req.get("id")
|
|
78
|
+
if method == "initialize":
|
|
79
|
+
result(i, {
|
|
80
|
+
"protocolVersion": (req.get("params") or {}).get("protocolVersion", PROTOCOL),
|
|
81
|
+
"capabilities": {"tools": {}},
|
|
82
|
+
"serverInfo": {"name": "agentprobe", "version": "1"},
|
|
83
|
+
"instructions": instructions,
|
|
84
|
+
})
|
|
85
|
+
elif method == "tools/list":
|
|
86
|
+
tools = [{"name": s["name"], "description": s.get("description", ""), "inputSchema": s["input_schema"]} for s in session.tool_schemas]
|
|
87
|
+
result(i, {"tools": tools + [_end_tool_schema()]})
|
|
88
|
+
elif method == "tools/call":
|
|
89
|
+
params = req.get("params") or {}
|
|
90
|
+
name, args = params.get("name"), params.get("arguments") or {}
|
|
91
|
+
if name == END_TOOL:
|
|
92
|
+
finish(str(args.get("message", "")))
|
|
93
|
+
result(i, {"content": [{"type": "text", "text": "The test is over. Thank you."}], "isError": False})
|
|
94
|
+
break
|
|
95
|
+
res, ok = session.call(name, args)
|
|
96
|
+
text = res if isinstance(res, str) else json.dumps(res, default=str)
|
|
97
|
+
if not ok and isinstance(res, dict):
|
|
98
|
+
text = "Error: " + str(res.get("error", text))
|
|
99
|
+
result(i, {"content": [{"type": "text", "text": text}], "isError": not ok})
|
|
100
|
+
elif method == "ping":
|
|
101
|
+
result(i, {})
|
|
102
|
+
elif i is not None and method:
|
|
103
|
+
send({"jsonrpc": "2.0", "id": i, "error": {"code": -32601, "message": f"unknown method {method}"}})
|
|
104
|
+
finish("") # the client disconnected without calling end_conversation
|
|
105
|
+
return session.trajectory
|