wehub-skill-sync
418 lines
14 KiB
Python
418 lines
14 KiB
Python
"""
|
|
TraceCapture: Wraps M2.1 API to capture interleaved thinking traces.
|
|
|
|
This module provides the core functionality for executing agent tasks
|
|
through MiniMax M2.1 while capturing all reasoning traces for analysis.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any, Callable
|
|
|
|
import anthropic
|
|
|
|
from reasoning_trace_optimizer.models import (
|
|
ReasoningTrace,
|
|
ThinkingBlock,
|
|
ToolCall,
|
|
)
|
|
|
|
|
|
class TraceCapture:
|
|
"""
|
|
Captures reasoning traces from MiniMax M2.1's interleaved thinking.
|
|
|
|
This class wraps the Anthropic SDK configured for M2.1 and captures
|
|
all thinking blocks, tool calls, and responses during agent execution.
|
|
|
|
Example:
|
|
```python
|
|
capture = TraceCapture()
|
|
trace = capture.run(
|
|
task="What's the weather in San Francisco?",
|
|
tools=[weather_tool],
|
|
tool_executor=execute_tool
|
|
)
|
|
print(f"Captured {len(trace.thinking_blocks)} thinking blocks")
|
|
```
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str | None = None,
|
|
base_url: str = "https://api.minimax.io/anthropic",
|
|
model: str = "MiniMax-M2.1",
|
|
):
|
|
"""
|
|
Initialize TraceCapture with M2.1 configuration.
|
|
|
|
Args:
|
|
api_key: MiniMax API key (defaults to ANTHROPIC_API_KEY env var)
|
|
base_url: API base URL (international or China endpoint)
|
|
model: Model to use (MiniMax-M2.1, MiniMax-M2.1-lightning, MiniMax-M2)
|
|
"""
|
|
self.model = model
|
|
self.client = anthropic.Anthropic(
|
|
api_key=api_key or os.environ.get("ANTHROPIC_API_KEY"),
|
|
base_url=base_url,
|
|
)
|
|
|
|
def run(
|
|
self,
|
|
task: str,
|
|
system_prompt: str = "You are a helpful assistant.",
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_executor: Callable[[str, dict], str] | None = None,
|
|
max_turns: int = 10,
|
|
max_tokens: int = 4096,
|
|
) -> ReasoningTrace:
|
|
"""
|
|
Execute a task and capture the full reasoning trace.
|
|
|
|
Args:
|
|
task: The user task/query to execute
|
|
system_prompt: System prompt for the agent
|
|
tools: List of tool definitions in Anthropic format
|
|
tool_executor: Function to execute tool calls (name, input) -> result
|
|
max_turns: Maximum conversation turns before stopping
|
|
max_tokens: Maximum tokens per response
|
|
|
|
Returns:
|
|
ReasoningTrace containing all thinking blocks, tool calls, and responses
|
|
"""
|
|
trace = ReasoningTrace(
|
|
session_id=str(uuid.uuid4()),
|
|
task=task,
|
|
system_prompt=system_prompt,
|
|
model=self.model,
|
|
started_at=datetime.now(),
|
|
)
|
|
|
|
messages = [{"role": "user", "content": task}]
|
|
turn = 0
|
|
|
|
try:
|
|
while turn < max_turns:
|
|
# Build request parameters
|
|
params = {
|
|
"model": self.model,
|
|
"max_tokens": max_tokens,
|
|
"system": system_prompt,
|
|
"messages": messages,
|
|
}
|
|
if tools:
|
|
params["tools"] = tools
|
|
|
|
# Make API call
|
|
response = self.client.messages.create(**params)
|
|
|
|
# Process response content blocks
|
|
thinking_blocks, text_blocks, tool_use_blocks = self._process_response(
|
|
response, turn, trace
|
|
)
|
|
|
|
# If no tool calls, we're done
|
|
if not tool_use_blocks:
|
|
trace.final_response = (
|
|
text_blocks[0].text if text_blocks else None
|
|
)
|
|
trace.success = True
|
|
break
|
|
|
|
# Append assistant response to history (CRITICAL for M2.1)
|
|
messages.append({"role": "assistant", "content": response.content})
|
|
|
|
# Execute tools and collect results
|
|
tool_results = []
|
|
for tool_block in tool_use_blocks:
|
|
result = self._execute_tool(
|
|
tool_block, tool_executor, turn, trace
|
|
)
|
|
tool_results.append(
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": tool_block.id,
|
|
"content": result,
|
|
}
|
|
)
|
|
|
|
# Add tool results to messages
|
|
messages.append({"role": "user", "content": tool_results})
|
|
|
|
turn += 1
|
|
trace.total_turns = turn
|
|
|
|
# Check if we hit max turns without completion
|
|
if turn >= max_turns and not trace.success:
|
|
trace.success = False
|
|
trace.error = f"Reached maximum turns ({max_turns}) without completion"
|
|
|
|
except Exception as e:
|
|
trace.success = False
|
|
trace.error = str(e)
|
|
|
|
trace.completed_at = datetime.now()
|
|
return trace
|
|
|
|
def _process_response(
|
|
self,
|
|
response: anthropic.types.Message,
|
|
turn: int,
|
|
trace: ReasoningTrace,
|
|
) -> tuple[list, list, list]:
|
|
"""Process response content blocks and update trace."""
|
|
thinking_blocks = []
|
|
text_blocks = []
|
|
tool_use_blocks = []
|
|
|
|
for block in response.content:
|
|
if block.type == "thinking":
|
|
thinking = ThinkingBlock(
|
|
content=block.thinking,
|
|
turn_index=turn,
|
|
signature=getattr(block, "signature", None),
|
|
)
|
|
trace.thinking_blocks.append(thinking)
|
|
thinking_blocks.append(block)
|
|
|
|
elif block.type == "text":
|
|
text_blocks.append(block)
|
|
|
|
elif block.type == "tool_use":
|
|
tool_use_blocks.append(block)
|
|
|
|
# Update token count
|
|
trace.total_tokens += response.usage.input_tokens + response.usage.output_tokens
|
|
|
|
return thinking_blocks, text_blocks, tool_use_blocks
|
|
|
|
def _execute_tool(
|
|
self,
|
|
tool_block: Any,
|
|
executor: Callable[[str, dict], str] | None,
|
|
turn: int,
|
|
trace: ReasoningTrace,
|
|
) -> str:
|
|
"""Execute a tool call and record it in the trace."""
|
|
tool_call = ToolCall(
|
|
id=tool_block.id,
|
|
name=tool_block.name,
|
|
input=tool_block.input,
|
|
turn_index=turn,
|
|
)
|
|
|
|
try:
|
|
if executor:
|
|
result = executor(tool_block.name, tool_block.input)
|
|
else:
|
|
result = f"[Mock result for {tool_block.name}]"
|
|
|
|
tool_call.result = result
|
|
tool_call.success = True
|
|
|
|
except Exception as e:
|
|
result = f"Error: {str(e)}"
|
|
tool_call.result = result
|
|
tool_call.success = False
|
|
tool_call.error = str(e)
|
|
|
|
trace.tool_calls.append(tool_call)
|
|
|
|
# Link thinking to tool call
|
|
if trace.thinking_blocks:
|
|
last_thinking = trace.thinking_blocks[-1]
|
|
if last_thinking.turn_index == turn:
|
|
last_thinking.following_action = f"tool_use:{tool_block.name}"
|
|
|
|
return result
|
|
|
|
def run_streaming(
|
|
self,
|
|
task: str,
|
|
system_prompt: str = "You are a helpful assistant.",
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_executor: Callable[[str, dict], str] | None = None,
|
|
max_turns: int = 10,
|
|
max_tokens: int = 4096,
|
|
on_thinking: Callable[[str], None] | None = None,
|
|
on_text: Callable[[str], None] | None = None,
|
|
on_tool_call: Callable[[str, dict], None] | None = None,
|
|
on_error: Callable[[str], None] | None = None,
|
|
) -> ReasoningTrace:
|
|
"""
|
|
Execute a task with streaming output and capture reasoning trace.
|
|
|
|
Similar to run() but streams thinking and text content in real-time
|
|
via callback functions.
|
|
|
|
Note: For multi-turn tool interactions, the non-streaming run() method
|
|
is recommended as it provides more reliable trace capture. Use this
|
|
method when you need real-time display of thinking/text content.
|
|
|
|
Args:
|
|
task: The user task/query to execute
|
|
system_prompt: System prompt for the agent
|
|
tools: List of tool definitions
|
|
tool_executor: Function to execute tool calls
|
|
max_turns: Maximum conversation turns
|
|
max_tokens: Maximum tokens per response
|
|
on_thinking: Callback for thinking content chunks
|
|
on_text: Callback for text content chunks
|
|
on_tool_call: Callback when tool is called (name, input)
|
|
on_error: Callback when an error occurs (error message)
|
|
|
|
Returns:
|
|
ReasoningTrace containing the full captured trace
|
|
"""
|
|
trace = ReasoningTrace(
|
|
session_id=str(uuid.uuid4()),
|
|
task=task,
|
|
system_prompt=system_prompt,
|
|
model=self.model,
|
|
started_at=datetime.now(),
|
|
)
|
|
|
|
messages = [{"role": "user", "content": task}]
|
|
turn = 0
|
|
|
|
try:
|
|
while turn < max_turns:
|
|
params = {
|
|
"model": self.model,
|
|
"max_tokens": max_tokens,
|
|
"system": system_prompt,
|
|
"messages": messages,
|
|
"stream": True,
|
|
}
|
|
if tools:
|
|
params["tools"] = tools
|
|
|
|
# Collect streamed content
|
|
thinking_buffer = ""
|
|
text_buffer = ""
|
|
tool_use_blocks = []
|
|
current_content = []
|
|
|
|
with self.client.messages.stream(**params) as stream:
|
|
for event in stream:
|
|
if event.type == "content_block_start":
|
|
if hasattr(event, "content_block"):
|
|
current_content.append(event.content_block)
|
|
|
|
elif event.type == "content_block_delta":
|
|
if hasattr(event, "delta"):
|
|
if event.delta.type == "thinking_delta":
|
|
chunk = event.delta.thinking
|
|
thinking_buffer += chunk
|
|
if on_thinking:
|
|
on_thinking(chunk)
|
|
|
|
elif event.delta.type == "text_delta":
|
|
chunk = event.delta.text
|
|
text_buffer += chunk
|
|
if on_text:
|
|
on_text(chunk)
|
|
|
|
# Get final message for tool_use blocks
|
|
final_message = stream.get_final_message()
|
|
for block in final_message.content:
|
|
if block.type == "tool_use":
|
|
tool_use_blocks.append(block)
|
|
if on_tool_call:
|
|
on_tool_call(block.name, block.input)
|
|
|
|
# Record thinking block
|
|
if thinking_buffer:
|
|
trace.thinking_blocks.append(
|
|
ThinkingBlock(
|
|
content=thinking_buffer,
|
|
turn_index=turn,
|
|
)
|
|
)
|
|
|
|
# Update tokens
|
|
trace.total_tokens += (
|
|
final_message.usage.input_tokens + final_message.usage.output_tokens
|
|
)
|
|
|
|
# If no tool calls, we're done
|
|
if not tool_use_blocks:
|
|
trace.final_response = text_buffer or None
|
|
trace.success = True
|
|
break
|
|
|
|
# Append to history
|
|
messages.append({"role": "assistant", "content": final_message.content})
|
|
|
|
# Execute tools
|
|
tool_results = []
|
|
for tool_block in tool_use_blocks:
|
|
result = self._execute_tool(tool_block, tool_executor, turn, trace)
|
|
tool_results.append(
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": tool_block.id,
|
|
"content": result,
|
|
}
|
|
)
|
|
|
|
messages.append({"role": "user", "content": tool_results})
|
|
turn += 1
|
|
trace.total_turns = turn
|
|
|
|
if turn >= max_turns and not trace.success:
|
|
trace.success = False
|
|
trace.error = f"Reached maximum turns ({max_turns})"
|
|
|
|
except Exception as e:
|
|
trace.success = False
|
|
trace.error = str(e)
|
|
if on_error:
|
|
on_error(str(e))
|
|
|
|
trace.completed_at = datetime.now()
|
|
return trace
|
|
|
|
|
|
def format_trace_for_display(trace: ReasoningTrace) -> str:
|
|
"""Format a reasoning trace for human-readable display."""
|
|
lines = [
|
|
f"Session: {trace.session_id}",
|
|
f"Task: {trace.task}",
|
|
f"Model: {trace.model}",
|
|
f"Status: {'Success' if trace.success else 'Failed'}",
|
|
f"Turns: {trace.total_turns}",
|
|
f"Tokens: {trace.total_tokens}",
|
|
"",
|
|
"=" * 60,
|
|
"REASONING TRACE",
|
|
"=" * 60,
|
|
]
|
|
|
|
for i, thinking in enumerate(trace.thinking_blocks):
|
|
lines.append(f"\n[Turn {thinking.turn_index}] Thinking:")
|
|
lines.append("-" * 40)
|
|
lines.append(thinking.content[:500] + "..." if len(thinking.content) > 500 else thinking.content)
|
|
|
|
# Show tool calls at this turn
|
|
turn_tools = trace.get_tool_calls_at_turn(thinking.turn_index)
|
|
for tool in turn_tools:
|
|
lines.append(f"\n Tool: {tool.name}({json.dumps(tool.input)})")
|
|
lines.append(f" Result: {tool.result[:100]}..." if tool.result and len(tool.result) > 100 else f" Result: {tool.result}")
|
|
|
|
if trace.final_response:
|
|
lines.append("\n" + "=" * 60)
|
|
lines.append("FINAL RESPONSE")
|
|
lines.append("=" * 60)
|
|
lines.append(trace.final_response)
|
|
|
|
if trace.error:
|
|
lines.append("\n" + "=" * 60)
|
|
lines.append("ERROR")
|
|
lines.append("=" * 60)
|
|
lines.append(trace.error)
|
|
|
|
return "\n".join(lines)
|