ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,395 @@
|
||||
"""
|
||||
Custom Agent for Ablation Study
|
||||
Extends ToolCallingAgent to support tone modifications
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import copy
|
||||
import traceback
|
||||
from datetime import datetime, timezone
|
||||
from litellm import completion
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
from tau_bench.agents.base import Agent
|
||||
from tau_bench.agents.tool_calling_agent import message_to_action
|
||||
from tau_bench.envs.base import Env
|
||||
from tau_bench.types import SolveResult, Action, RESPOND_ACTION_NAME
|
||||
|
||||
|
||||
def completion_token_limit(model: str) -> int:
|
||||
"""Return enough output budget for reasoning models to emit an action."""
|
||||
return 8192 if "kimi-k3" in str(model).lower() else 4096
|
||||
|
||||
|
||||
class AblationAgent(Agent):
|
||||
"""
|
||||
Agent that supports tone modifications for ablation studies
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tools_info: List[Dict[str, Any]],
|
||||
wiki: str,
|
||||
model: str,
|
||||
provider: str,
|
||||
temperature: float = 0.0,
|
||||
verbose: bool = True,
|
||||
seed: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the ablation agent
|
||||
|
||||
Args:
|
||||
tools_info: Information about available tools
|
||||
wiki: Wiki/system prompt text (may have tone modifications already applied)
|
||||
model: Model name
|
||||
provider: Model provider
|
||||
temperature: Sampling temperature
|
||||
verbose: Whether to show detailed output (default: True)
|
||||
"""
|
||||
self.tools_info = tools_info
|
||||
self.wiki = wiki
|
||||
self.model = model
|
||||
self.provider = provider
|
||||
self.temperature = temperature
|
||||
self.verbose = verbose
|
||||
self.seed = seed
|
||||
|
||||
def solve(
|
||||
self, env: Env, task_index: Optional[int] = None, max_num_steps: int = 30
|
||||
) -> SolveResult:
|
||||
"""
|
||||
Solve a task with potential tone modifications
|
||||
|
||||
Args:
|
||||
env: The environment
|
||||
task_index: Optional task index
|
||||
max_num_steps: Maximum number of steps
|
||||
|
||||
Returns:
|
||||
SolveResult with the outcome
|
||||
"""
|
||||
if self.verbose:
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🎯 STARTING TASK {task_index if task_index is not None else 'N/A'}")
|
||||
print(f"{'='*80}")
|
||||
print(f"\n📜 SYSTEM PROMPT (Wiki) - {len(self.wiki)} characters:")
|
||||
print("─"*40)
|
||||
# Show first 500 chars of wiki to see tone modifications
|
||||
if len(self.wiki) > 500:
|
||||
print(self.wiki[:500])
|
||||
print(f"... [{len(self.wiki) - 500} more characters]")
|
||||
else:
|
||||
print(self.wiki)
|
||||
print("─"*40)
|
||||
|
||||
total_cost = 0.0
|
||||
env_reset_res = env.reset(task_index=task_index)
|
||||
obs = env_reset_res.observation
|
||||
info = env_reset_res.info.model_dump()
|
||||
reward = 0.0
|
||||
api_records: List[Dict[str, Any]] = []
|
||||
tool_call_count = 0
|
||||
tool_error_count = 0
|
||||
failure = None
|
||||
|
||||
if self.verbose:
|
||||
print(f"\n📝 Initial User Message:")
|
||||
print(f"{'─'*40}")
|
||||
print(obs)
|
||||
print(f"{'─'*40}")
|
||||
|
||||
# Initialize messages
|
||||
messages: List[Dict[str, Any]] = [
|
||||
{"role": "system", "content": self.wiki},
|
||||
{"role": "user", "content": obs},
|
||||
]
|
||||
|
||||
for step in range(max_num_steps):
|
||||
if self.verbose:
|
||||
print(f"\n{'━'*80}")
|
||||
print(f"📍 STEP {step + 1}/{max_num_steps}")
|
||||
print(f"{'━'*80}")
|
||||
|
||||
# Debug: Print request details
|
||||
if self.verbose: # Show full API request details when verbose
|
||||
print(f"\n{'='*60}")
|
||||
print(f"🚀 API CALL #{step + 1} to {self.provider} / {self.model}")
|
||||
print(f"{'='*60}")
|
||||
print(f"📤 SENDING {len(messages)} messages:")
|
||||
print("\n" + "─"*50)
|
||||
for i, msg in enumerate(messages): # Show ALL messages
|
||||
role = msg.get('role', 'unknown')
|
||||
content = msg.get('content', '')
|
||||
print(f"\n📨 Message [{i+1}] - Role: {role.upper()}")
|
||||
print("─"*50)
|
||||
if content:
|
||||
print(content)
|
||||
if 'tool_calls' in msg and msg['tool_calls']:
|
||||
print(f"\n🔧 Tool Calls:")
|
||||
for tc in msg['tool_calls']:
|
||||
if isinstance(tc, dict):
|
||||
print(f" - Function: {tc.get('function', {}).get('name', 'unknown')}")
|
||||
print(f" Args: {tc.get('function', {}).get('arguments', 'none')}")
|
||||
if 'tool_call_id' in msg:
|
||||
print(f"\n🔧 Tool Response ID: {msg['tool_call_id']}")
|
||||
print("─"*50)
|
||||
print("\n" + "="*60)
|
||||
print(f"🔧 Temperature: {self.temperature}")
|
||||
print(f"🛠️ Tools: {len(self.tools_info) if self.tools_info else 0} tools available")
|
||||
if self.tools_info:
|
||||
print("\n📋 COMPLETE TOOL DEFINITIONS (JSON):")
|
||||
print("─"*50)
|
||||
import json
|
||||
for i, tool in enumerate(self.tools_info, 1):
|
||||
print(f"\n[Tool {i}] {tool.get('function', {}).get('name', 'unknown')}:")
|
||||
print(json.dumps(tool, indent=2))
|
||||
print("─"*50)
|
||||
print("="*60)
|
||||
|
||||
# Get completion from model
|
||||
try:
|
||||
# Prepare completion kwargs
|
||||
# Kimi K3 can spend most of a 4K completion budget on hidden
|
||||
# reasoning in the longer Tau-Bench tasks and then return an
|
||||
# empty visible message with no tool call. That is not a
|
||||
# usable Agent action and caused the otherwise complete 60-cell
|
||||
# campaign to fail at the simulator boundary. Reserve the same
|
||||
# reasoning headroom used by the paired Kimi user simulator;
|
||||
# ordinary non-reasoning models retain the historical limit.
|
||||
completion_limit = completion_token_limit(self.model)
|
||||
completion_kwargs = {
|
||||
"messages": messages,
|
||||
"model": self.model,
|
||||
"custom_llm_provider": self.provider,
|
||||
"tools": self.tools_info,
|
||||
"temperature": self.temperature,
|
||||
"max_tokens": completion_limit,
|
||||
}
|
||||
requested_seed = (
|
||||
self.seed + (task_index or 0) * 1000 + step
|
||||
if self.seed is not None else None
|
||||
)
|
||||
if requested_seed is not None:
|
||||
completion_kwargs["seed"] = requested_seed
|
||||
|
||||
# Add reasoning_effort for gpt-5 to minimize thinking tokens
|
||||
if "gpt-5" in self.model:
|
||||
completion_kwargs["extra_body"] = {"reasoning_effort": "low"}
|
||||
if self.verbose:
|
||||
print("💭 Using reasoning_effort='low' to minimize thinking tokens")
|
||||
|
||||
requested_at = datetime.now(timezone.utc).isoformat()
|
||||
started = time.perf_counter()
|
||||
res = completion(**completion_kwargs)
|
||||
choice = res.choices[0]
|
||||
usage = getattr(res, "usage", None)
|
||||
usage_payload = (
|
||||
usage.model_dump()
|
||||
if usage is not None and hasattr(usage, "model_dump")
|
||||
else None
|
||||
)
|
||||
hidden_cost = getattr(res, "_hidden_params", {}).get("response_cost")
|
||||
api_records.append({
|
||||
"requested_at": requested_at,
|
||||
"provider": self.provider,
|
||||
"model": self.model,
|
||||
"task_index": task_index,
|
||||
"step": step + 1,
|
||||
"requested_seed": requested_seed,
|
||||
"request": {
|
||||
"messages": copy.deepcopy(messages),
|
||||
"tools": copy.deepcopy(self.tools_info),
|
||||
"temperature": self.temperature,
|
||||
"max_tokens": completion_limit,
|
||||
},
|
||||
"elapsed_ms": round((time.perf_counter() - started) * 1000, 3),
|
||||
"response": {
|
||||
"id": getattr(res, "id", None),
|
||||
"model": getattr(res, "model", None),
|
||||
"created": getattr(res, "created", None),
|
||||
"finish_reason": getattr(choice, "finish_reason", None),
|
||||
"content": choice.message.content,
|
||||
"reasoning_content": getattr(choice.message, "reasoning_content", None),
|
||||
"tool_calls": [
|
||||
item.model_dump() if hasattr(item, "model_dump") else item
|
||||
for item in (getattr(choice.message, "tool_calls", None) or [])
|
||||
],
|
||||
"usage": usage_payload,
|
||||
"litellm_estimated_cost": hidden_cost,
|
||||
},
|
||||
})
|
||||
|
||||
# Debug: Print response
|
||||
if self.verbose: # Show full API response details when verbose
|
||||
print(f"\n📥 RESPONSE received:")
|
||||
print("─"*50)
|
||||
if res.choices[0].message.content:
|
||||
print("📝 Response Content:")
|
||||
print("─"*50)
|
||||
print(res.choices[0].message.content) # Show FULL content
|
||||
print("─"*50)
|
||||
if hasattr(res.choices[0].message, 'tool_calls') and res.choices[0].message.tool_calls:
|
||||
print(f"\n🔧 Tool calls: {len(res.choices[0].message.tool_calls)} tool(s) called")
|
||||
for idx, tc in enumerate(res.choices[0].message.tool_calls): # Show ALL tool calls
|
||||
print(f"\n Tool Call [{idx+1}]:")
|
||||
print(f" - Function: {tc.function.name}")
|
||||
print(f" - Arguments (FULL):")
|
||||
print(f" {tc.function.arguments}") # Show FULL arguments
|
||||
print(f"{'='*60}\n")
|
||||
except Exception as e:
|
||||
if "requested_at" in locals() and (
|
||||
not api_records or api_records[-1].get("step") != step + 1
|
||||
):
|
||||
api_records.append({
|
||||
"requested_at": requested_at,
|
||||
"provider": self.provider,
|
||||
"model": self.model,
|
||||
"task_index": task_index,
|
||||
"step": step + 1,
|
||||
"requested_seed": requested_seed,
|
||||
"request": completion_kwargs,
|
||||
"elapsed_ms": round((time.perf_counter() - started) * 1000, 3),
|
||||
"error": {"type": type(e).__name__, "message": str(e)},
|
||||
})
|
||||
print(f"\n❌ ERROR calling API:")
|
||||
print(f" Provider: {self.provider}")
|
||||
print(f" Model: {self.model}")
|
||||
print(f" Error: {str(e)}")
|
||||
print(f" Error type: {type(e).__name__}")
|
||||
print(f" Traceback:\n{traceback.format_exc()}")
|
||||
failure = {
|
||||
"type": type(e).__name__,
|
||||
"message": str(e),
|
||||
"traceback": traceback.format_exc(),
|
||||
}
|
||||
# Return a scored failure with every accepted receipt retained.
|
||||
# Raising here made the outer runner discard the complete
|
||||
# in-memory trajectory and all calls made before a late error.
|
||||
reward = 0.0
|
||||
break
|
||||
|
||||
next_message = res.choices[0].message.model_dump()
|
||||
cost = res._hidden_params.get("response_cost", 0)
|
||||
if cost is not None:
|
||||
total_cost += cost
|
||||
|
||||
# Show assistant response if verbose
|
||||
if self.verbose:
|
||||
print(f"\n🤖 Assistant Response:")
|
||||
print(f"{'─'*40}")
|
||||
if next_message.get("content"):
|
||||
print(f"💬 Message: {next_message['content']}")
|
||||
if next_message.get("tool_calls"):
|
||||
print(f"\n🔧 Tool Calls ({len(next_message['tool_calls'])} tool(s)):")
|
||||
for i, tc in enumerate(next_message["tool_calls"], 1):
|
||||
func_name = tc.get('function', {}).get('name', 'unknown')
|
||||
func_args = tc.get('function', {}).get('arguments', '')
|
||||
print(f" [{i}] {func_name}")
|
||||
try:
|
||||
import json
|
||||
args_dict = json.loads(func_args) if isinstance(func_args, str) else func_args
|
||||
for key, value in args_dict.items():
|
||||
value_str = str(value)
|
||||
print(f" • {key}: {value_str}")
|
||||
except Exception:
|
||||
print(f" Args: {func_args}")
|
||||
print(f"{'─'*40}")
|
||||
|
||||
|
||||
# Convert message to action
|
||||
action = message_to_action(next_message)
|
||||
if action.name != RESPOND_ACTION_NAME:
|
||||
tool_call_count += 1
|
||||
|
||||
# Step in environment
|
||||
env_response = env.step(action)
|
||||
if action.name != RESPOND_ACTION_NAME and str(
|
||||
env_response.observation
|
||||
).startswith(("Error:", "Unknown action")):
|
||||
tool_error_count += 1
|
||||
reward = env_response.reward
|
||||
info = {**info, **env_response.info.model_dump()}
|
||||
|
||||
# Show environment response if verbose
|
||||
if self.verbose:
|
||||
print(f"\n🌍 Environment Response:")
|
||||
print(f"{'─'*40}")
|
||||
print(f" Action: {action.name}")
|
||||
if env_response.observation:
|
||||
obs_str = env_response.observation
|
||||
if action.name != RESPOND_ACTION_NAME:
|
||||
print(f" Tool Output: {obs_str}")
|
||||
else:
|
||||
print(f" User Reply: {obs_str}")
|
||||
print(f" Reward: {reward}")
|
||||
print(f" Done: {env_response.done}")
|
||||
print(f"{'─'*40}")
|
||||
|
||||
# Update messages based on action type
|
||||
if action.name != RESPOND_ACTION_NAME:
|
||||
# Tool call - limit to first tool call
|
||||
next_message["tool_calls"] = next_message["tool_calls"][:1]
|
||||
messages.extend(
|
||||
[
|
||||
next_message,
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": next_message["tool_calls"][0]["id"],
|
||||
"name": next_message["tool_calls"][0]["function"]["name"],
|
||||
"content": env_response.observation,
|
||||
},
|
||||
]
|
||||
)
|
||||
else:
|
||||
# Response to user
|
||||
messages.extend(
|
||||
[
|
||||
next_message,
|
||||
{"role": "user", "content": env_response.observation},
|
||||
]
|
||||
)
|
||||
|
||||
# Check if done
|
||||
if env_response.done:
|
||||
if self.verbose:
|
||||
if reward == 1:
|
||||
print(f"\n✅ Task completed successfully! (Reward = {reward})")
|
||||
else:
|
||||
print(f"\n🏁 Task ended (Reward = {reward})")
|
||||
break
|
||||
|
||||
if self.verbose:
|
||||
print(f"\n{'='*80}")
|
||||
print(f"📊 TASK SUMMARY")
|
||||
print(f"{'='*80}")
|
||||
print(f" Final Reward: {reward}")
|
||||
print(f" Total Steps: {step + 1}")
|
||||
print(f" Total Cost: ${total_cost:.4f}")
|
||||
print(f" Messages Exchanged: {len(messages)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
info["experiment_metrics"] = {
|
||||
"agent_steps": step + 1,
|
||||
"agent_model_calls": len(api_records),
|
||||
"tool_calls": tool_call_count,
|
||||
"tool_errors": tool_error_count,
|
||||
}
|
||||
info["agent_api_records"] = api_records
|
||||
info["user_api_records"] = (
|
||||
env.user.get_api_records()
|
||||
if hasattr(env.user, "get_api_records") else []
|
||||
)
|
||||
if failure is not None:
|
||||
info["error"] = failure["message"]
|
||||
info["error_type"] = failure["type"]
|
||||
info["traceback"] = failure["traceback"]
|
||||
|
||||
return SolveResult(
|
||||
reward=reward,
|
||||
info=info,
|
||||
messages=messages,
|
||||
total_cost=total_cost,
|
||||
)
|
||||
Reference in New Issue
Block a user