1"""
2CLI interface for Reasoning Trace Optimizer.
3 
4Provides command-line access to the optimization tools.
5"""
6 
7import argparse
8import json
9import sys
10from pathlib import Path
11 
12from rich.console import Console
13 
14from reasoning_trace_optimizer.analyzer import TraceAnalyzer, format_analysis_report
15from reasoning_trace_optimizer.capture import TraceCapture, format_trace_for_display
16from reasoning_trace_optimizer.loop import OptimizationLoop, LoopConfig
17from reasoning_trace_optimizer.skill_generator import SkillGenerator
18 
19 
20console = Console()
21 
22 
23def cmd_capture(args: argparse.Namespace) -> None:
24    """Run a task and capture reasoning trace."""
25    capture = TraceCapture(
26        api_key=args.api_key,
27        base_url=args.base_url,
28        model=args.model,
29    )
30 
31    console.print(f"[cyan]Capturing trace for task: {args.task}[/cyan]")
32 
33    trace = capture.run(
34        task=args.task,
35        system_prompt=args.system_prompt or "You are a helpful assistant.",
36        max_turns=args.max_turns,
37    )
38 
39    # Output trace
40    output = format_trace_for_display(trace)
41    if args.output:
42        Path(args.output).write_text(output)
43        console.print(f"[green]Trace saved to: {args.output}[/green]")
44    else:
45        console.print(output)
46 
47 
48def cmd_analyze(args: argparse.Namespace) -> None:
49    """Analyze a captured reasoning trace."""
50    # For now, run capture + analyze together
51    # In future, could load trace from file
52 
53    capture = TraceCapture(
54        api_key=args.api_key,
55        base_url=args.base_url,
56        model=args.model,
57    )
58    analyzer = TraceAnalyzer(
59        api_key=args.api_key,
60        base_url=args.base_url,
61        model=args.model,
62    )
63 
64    console.print(f"[cyan]Capturing and analyzing: {args.task}[/cyan]")
65 
66    trace = capture.run(
67        task=args.task,
68        system_prompt=args.system_prompt or "You are a helpful assistant.",
69    )
70 
71    analysis = analyzer.analyze(trace)
72 
73    # Output analysis
74    output = format_analysis_report(analysis)
75    if args.output:
76        Path(args.output).write_text(output)
77        console.print(f"[green]Analysis saved to: {args.output}[/green]")
78    else:
79        console.print(output)
80 
81 
82def cmd_optimize(args: argparse.Namespace) -> None:
83    """Run full optimization loop."""
84    config = LoopConfig(
85        max_iterations=args.max_iterations,
86        convergence_threshold=args.convergence_threshold,
87        min_score_threshold=args.min_score,
88        save_artifacts=True,
89        artifacts_dir=args.artifacts_dir,
90        verbose=True,
91    )
92 
93    loop = OptimizationLoop(
94        config=config,
95        api_key=args.api_key,
96        base_url=args.base_url,
97        model=args.model,
98    )
99 
100    console.print(f"[cyan]Starting optimization for: {args.task}[/cyan]")
101 
102    result = loop.run(
103        task=args.task,
104        initial_prompt=args.system_prompt or "You are a helpful assistant.",
105    )
106 
107    # Output final prompt
108    if args.output:
109        Path(args.output).write_text(result.final_prompt)
110        console.print(f"[green]Optimized prompt saved to: {args.output}[/green]")
111 
112    # Generate skill if requested
113    if args.generate_skill:
114        generator = SkillGenerator(
115            api_key=args.api_key,
116            base_url=args.base_url,
117            model=args.model,
118        )
119        skill_path = generator.generate(
120            result=result,
121            skill_name=args.skill_name or "optimized-agent",
122            output_dir=args.skills_dir,
123        )
124        console.print(f"[green]Generated skill at: {skill_path}[/green]")
125 
126 
127def cmd_generate_skill(args: argparse.Namespace) -> None:
128    """Generate a skill from optimization artifacts."""
129    # Load summary from artifacts
130    artifacts_dir = Path(args.artifacts_dir)
131    summary_path = artifacts_dir / "summary.json"
132 
133    if not summary_path.exists():
134        console.print("[red]Error: No optimization summary found. Run optimize first.[/red]")
135        sys.exit(1)
136 
137    with open(summary_path) as f:
138        summary = json.load(f)
139 
140    # Create minimal loop result from summary
141    from reasoning_trace_optimizer.models import LoopResult, LoopIteration, ReasoningTrace, AnalysisResult
142 
143    # Load final prompt
144    final_prompt_path = artifacts_dir / "final_prompt.txt"
145    final_prompt = final_prompt_path.read_text() if final_prompt_path.exists() else ""
146 
147    result = LoopResult(
148        task=summary.get("task", "Unknown task"),
149        final_prompt=final_prompt,
150        total_iterations=summary.get("total_iterations", 0),
151        initial_score=summary.get("initial_score", 0),
152        final_score=summary.get("final_score", 0),
153        improvement_percentage=summary.get("improvement_percentage", 0),
154        converged=summary.get("converged", False),
155    )
156 
157    generator = SkillGenerator(
158        api_key=args.api_key,
159        base_url=args.base_url,
160        model=args.model,
161    )
162 
163    skill_path = generator.generate(
164        result=result,
165        skill_name=args.skill_name,
166        output_dir=args.output_dir,
167    )
168 
169    console.print(f"[green]Generated skill at: {skill_path}[/green]")
170 
171 
172def main() -> None:
173    """Main CLI entry point."""
174    parser = argparse.ArgumentParser(
175        prog="rto",
176        description="Reasoning Trace Optimizer - Debug and optimize AI agents using M2.1's interleaved thinking",
177    )
178 
179    # Global arguments
180    parser.add_argument(
181        "--api-key",
182        help="MiniMax API key (or set ANTHROPIC_API_KEY env var)",
183    )
184    parser.add_argument(
185        "--base-url",
186        default="https://api.minimax.io/anthropic",
187        help="API base URL",
188    )
189    parser.add_argument(
190        "--model",
191        default="MiniMax-M2.1",
192        choices=["MiniMax-M2.1", "MiniMax-M2.1-lightning", "MiniMax-M2"],
193        help="Model to use",
194    )
195 
196    subparsers = parser.add_subparsers(dest="command", required=True)
197 
198    # Capture command
199    capture_parser = subparsers.add_parser(
200        "capture",
201        help="Capture reasoning trace for a task",
202    )
203    capture_parser.add_argument("task", help="Task to execute")
204    capture_parser.add_argument("--system-prompt", "-s", help="System prompt")
205    capture_parser.add_argument("--max-turns", type=int, default=10)
206    capture_parser.add_argument("--output", "-o", help="Output file path")
207    capture_parser.set_defaults(func=cmd_capture)
208 
209    # Analyze command
210    analyze_parser = subparsers.add_parser(
211        "analyze",
212        help="Capture and analyze reasoning trace",
213    )
214    analyze_parser.add_argument("task", help="Task to analyze")
215    analyze_parser.add_argument("--system-prompt", "-s", help="System prompt")
216    analyze_parser.add_argument("--output", "-o", help="Output file path")
217    analyze_parser.set_defaults(func=cmd_analyze)
218 
219    # Optimize command
220    optimize_parser = subparsers.add_parser(
221        "optimize",
222        help="Run full optimization loop",
223    )
224    optimize_parser.add_argument("task", help="Task to optimize for")
225    optimize_parser.add_argument("--system-prompt", "-s", help="Initial system prompt")
226    optimize_parser.add_argument("--max-iterations", type=int, default=5)
227    optimize_parser.add_argument("--convergence-threshold", type=float, default=5.0)
228    optimize_parser.add_argument("--min-score", type=float, default=80.0)
229    optimize_parser.add_argument(
230        "--artifacts-dir",
231        default="./optimization_artifacts",
232        help="Directory for artifacts",
233    )
234    optimize_parser.add_argument("--output", "-o", help="Output file for final prompt")
235    optimize_parser.add_argument(
236        "--generate-skill",
237        action="store_true",
238        help="Generate Agent Skill from results",
239    )
240    optimize_parser.add_argument("--skill-name", help="Name for generated skill")
241    optimize_parser.add_argument(
242        "--skills-dir",
243        default="./generated_skills",
244        help="Directory for generated skills",
245    )
246    optimize_parser.set_defaults(func=cmd_optimize)
247 
248    # Generate skill command
249    skill_parser = subparsers.add_parser(
250        "generate-skill",
251        help="Generate skill from optimization artifacts",
252    )
253    skill_parser.add_argument("skill_name", help="Name for the skill")
254    skill_parser.add_argument(
255        "--artifacts-dir",
256        default="./optimization_artifacts",
257        help="Directory with optimization artifacts",
258    )
259    skill_parser.add_argument(
260        "--output-dir",
261        default="./generated_skills",
262        help="Output directory for skill",
263    )
264    skill_parser.set_defaults(func=cmd_generate_skill)
265 
266    args = parser.parse_args()
267    args.func(args)
268 
269 
270if __name__ == "__main__":
271    main()
272