1"""
2Example 2: Tool Usage with Trace Capture
3 
4Demonstrates how M2.1's interleaved thinking reasons between tool calls.
5This is where interleaved thinking really shines - you can see the model
6adapting to tool outputs in real-time.
7"""
8 
9import json
10import os
11from pathlib import Path
12from dotenv import load_dotenv
13 
14from reasoning_trace_optimizer import TraceCapture
15from reasoning_trace_optimizer.capture import format_trace_for_display
16 
17# Load environment variables from the project root
18env_path = Path(__file__).parent.parent / ".env"
19load_dotenv(env_path)
20 
21 
22# Define mock tools
23TOOLS = [
24    {
25        "name": "get_weather",
26        "description": "Get current weather for a location. Returns temperature and conditions.",
27        "input_schema": {
28            "type": "object",
29            "properties": {
30                "location": {
31                    "type": "string",
32                    "description": "City name, e.g., 'San Francisco, CA'",
33                }
34            },
35            "required": ["location"],
36        },
37    },
38    {
39        "name": "get_forecast",
40        "description": "Get 3-day weather forecast for a location.",
41        "input_schema": {
42            "type": "object",
43            "properties": {
44                "location": {
45                    "type": "string",
46                    "description": "City name",
47                },
48                "days": {
49                    "type": "integer",
50                    "description": "Number of days (1-3)",
51                    "default": 3,
52                },
53            },
54            "required": ["location"],
55        },
56    },
57]
58 
59 
60# Mock tool executor
61def execute_tool(name: str, input_data: dict) -> str:
62    """Execute a mock tool and return results."""
63    if name == "get_weather":
64        location = input_data.get("location", "Unknown")
65        # Simulate different weather for different cities
66        if "san francisco" in location.lower():
67            return json.dumps({
68                "location": location,
69                "temperature": "18°C",
70                "conditions": "Foggy",
71                "humidity": "85%",
72            })
73        elif "new york" in location.lower():
74            return json.dumps({
75                "location": location,
76                "temperature": "22°C",
77                "conditions": "Partly cloudy",
78                "humidity": "60%",
79            })
80        else:
81            return json.dumps({
82                "location": location,
83                "temperature": "20°C",
84                "conditions": "Clear",
85                "humidity": "50%",
86            })
87 
88    elif name == "get_forecast":
89        location = input_data.get("location", "Unknown")
90        days = input_data.get("days", 3)
91        forecast = []
92        for i in range(days):
93            forecast.append({
94                "day": i + 1,
95                "high": f"{20 + i * 2}°C",
96                "low": f"{12 + i}°C",
97                "conditions": ["Sunny", "Cloudy", "Rainy"][i % 3],
98            })
99        return json.dumps({
100            "location": location,
101            "forecast": forecast,
102        })
103 
104    return json.dumps({"error": f"Unknown tool: {name}"})
105 
106 
107def main():
108    """Run a task with tools and observe interleaved thinking."""
109 
110    capture = TraceCapture(
111        api_key=os.getenv("ANTHROPIC_API_KEY"),
112        base_url="https://api.minimax.io/anthropic",
113        model="MiniMax-M2.1",
114    )
115 
116    task = """Compare the current weather in San Francisco and New York City.
117    Then tell me which city would be better for outdoor activities this weekend."""
118 
119    system_prompt = """You are a helpful weather assistant.
120    Use the available tools to get accurate weather information.
121    Always provide specific data to support your recommendations."""
122 
123    print("=" * 60)
124    print("TOOL USAGE WITH INTERLEAVED THINKING")
125    print("=" * 60)
126    print(f"\nTask: {task}")
127    print(f"\nTools available: {', '.join(t['name'] for t in TOOLS)}")
128    print("\nCapturing trace with tool usage...\n")
129 
130    # Capture the trace (using non-streaming for reliability)
131    trace = capture.run(
132        task=task,
133        system_prompt=system_prompt,
134        tools=TOOLS,
135        tool_executor=execute_tool,
136        max_turns=10,
137    )
138 
139    print("\n\n" + "=" * 60)
140    print("TRACE ANALYSIS")
141    print("=" * 60)
142 
143    print(f"\nSuccess: {trace.success}")
144    print(f"Total Turns: {trace.total_turns}")
145    print(f"Thinking Blocks: {len(trace.thinking_blocks)}")
146    print(f"Tool Calls: {len(trace.tool_calls)}")
147 
148    # Show how thinking evolved between tool calls
149    print("\n" + "=" * 60)
150    print("THINKING EVOLUTION ACROSS TOOL CALLS")
151    print("=" * 60)
152 
153    for i, thinking in enumerate(trace.thinking_blocks):
154        print(f"\n[Turn {thinking.turn_index}] Thinking Block {i + 1}")
155        print("-" * 40)
156 
157        # Show what tool was called after this thinking
158        turn_tools = trace.get_tool_calls_at_turn(thinking.turn_index)
159        if turn_tools:
160            print(f"Following action: Called {', '.join(t.name for t in turn_tools)}")
161        else:
162            print("Following action: Generated response")
163 
164        # Show key reasoning points (first 300 chars)
165        print(f"\nReasoning preview:\n{thinking.content[:300]}...")
166 
167    # Show tool call results
168    print("\n" + "=" * 60)
169    print("TOOL CALL SUMMARY")
170    print("=" * 60)
171 
172    for tc in trace.tool_calls:
173        status = "✅" if tc.success else "❌"
174        print(f"\n{status} {tc.name}")
175        print(f"   Input: {json.dumps(tc.input)}")
176        print(f"   Result: {tc.result[:100]}..." if tc.result and len(tc.result) > 100 else f"   Result: {tc.result}")
177 
178    # Final response
179    if trace.final_response:
180        print("\n" + "=" * 60)
181        print("FINAL RESPONSE")
182        print("=" * 60)
183        print(trace.final_response)
184 
185 
186if __name__ == "__main__":
187    main()
188