1"""
2TraceCapture: Wraps M2.1 API to capture interleaved thinking traces.
3 
4This module provides the core functionality for executing agent tasks
5through MiniMax M2.1 while capturing all reasoning traces for analysis.
6"""
7 
8import json
9import os
10import uuid
11from datetime import datetime
12from typing import Any, Callable
13 
14import anthropic
15 
16from reasoning_trace_optimizer.models import (
17    ReasoningTrace,
18    ThinkingBlock,
19    ToolCall,
20)
21 
22 
23class TraceCapture:
24    """
25    Captures reasoning traces from MiniMax M2.1's interleaved thinking.
26 
27    This class wraps the Anthropic SDK configured for M2.1 and captures
28    all thinking blocks, tool calls, and responses during agent execution.
29 
30    Example:
31        ```python
32        capture = TraceCapture()
33        trace = capture.run(
34            task="What's the weather in San Francisco?",
35            tools=[weather_tool],
36            tool_executor=execute_tool
37        )
38        print(f"Captured {len(trace.thinking_blocks)} thinking blocks")
39        ```
40    """
41 
42    def __init__(
43        self,
44        api_key: str | None = None,
45        base_url: str = "https://api.minimax.io/anthropic",
46        model: str = "MiniMax-M2.1",
47    ):
48        """
49        Initialize TraceCapture with M2.1 configuration.
50 
51        Args:
52            api_key: MiniMax API key (defaults to ANTHROPIC_API_KEY env var)
53            base_url: API base URL (international or China endpoint)
54            model: Model to use (MiniMax-M2.1, MiniMax-M2.1-lightning, MiniMax-M2)
55        """
56        self.model = model
57        self.client = anthropic.Anthropic(
58            api_key=api_key or os.environ.get("ANTHROPIC_API_KEY"),
59            base_url=base_url,
60        )
61 
62    def run(
63        self,
64        task: str,
65        system_prompt: str = "You are a helpful assistant.",
66        tools: list[dict[str, Any]] | None = None,
67        tool_executor: Callable[[str, dict], str] | None = None,
68        max_turns: int = 10,
69        max_tokens: int = 4096,
70    ) -> ReasoningTrace:
71        """
72        Execute a task and capture the full reasoning trace.
73 
74        Args:
75            task: The user task/query to execute
76            system_prompt: System prompt for the agent
77            tools: List of tool definitions in Anthropic format
78            tool_executor: Function to execute tool calls (name, input) -> result
79            max_turns: Maximum conversation turns before stopping
80            max_tokens: Maximum tokens per response
81 
82        Returns:
83            ReasoningTrace containing all thinking blocks, tool calls, and responses
84        """
85        trace = ReasoningTrace(
86            session_id=str(uuid.uuid4()),
87            task=task,
88            system_prompt=system_prompt,
89            model=self.model,
90            started_at=datetime.now(),
91        )
92 
93        messages = [{"role": "user", "content": task}]
94        turn = 0
95 
96        try:
97            while turn < max_turns:
98                # Build request parameters
99                params = {
100                    "model": self.model,
101                    "max_tokens": max_tokens,
102                    "system": system_prompt,
103                    "messages": messages,
104                }
105                if tools:
106                    params["tools"] = tools
107 
108                # Make API call
109                response = self.client.messages.create(**params)
110 
111                # Process response content blocks
112                thinking_blocks, text_blocks, tool_use_blocks = self._process_response(
113                    response, turn, trace
114                )
115 
116                # If no tool calls, we're done
117                if not tool_use_blocks:
118                    trace.final_response = (
119                        text_blocks[0].text if text_blocks else None
120                    )
121                    trace.success = True
122                    break
123 
124                # Append assistant response to history (CRITICAL for M2.1)
125                messages.append({"role": "assistant", "content": response.content})
126 
127                # Execute tools and collect results
128                tool_results = []
129                for tool_block in tool_use_blocks:
130                    result = self._execute_tool(
131                        tool_block, tool_executor, turn, trace
132                    )
133                    tool_results.append(
134                        {
135                            "type": "tool_result",
136                            "tool_use_id": tool_block.id,
137                            "content": result,
138                        }
139                    )
140 
141                # Add tool results to messages
142                messages.append({"role": "user", "content": tool_results})
143 
144                turn += 1
145                trace.total_turns = turn
146 
147            # Check if we hit max turns without completion
148            if turn >= max_turns and not trace.success:
149                trace.success = False
150                trace.error = f"Reached maximum turns ({max_turns}) without completion"
151 
152        except Exception as e:
153            trace.success = False
154            trace.error = str(e)
155 
156        trace.completed_at = datetime.now()
157        return trace
158 
159    def _process_response(
160        self,
161        response: anthropic.types.Message,
162        turn: int,
163        trace: ReasoningTrace,
164    ) -> tuple[list, list, list]:
165        """Process response content blocks and update trace."""
166        thinking_blocks = []
167        text_blocks = []
168        tool_use_blocks = []
169 
170        for block in response.content:
171            if block.type == "thinking":
172                thinking = ThinkingBlock(
173                    content=block.thinking,
174                    turn_index=turn,
175                    signature=getattr(block, "signature", None),
176                )
177                trace.thinking_blocks.append(thinking)
178                thinking_blocks.append(block)
179 
180            elif block.type == "text":
181                text_blocks.append(block)
182 
183            elif block.type == "tool_use":
184                tool_use_blocks.append(block)
185 
186        # Update token count
187        trace.total_tokens += response.usage.input_tokens + response.usage.output_tokens
188 
189        return thinking_blocks, text_blocks, tool_use_blocks
190 
191    def _execute_tool(
192        self,
193        tool_block: Any,
194        executor: Callable[[str, dict], str] | None,
195        turn: int,
196        trace: ReasoningTrace,
197    ) -> str:
198        """Execute a tool call and record it in the trace."""
199        tool_call = ToolCall(
200            id=tool_block.id,
201            name=tool_block.name,
202            input=tool_block.input,
203            turn_index=turn,
204        )
205 
206        try:
207            if executor:
208                result = executor(tool_block.name, tool_block.input)
209            else:
210                result = f"[Mock result for {tool_block.name}]"
211 
212            tool_call.result = result
213            tool_call.success = True
214 
215        except Exception as e:
216            result = f"Error: {str(e)}"
217            tool_call.result = result
218            tool_call.success = False
219            tool_call.error = str(e)
220 
221        trace.tool_calls.append(tool_call)
222 
223        # Link thinking to tool call
224        if trace.thinking_blocks:
225            last_thinking = trace.thinking_blocks[-1]
226            if last_thinking.turn_index == turn:
227                last_thinking.following_action = f"tool_use:{tool_block.name}"
228 
229        return result
230 
231    def run_streaming(
232        self,
233        task: str,
234        system_prompt: str = "You are a helpful assistant.",
235        tools: list[dict[str, Any]] | None = None,
236        tool_executor: Callable[[str, dict], str] | None = None,
237        max_turns: int = 10,
238        max_tokens: int = 4096,
239        on_thinking: Callable[[str], None] | None = None,
240        on_text: Callable[[str], None] | None = None,
241        on_tool_call: Callable[[str, dict], None] | None = None,
242        on_error: Callable[[str], None] | None = None,
243    ) -> ReasoningTrace:
244        """
245        Execute a task with streaming output and capture reasoning trace.
246 
247        Similar to run() but streams thinking and text content in real-time
248        via callback functions.
249 
250        Note: For multi-turn tool interactions, the non-streaming run() method
251        is recommended as it provides more reliable trace capture. Use this
252        method when you need real-time display of thinking/text content.
253 
254        Args:
255            task: The user task/query to execute
256            system_prompt: System prompt for the agent
257            tools: List of tool definitions
258            tool_executor: Function to execute tool calls
259            max_turns: Maximum conversation turns
260            max_tokens: Maximum tokens per response
261            on_thinking: Callback for thinking content chunks
262            on_text: Callback for text content chunks
263            on_tool_call: Callback when tool is called (name, input)
264            on_error: Callback when an error occurs (error message)
265 
266        Returns:
267            ReasoningTrace containing the full captured trace
268        """
269        trace = ReasoningTrace(
270            session_id=str(uuid.uuid4()),
271            task=task,
272            system_prompt=system_prompt,
273            model=self.model,
274            started_at=datetime.now(),
275        )
276 
277        messages = [{"role": "user", "content": task}]
278        turn = 0
279 
280        try:
281            while turn < max_turns:
282                params = {
283                    "model": self.model,
284                    "max_tokens": max_tokens,
285                    "system": system_prompt,
286                    "messages": messages,
287                    "stream": True,
288                }
289                if tools:
290                    params["tools"] = tools
291 
292                # Collect streamed content
293                thinking_buffer = ""
294                text_buffer = ""
295                tool_use_blocks = []
296                current_content = []
297 
298                with self.client.messages.stream(**params) as stream:
299                    for event in stream:
300                        if event.type == "content_block_start":
301                            if hasattr(event, "content_block"):
302                                current_content.append(event.content_block)
303 
304                        elif event.type == "content_block_delta":
305                            if hasattr(event, "delta"):
306                                if event.delta.type == "thinking_delta":
307                                    chunk = event.delta.thinking
308                                    thinking_buffer += chunk
309                                    if on_thinking:
310                                        on_thinking(chunk)
311 
312                                elif event.delta.type == "text_delta":
313                                    chunk = event.delta.text
314                                    text_buffer += chunk
315                                    if on_text:
316                                        on_text(chunk)
317 
318                    # Get final message for tool_use blocks
319                    final_message = stream.get_final_message()
320                    for block in final_message.content:
321                        if block.type == "tool_use":
322                            tool_use_blocks.append(block)
323                            if on_tool_call:
324                                on_tool_call(block.name, block.input)
325 
326                # Record thinking block
327                if thinking_buffer:
328                    trace.thinking_blocks.append(
329                        ThinkingBlock(
330                            content=thinking_buffer,
331                            turn_index=turn,
332                        )
333                    )
334 
335                # Update tokens
336                trace.total_tokens += (
337                    final_message.usage.input_tokens + final_message.usage.output_tokens
338                )
339 
340                # If no tool calls, we're done
341                if not tool_use_blocks:
342                    trace.final_response = text_buffer or None
343                    trace.success = True
344                    break
345 
346                # Append to history
347                messages.append({"role": "assistant", "content": final_message.content})
348 
349                # Execute tools
350                tool_results = []
351                for tool_block in tool_use_blocks:
352                    result = self._execute_tool(tool_block, tool_executor, turn, trace)
353                    tool_results.append(
354                        {
355                            "type": "tool_result",
356                            "tool_use_id": tool_block.id,
357                            "content": result,
358                        }
359                    )
360 
361                messages.append({"role": "user", "content": tool_results})
362                turn += 1
363                trace.total_turns = turn
364 
365            if turn >= max_turns and not trace.success:
366                trace.success = False
367                trace.error = f"Reached maximum turns ({max_turns})"
368 
369        except Exception as e:
370            trace.success = False
371            trace.error = str(e)
372            if on_error:
373                on_error(str(e))
374 
375        trace.completed_at = datetime.now()
376        return trace
377 
378 
379def format_trace_for_display(trace: ReasoningTrace) -> str:
380    """Format a reasoning trace for human-readable display."""
381    lines = [
382        f"Session: {trace.session_id}",
383        f"Task: {trace.task}",
384        f"Model: {trace.model}",
385        f"Status: {'Success' if trace.success else 'Failed'}",
386        f"Turns: {trace.total_turns}",
387        f"Tokens: {trace.total_tokens}",
388        "",
389        "=" * 60,
390        "REASONING TRACE",
391        "=" * 60,
392    ]
393 
394    for i, thinking in enumerate(trace.thinking_blocks):
395        lines.append(f"\n[Turn {thinking.turn_index}] Thinking:")
396        lines.append("-" * 40)
397        lines.append(thinking.content[:500] + "..." if len(thinking.content) > 500 else thinking.content)
398 
399        # Show tool calls at this turn
400        turn_tools = trace.get_tool_calls_at_turn(thinking.turn_index)
401        for tool in turn_tools:
402            lines.append(f"\n  Tool: {tool.name}({json.dumps(tool.input)})")
403            lines.append(f"  Result: {tool.result[:100]}..." if tool.result and len(tool.result) > 100 else f"  Result: {tool.result}")
404 
405    if trace.final_response:
406        lines.append("\n" + "=" * 60)
407        lines.append("FINAL RESPONSE")
408        lines.append("=" * 60)
409        lines.append(trace.final_response)
410 
411    if trace.error:
412        lines.append("\n" + "=" * 60)
413        lines.append("ERROR")
414        lines.append("=" * 60)
415        lines.append(trace.error)
416 
417    return "\n".join(lines)
418