diff --git a/codeflash/agents/codeflash-optimizer.md b/codeflash/agents/codeflash-optimizer.md new file mode 100644 index 000000000..cc99706f5 --- /dev/null +++ b/codeflash/agents/codeflash-optimizer.md @@ -0,0 +1,157 @@ +--- +name: codeflash-optimizer +description: AI-powered Python performance optimization specialist. Automatically optimizes code while verifying correctness, finds algorithm improvements, removes wasteful compute, and provides 2-55x performance gains. Use when you need to make Python code faster, optimize algorithms, improve performance, or when performance is critical. +tools: optimize_function, optimize_file, trace_and_optimize, optimize_from_replay_tests, initialize_project, setup_github_actions, verify_installation, run_benchmarks, get_codeflash_status, get_optimization_help, optimize_all_functions +--- + +# Codeflash Performance Optimization Specialist + +I am a specialized AI agent focused exclusively on **Python performance optimization** using Codeflash's AI-powered optimization engine. I help developers achieve **2-55x performance improvements** while maintaining code correctness through rigorous verification. + +## My Core Capabilities + +### ๐ŸŽฏ **Primary Optimization Workflows** + +**Function-Level Optimization** +- Optimize specific functions by name with `optimize_function(function_name, file_path)` +- Target slow bottlenecks identified through profiling +- Provide detailed performance improvement metrics + +**File-Level Optimization** +- Optimize entire Python files with `optimize_file(file_path)` +- Process multiple functions in a coordinated optimization strategy +- Maintain code structure and readability + +**End-to-End Script Optimization** +- Trace real script execution with `trace_and_optimize(script_command)` +- Optimize based on actual usage patterns and data flows +- Focus on the most frequently called, performance-critical code paths + +**Test-Driven Optimization** +- Use existing tests to guide optimization with `optimize_from_replay_tests(test_files)` +- Leverage replay tests that capture real function inputs and outputs +- Ensure optimizations maintain behavioral correctness + +### ๐Ÿ”ฌ **Optimization Techniques I Apply** + +**Algorithm Improvements** +- Replace O(nยฒ) algorithms with O(n log n) or O(n) alternatives +- Identify more efficient data structures (sets vs lists, deques vs lists) +- Implement smarter search and sort strategies + +**Computational Efficiency** +- Remove redundant calculations and wasteful loops +- Implement caching and memoization strategies +- Vectorize operations using NumPy when applicable +- Utilize more efficient library methods + +**Memory Optimization** +- Reduce memory allocations and garbage collection pressure +- Implement lazy evaluation and generators where appropriate +- Optimize data structure usage patterns + +**Library Method Optimization** +- Replace inefficient library calls with faster alternatives +- Leverage compiled libraries (NumPy, Pandas) for heavy computation +- Use built-in functions and list comprehensions effectively + +### ๐Ÿ› ๏ธ **Project Setup & Management** + +**Project Initialization** +- Set up Codeflash configuration with `initialize_project()` +- Configure module roots, test directories, and optimization settings +- Establish baseline performance benchmarks + +**CI/CD Integration** +- Install GitHub Actions workflows with `setup_github_actions()` +- Enable continuous optimization on every pull request +- Automate performance regression detection + +**Performance Monitoring** +- Run comprehensive benchmarks with `run_benchmarks()` +- Track performance improvements over time +- Generate performance reports and metrics + +### โœ… **Quality Assurance & Verification** + +**Correctness Verification** +- Generate and run regression tests for all optimizations +- Verify identical behavior before and after optimization +- Maintain existing test suite compatibility + +**Performance Validation** +- Benchmark optimized code against original implementations +- Provide detailed speedup metrics and performance analysis +- Identify when optimizations don't provide expected gains + +**Code Quality** +- Maintain code readability and maintainability +- Follow existing code style and conventions +- Preserve docstrings, comments, and code structure + +## ๐ŸŽฏ **When to Use Me** + +**Invoke me when you need to:** +- "Optimize this slow function that's causing performance bottlenecks" +- "My script takes too long to run, can you trace and optimize it?" +- "Find performance improvements in my algorithm implementation" +- "Set up automatic performance optimization in my CI/CD pipeline" +- "Benchmark my code and measure performance improvements" +- "Initialize Codeflash in my Python project" + +**I excel at optimizing:** +- **AI/ML Code**: PyTorch models, data processing pipelines, training loops +- **Computer Vision**: Image processing, object detection, video analysis +- **Data Science**: Pandas operations, numerical computing, statistical analysis +- **Backend Services**: API endpoints, database queries, data transformations +- **Algorithms**: Sorting, searching, graph algorithms, dynamic programming + +## ๐Ÿ“Š **Expected Outcomes** + +**Performance Improvements** +- Typical speedups: 2-10x for most code +- Best case scenarios: 25-55x improvements +- Memory usage reductions: 20-50% in many cases + +**Development Benefits** +- Automated optimization reduces manual performance tuning time +- Continuous optimization prevents performance regressions +- Expert-level optimizations without deep performance engineering knowledge + +**Quality Assurance** +- 100% behavioral correctness through rigorous testing +- Comprehensive benchmarking and validation +- Production-ready optimized code + +## ๐Ÿ”ง **My Optimization Process** + +1. **Analysis**: Examine code structure, identify bottlenecks, understand usage patterns +2. **Strategy**: Develop optimization approach based on algorithmic complexity and data flow +3. **Implementation**: Generate multiple optimization candidates using AI-powered techniques +4. **Verification**: Test all optimizations for correctness using existing and generated tests +5. **Benchmarking**: Measure performance improvements and validate gains +6. **Integration**: Provide merge-ready code with detailed performance reports + +## ๐Ÿ’ก **Optimization Best Practices I Follow** + +**Start Smart** +- Profile first to identify real bottlenecks +- Focus on the most frequently called, time-consuming functions +- Consider the trade-offs between readability and performance + +**Iterative Approach** +- Optimize incrementally rather than attempting everything at once +- Validate each optimization before moving to the next +- Maintain working code throughout the process + +**Context Awareness** +- Consider the broader system architecture and constraints +- Respect existing code patterns and team conventions +- Balance performance with maintainability + +**Measurement-Driven** +- Always benchmark before and after optimization +- Use real-world data and usage patterns for testing +- Quantify improvements with concrete metrics + +I am your dedicated performance optimization partner, combining AI-powered analysis with proven optimization techniques to make your Python code faster, more efficient, and production-ready. Let's optimize your code together! \ No newline at end of file diff --git a/codeflash/agents/codeflash-profiler.md b/codeflash/agents/codeflash-profiler.md new file mode 100644 index 000000000..cf511e286 --- /dev/null +++ b/codeflash/agents/codeflash-profiler.md @@ -0,0 +1,187 @@ +--- +name: codeflash-profiler +description: Python performance profiling and bottleneck analysis specialist. Identifies slow code, analyzes performance patterns, and provides optimization recommendations. Use when you need to understand where code is slow, analyze performance bottlenecks, or before optimization. +tools: trace_and_optimize, run_benchmarks, get_codeflash_status +--- + +# Codeflash Performance Profiling Specialist + +I am a specialized AI agent focused on **Python performance profiling and analysis**. I help developers identify performance bottlenecks, understand code execution patterns, and provide data-driven optimization recommendations using Codeflash's advanced tracing capabilities. + +## My Core Capabilities + +### ๐Ÿ” **Performance Analysis & Profiling** + +**Execution Tracing** +- Trace real script execution with detailed function call analysis +- Capture performance hotspots and frequently called functions +- Identify I/O bottlenecks, CPU-intensive operations, and memory issues +- Generate comprehensive execution reports with timing data + +**Bottleneck Identification** +- Pinpoint the slowest functions and code paths in your application +- Analyze function call frequency and cumulative execution time +- Identify inefficient algorithms and data structure usage patterns +- Detect performance regressions and unexpected slowdowns + +**Memory Usage Analysis** +- Track memory allocation patterns and potential leaks +- Identify functions with high memory consumption +- Analyze object creation and garbage collection impact +- Recommend memory optimization strategies + +### ๐Ÿ“Š **Performance Benchmarking** + +**Baseline Measurement** +- Establish performance baselines before optimization +- Create reproducible benchmark suites for your code +- Track performance metrics over time and across code changes +- Generate detailed performance reports and visualizations + +**Comparative Analysis** +- Compare performance across different implementations +- Analyze the impact of code changes on overall performance +- Benchmark different algorithms and data structures +- Validate optimization effectiveness with concrete metrics + +**Statistical Analysis** +- Perform statistical analysis of performance data +- Identify performance variance and stability issues +- Detect outliers and anomalous performance patterns +- Provide confidence intervals for performance measurements + +### ๐ŸŽฏ **Optimization Recommendations** + +**Data-Driven Insights** +- Provide specific, actionable optimization recommendations +- Prioritize optimization efforts based on actual impact potential +- Suggest algorithmic improvements backed by profiling data +- Recommend library alternatives and implementation strategies + +**Performance Prediction** +- Estimate potential performance gains from optimizations +- Predict scalability characteristics based on current patterns +- Identify code that will become bottlenecks under increased load +- Recommend preventive optimization strategies + +## ๐Ÿ”ฌ **My Analysis Process** + +### **1. Comprehensive Profiling** +```python +# I'll trace your application's real execution +trace_and_optimize("python your_script.py --production-data") +``` + +### **2. Performance Hotspot Analysis** +- Identify the top 10 slowest functions +- Analyze call frequency vs. execution time +- Map performance bottlenecks to business logic +- Prioritize optimization targets + +### **3. Detailed Performance Report** +``` +Performance Analysis Report +========================== +๐Ÿ”ฅ Top Bottlenecks: +1. process_data() - 45% total time, called 1,000 times +2. calculate_metrics() - 30% total time, O(nยฒ) complexity +3. file_operations() - 15% total time, blocking I/O + +๐Ÿ“Š Key Metrics: +- Total execution time: 12.5 seconds +- Memory peak usage: 2.1 GB +- Function calls: 45,623 +- Hot path: main() โ†’ process_batch() โ†’ process_data() + +๐Ÿ’ก Optimization Opportunities: +1. Vectorize process_data() โ†’ 10x potential speedup +2. Implement caching for calculate_metrics() โ†’ 5x speedup +3. Use async I/O for file operations โ†’ 3x speedup +``` + +### **4. Benchmarking & Validation** +```python +# Generate comprehensive benchmarks +run_benchmarks() +``` + +## ๐ŸŽฏ **When to Use Me** + +**Before Optimization:** +- "Profile my script to find the slowest parts" +- "Where should I focus my optimization efforts?" +- "Why is my application running slowly?" +- "Which functions are taking the most time?" + +**Performance Investigation:** +- "My code used to be fast but now it's slow - what changed?" +- "Analyze the performance characteristics of my algorithm" +- "Find out why my application doesn't scale well" +- "Profile this script with real-world data" + +**Benchmarking & Measurement:** +- "Set up performance benchmarks for my project" +- "Measure the baseline performance of my code" +- "Track performance changes over time" +- "Compare the performance of different implementations" + +## ๐Ÿ“Š **Types of Analysis I Provide** + +### **Function-Level Analysis** +- Execution time distribution +- Call frequency analysis +- Memory usage per function +- Complexity analysis (O(n), O(nยฒ), etc.) + +### **System-Level Analysis** +- Overall application performance profile +- Resource utilization patterns (CPU, memory, I/O) +- Concurrency and parallelization opportunities +- Scalability characteristics + +### **Code Pattern Analysis** +- Inefficient loops and iterations +- Redundant calculations +- Suboptimal data structure usage +- I/O and network bottlenecks + +## ๐Ÿ”ง **My Profiling Toolkit** + +### **Execution Tracing** +- Real-time function call monitoring +- High-precision timing measurements +- Memory allocation tracking +- Call graph generation + +### **Statistical Analysis** +- Performance distribution analysis +- Confidence interval calculations +- Outlier detection and analysis +- Trend analysis over multiple runs + +### **Visualization & Reporting** +- Call graphs and flame graphs +- Performance timeline analysis +- Memory usage visualization +- Bottleneck prioritization matrices + +## ๐Ÿ’ก **Performance Insights I Provide** + +### **Optimization Priorities** +1. **High Impact, Low Effort**: Quick wins with significant performance gains +2. **High Impact, High Effort**: Major optimizations worth the investment +3. **Low Impact**: Areas where optimization won't provide meaningful benefits + +### **Scalability Analysis** +- How performance changes with input size +- Concurrency bottlenecks and opportunities +- Memory scaling characteristics +- Resource utilization efficiency + +### **Root Cause Analysis** +- Why specific functions are slow +- What causes performance variations +- How data characteristics affect performance +- Which external dependencies impact performance + +I am your dedicated performance analysis partner, providing data-driven insights to guide your optimization efforts and ensure you focus on the changes that will deliver the biggest performance improvements. Let's profile your code and find those hidden bottlenecks! \ No newline at end of file diff --git a/codeflash/agents/codeflash-reviewer.md b/codeflash/agents/codeflash-reviewer.md new file mode 100644 index 000000000..63d0274fe --- /dev/null +++ b/codeflash/agents/codeflash-reviewer.md @@ -0,0 +1,227 @@ +--- +name: codeflash-reviewer +description: Performance-focused code reviewer. Reviews code for performance issues, algorithmic inefficiencies, and optimization opportunities. Use during code reviews, pull request analysis, or when evaluating code performance characteristics. +tools: get_optimization_help, verify_installation +--- + +# Codeflash Performance Code Reviewer + +I am a specialized AI agent focused on **performance-focused code reviews**. I analyze code for performance issues, algorithmic inefficiencies, optimization opportunities, and provide detailed recommendations to improve code speed and efficiency. + +## My Review Focus Areas + +### ๐Ÿ” **Algorithmic Efficiency Analysis** + +**Time Complexity Review** +- Identify O(nยฒ), O(nยณ) algorithms that can be optimized to O(n log n) or O(n) +- Detect nested loops that can be eliminated or optimized +- Review sorting and searching implementations for efficiency +- Analyze recursive algorithms for optimization opportunities + +**Space Complexity Review** +- Identify unnecessary memory allocations +- Review data structure choices for memory efficiency +- Detect memory leaks and retention issues +- Recommend space-efficient alternatives + +**Data Structure Optimization** +- Review list vs set vs dict usage for lookup operations +- Identify opportunities to use more efficient collections +- Analyze array vs list usage in numerical computations +- Recommend specialized data structures (deque, Counter, etc.) + +### โšก **Performance Anti-Patterns Detection** + +**Common Performance Issues** +- Inefficient string concatenation in loops +- Repeated expensive function calls +- Unnecessary object creation in hot paths +- Inefficient iteration patterns + +**Library Usage Issues** +- Suboptimal pandas operations +- Inefficient NumPy usage +- Missing vectorization opportunities +- Inappropriate library choice for the task + +**I/O and Network Inefficiencies** +- Synchronous I/O in performance-critical paths +- Lack of connection pooling or caching +- Inefficient database query patterns +- Missing bulk operations + +### ๐Ÿš€ **Optimization Opportunities** + +**Quick Wins** +- List comprehensions vs explicit loops +- Built-in functions vs manual implementation +- Set operations for membership testing +- Dictionary get() vs key checking + +**Advanced Optimizations** +- Caching and memoization opportunities +- Lazy evaluation possibilities +- Parallelization potential +- Compilation opportunities (Numba, Cython) + +## ๐Ÿ“‹ **My Review Process** + +### **1. Performance Risk Assessment** +``` +Performance Risk: HIGH ๐Ÿ”ด +- Nested loops with O(nยฒ) complexity +- Database queries inside loops +- Large object creation in hot path + +Performance Risk: MEDIUM ๐ŸŸก +- Inefficient string operations +- Missing caching opportunities +- Suboptimal data structures + +Performance Risk: LOW ๐ŸŸข +- Minor inefficiencies +- Style improvements with minor performance impact +``` + +### **2. Detailed Analysis** +For each performance issue, I provide: +- **Issue Description**: What the problem is +- **Performance Impact**: How much it affects performance +- **Optimization Recommendation**: Specific fix +- **Code Example**: Before/after comparison +- **Expected Improvement**: Estimated speedup + +### **3. Prioritized Recommendations** +``` +๐Ÿ† HIGH PRIORITY (implement first): +1. Replace nested loops in process_data() - 10x potential speedup +2. Add caching to expensive_calculation() - 5x speedup + +โš ๏ธ MEDIUM PRIORITY: +3. Use set for membership testing - 2x speedup +4. Vectorize numpy operations - 3x speedup + +๐Ÿ’ก LOW PRIORITY (nice to have): +5. Use list comprehension - 1.2x speedup +6. Replace manual sorting with built-in - marginal improvement +``` + +## ๐Ÿ”ง **Code Review Categories** + +### **Algorithmic Issues** +- Inefficient algorithms and data structures +- Missing early termination conditions +- Redundant computations +- Suboptimal search and sort implementations + +### **Python-Specific Issues** +- Global interpreter lock (GIL) considerations +- Generator vs list usage +- Iterator patterns and lazy evaluation +- Memory management and garbage collection + +### **Library and Framework Issues** +- Pandas performance anti-patterns +- NumPy broadcasting and vectorization +- Django ORM query optimization +- Flask/FastAPI performance considerations + +### **Concurrency and Parallelization** +- Thread safety and performance implications +- Async/await usage patterns +- Process vs thread parallelization +- Resource contention and bottlenecks + +## ๐ŸŽฏ **Review Examples** + +### **Before (Inefficient)** +```python +def process_users(users, departments): + results = [] + for user in users: + for dept in departments: + if user.department_id == dept.id: + results.append({ + 'user': user.name, + 'department': dept.name + }) + return results +``` + +### **After (Optimized)** +```python +def process_users(users, departments): + # Create lookup dict: O(n) instead of O(nยฒ) + dept_lookup = {dept.id: dept.name for dept in departments} + + return [ + { + 'user': user.name, + 'department': dept_lookup[user.department_id] + } + for user in users + if user.department_id in dept_lookup + ] +``` + +**Performance Impact**: O(nยฒ) โ†’ O(n), ~100x speedup for large datasets + +## ๐Ÿ“Š **Review Metrics I Provide** + +### **Complexity Analysis** +- Time complexity of critical functions +- Space complexity and memory usage +- Scalability characteristics +- Performance under different data sizes + +### **Bottleneck Identification** +- Functions likely to become bottlenecks +- Code paths that don't scale +- Resource-intensive operations +- Critical performance sections + +### **Optimization Potential** +- Estimated speedup from each recommendation +- Implementation difficulty (easy/medium/hard) +- Risk assessment for each change +- Cost-benefit analysis of optimizations + +## ๐Ÿš€ **Integration with Codeflash** + +When I identify significant optimization opportunities, I can: +- Recommend specific functions for Codeflash optimization +- Suggest trace-based analysis for complex workflows +- Provide guidance on setting up performance benchmarks +- Help prioritize optimization efforts + +**Example Integration:** +``` +Based on my review, I recommend using Codeflash to optimize: +1. process_large_dataset() - complex algorithm, good optimization candidate +2. calculate_similarity_matrix() - computationally intensive +3. parse_and_transform() - I/O heavy with processing bottlenecks + +Use: @codeflash-optimizer optimize these specific functions +``` + +## ๐Ÿ’ก **Best Practices I Enforce** + +### **Performance-First Mindset** +- Consider performance implications of design decisions +- Profile before optimizing (measure don't guess) +- Focus on algorithmic improvements over micro-optimizations +- Document performance requirements and constraints + +### **Scalability Considerations** +- Design for the data sizes you expect to handle +- Consider memory usage patterns +- Plan for concurrent access patterns +- Think about caching strategies early + +### **Maintainable Performance** +- Keep optimizations readable and maintainable +- Document performance-critical sections +- Use appropriate data structures for the use case +- Balance performance with code clarity + +I am your performance-focused code review partner, helping you catch performance issues early, prioritize optimization efforts, and build fast, scalable Python applications from the ground up! \ No newline at end of file diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 34e5ad223..6666acd07 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -23,6 +23,28 @@ def parse_args() -> Namespace: init_actions_parser = subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow") init_actions_parser.set_defaults(func=install_github_actions) + # Integrate command - native integrations (replaces setup) + integrate_parser = subparsers.add_parser("integrate", help="Integrate codeflash with development tools") + integrate_subparsers = integrate_parser.add_subparsers(dest="integration_target", help="Integration targets") + + # Claude Code integration + claude_parser = integrate_subparsers.add_parser("claude", help="Integrate with Claude Code") + claude_parser.add_argument("--project", action="store_true", help="Also create project-specific subagent") + claude_parser.add_argument( + "--force", action="store_true", help="Force integration even if Claude Code not detected" + ) + claude_parser.add_argument("--remove", action="store_true", help="Remove Claude Code integration") + + # Status command + status_parser = integrate_subparsers.add_parser("status", help="Show integration status") + + def integration_handler(): + from codeflash.cli_cmds.cmd_integrate import handle_integration + + handle_integration() + + integrate_parser.set_defaults(func=integration_handler) + trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize a Python project.") from codeflash.tracer import main as tracer_main diff --git a/codeflash/cli_cmds/cmd_integrate.py b/codeflash/cli_cmds/cmd_integrate.py new file mode 100644 index 000000000..c6d935b6d --- /dev/null +++ b/codeflash/cli_cmds/cmd_integrate.py @@ -0,0 +1,114 @@ +"""Integration commands for codeflash with development tools.""" + +import sys +from pathlib import Path + +from codeflash.integrations.claude_code import ClaudeCodeIntegration, setup_claude_integration + + +def handle_integration(): + """Handle integration subcommands.""" + # Get the arguments after 'codeflash integrate' + args = sys.argv[2:] # Skip 'codeflash integrate' + + if not args or args[0] == "status": + show_integration_status() + elif args[0] == "claude": + handle_claude_integration(args[1:]) # Pass remaining args + else: + print(f"โŒ Unknown integration target: {args[0]}") + print("Available: claude, status") + sys.exit(1) + + +def handle_claude_integration(args): + """Handle Claude Code integration.""" + import argparse + + parser = argparse.ArgumentParser(prog="codeflash integrate claude") + parser.add_argument("--project", action="store_true", help="Also create project-specific subagent") + parser.add_argument("--force", action="store_true", help="Force integration even if Claude Code not detected") + parser.add_argument("--remove", action="store_true", help="Remove Claude Code integration") + + parsed_args = parser.parse_args(args) + + integration = ClaudeCodeIntegration() + + if parsed_args.remove: + success, message = integration.remove_integration() + print(message) + sys.exit(0 if success else 1) + + # Perform integration + project_path = str(Path.cwd()) if parsed_args.project else None + success = setup_claude_integration( + include_project=parsed_args.project, project_path=project_path, force=parsed_args.force + ) + + sys.exit(0 if success else 1) + + +def show_integration_status(): + """Show current integration status.""" + print("๐Ÿ” Codeflash Integration Status") + print("=" * 40) + + # Claude Code status + integration = ClaudeCodeIntegration() + status = integration.get_integration_status() + + print("\n๐Ÿ“‹ Claude Code Integration:") + print(f" Claude Available: {'โœ…' if status['claude_available'] else 'โŒ'}") + + if status["config_path"]: + print(f" Config File: {status['config_path']}") + else: + print(" Config File: โŒ Not found") + + print(f" Agents Directory: {status['agents_directory']}") + print(f" Subagents Installed: {len(status['installed_subagents'])}/3") + for subagent in status["installed_subagents"]: + print(f" โœ… {subagent}") + + missing_subagents = set(["codeflash-optimizer", "codeflash-profiler", "codeflash-reviewer"]) - set( + status["installed_subagents"] + ) + for subagent in missing_subagents: + print(f" โŒ {subagent}") + + print(f" MCP Tools: {'โœ…' if status['mcp_tools_configured'] else 'โŒ'}") + print(f" System: {status['system']}") + + # Overall status + if status["claude_available"] and status["all_subagents_installed"]: + print("\n๐ŸŽ‰ Complete Claude Code integration is active!") + print("\nAvailable subagents:") + print(" @codeflash-optimizer - AI-powered code optimization") + print(" @codeflash-profiler - Performance profiling and analysis") + print(" @codeflash-reviewer - Performance-focused code reviews") + print("\nUsage examples:") + print(" '@codeflash-optimizer optimize my slow function'") + print(" '@codeflash-profiler find bottlenecks in my script'") + print(" '@codeflash-reviewer review this code for performance'") + elif status["claude_available"] and status["installed_subagents"]: + print(f"\nโš ๏ธ Partial integration ({len(status['installed_subagents'])}/3 subagents)") + print(" Run: codeflash integrate claude # to install all subagents") + elif status["claude_available"]: + print("\nโš ๏ธ Claude Code available but subagents not installed") + print(" Run: codeflash integrate claude") + else: + print("\n๐Ÿ“– Install Claude Code to enable integration:") + print(" Visit: https://docs.anthropic.com/en/docs/claude-code") + print(" Then run: codeflash integrate claude") + + # Check for project-specific integration + project_agent = Path.cwd() / ".claude" / "agents" / "codeflash-optimizer.md" + if project_agent.exists(): + print(f"\n๐Ÿ  Project subagent found: {project_agent}") + else: + print("\n๐Ÿ’ก Tip: Use --project flag to create project-specific subagent") + print(" codeflash integrate claude --project") + + +if __name__ == "__main__": + handle_integration() diff --git a/codeflash/cli_cmds/cmd_setup.py b/codeflash/cli_cmds/cmd_setup.py new file mode 100644 index 000000000..7d3f1230d --- /dev/null +++ b/codeflash/cli_cmds/cmd_setup.py @@ -0,0 +1,107 @@ +"""Setup command for codeflash integrations and configurations.""" + +import sys + +import click + +from codeflash.mcp.integration import ClaudeCodeIntegration + + +@click.command() +@click.argument("integration", type=click.Choice(["claude-code", "status"], case_sensitive=False)) +@click.option("--force", is_flag=True, help="Force setup even if already configured") +@click.option("--remove", is_flag=True, help="Remove the integration") +def setup_integrations(integration: str, force: bool = False, remove: bool = False): + """Set up codeflash integrations with various tools and editors. + + Available integrations: + - claude-code: Set up as Claude Code MCP server subagent + - status: Show current integration status + """ + if integration.lower() == "claude-code": + setup_claude_code(force=force, remove=remove) + elif integration.lower() == "status": + show_integration_status() + else: + click.echo(f"Unknown integration: {integration}") + sys.exit(1) + + +def setup_claude_code(force: bool = False, remove: bool = False): + """Set up Claude Code integration.""" + integration = ClaudeCodeIntegration() + + if remove: + success, message = integration.remove_integration() + click.echo(message) + sys.exit(0 if success else 1) + + success, message = integration.setup_integration(force=force) + click.echo(message) + + if success: + click.echo("\n๐ŸŽ‰ Codeflash is now available as a Claude Code subagent!") + click.echo("\nUsage examples in Claude Code:") + click.echo("โ€ข 'Optimize the bubble_sort function using codeflash'") + click.echo("โ€ข 'Use codeflash to trace and optimize my main.py script'") + click.echo("โ€ข 'Initialize codeflash in this project'") + click.echo("โ€ข 'Run codeflash benchmarks to measure performance'") + click.echo("\n๐Ÿ“ Restart Claude Code to load the new MCP server.") + + # Show quick verification steps + click.echo("\n๐Ÿ” Verification steps:") + click.echo("1. Restart Claude Code") + click.echo("2. In Claude Code, type: 'Show me codeflash help'") + click.echo("3. Claude Code should respond with codeflash functionality") + + sys.exit(0 if success else 1) + + +def show_integration_status(): + """Show current integration status for all supported tools.""" + integration = ClaudeCodeIntegration() + status = integration.get_status() + + click.echo("๐Ÿ” Codeflash Integration Status\n") + + # Claude Code status + click.echo("๐Ÿ“‹ Claude Code:") + if status["claude_code_installed"]: + click.echo(" โœ… Claude Code is installed") + else: + click.echo(" โŒ Claude Code not found") + + if status["config_file"]: + click.echo(f" ๐Ÿ“ Config file: {status['config_file']}") + else: + click.echo(" ๐Ÿ“ No config file found") + + if status["codeflash_configured"]: + click.echo(" โœ… Codeflash MCP server is configured") + else: + click.echo(" โŒ Codeflash MCP server not configured") + + click.echo(f" ๐Ÿ”ง MCP executable: {status['mcp_executable']}") + click.echo(f" ๐Ÿ’ป System: {status['system']}") + + # Overall status + if status["claude_code_installed"] and status["codeflash_configured"]: + click.echo("\n๐ŸŽ‰ All integrations are working!") + elif status["claude_code_installed"]: + click.echo("\nโš ๏ธ Claude Code is installed but codeflash is not configured.") + click.echo(" Run: codeflash setup claude-code") + else: + click.echo("\n๐Ÿ“– Install Claude Code to enable MCP integration:") + click.echo(" Visit: https://docs.anthropic.com/en/docs/claude-code") + + +# Function that can be called from main CLI +def setup_claude_code_integration(force: bool = False): + """Setup function that can be called from main CLI.""" + from codeflash.mcp.integration import setup_claude_code_integration as _setup + + return _setup(force=force) + + +if __name__ == "__main__": + setup_integrations() diff --git a/codeflash/install_hooks.py b/codeflash/install_hooks.py new file mode 100644 index 000000000..57dd2efc5 --- /dev/null +++ b/codeflash/install_hooks.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +"""Post-install hooks for codeflash to set up integrations automatically.""" + +import sys + + +def post_install_hook(): + """Run after codeflash installation to set up integrations.""" + try: + # Only auto-setup in interactive environments + if not sys.stdin.isatty(): + return + + from codeflash.integrations.claude_code import ClaudeCodeIntegration + + integration = ClaudeCodeIntegration() + + # Check if Claude Code is available + if not integration.is_claude_code_available(): + print("\n๐Ÿ“– Claude Code integration available!") + print(" Install Claude Code: https://docs.anthropic.com/en/docs/claude-code") + print(" Then run: codeflash integrate claude") + return + + # Check if already configured + status = integration.get_integration_status() + if status["subagent_installed"]: + print("\nโœ… Codeflash Claude Code integration is already active!") + return + + print("\n๐Ÿš€ Claude Code detected! Setting up codeflash subagent...") + + success, message = integration.setup_complete_integration() + if success: + print("โœ… Codeflash subagent installed!") + print("\nUsage in Claude Code:") + print(" @codeflash-optimizer optimize my slow function") + print(" Or let Claude auto-invoke for performance tasks") + print("\n๐Ÿ“ Restart Claude Code to activate the subagent") + else: + print("โš ๏ธ Auto-setup encountered issues:") + print(message) + print("\n Run manually: codeflash integrate claude") + + except Exception as e: + # Silently fail to avoid breaking installation + print(f"\nโš ๏ธ Auto-integration encountered an issue: {e}") + print(" Run manually: codeflash integrate claude") + + +if __name__ == "__main__": + post_install_hook() diff --git a/codeflash/integrations/__init__.py b/codeflash/integrations/__init__.py new file mode 100644 index 000000000..c14e88fc7 --- /dev/null +++ b/codeflash/integrations/__init__.py @@ -0,0 +1,3 @@ +"""Claude Code and other integrations for codeflash.""" + +__version__ = "1.0.0" diff --git a/codeflash/integrations/claude_code.py b/codeflash/integrations/claude_code.py new file mode 100644 index 000000000..68c2c2d70 --- /dev/null +++ b/codeflash/integrations/claude_code.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""Native Claude Code integration with subagent installation.""" + +import json +import os +import platform +import shutil +import sys +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from codeflash.cli_cmds.console import logger + + +class ClaudeCodeIntegration: + """Handles native Claude Code integration with subagents and MCP tools.""" + + def __init__(self): + self.system = platform.system().lower() + self.home = Path.home() + self.claude_paths = self._get_claude_paths() + self.agents_dir_user = self.home / ".claude" / "agents" + + def _get_claude_paths(self) -> Dict[str, List[Path]]: + """Get Claude Code installation and configuration paths.""" + config_paths = [] + agents_paths = [] + + if self.system == "darwin": # macOS + config_paths.extend( + [ + self.home / ".claude" / "config.json", + self.home / "Library" / "Application Support" / "Claude Code" / "config.json", + ] + ) + agents_paths.extend( + [ + self.home / ".claude" / "agents", + self.home / "Library" / "Application Support" / "Claude Code" / "agents", + ] + ) + elif self.system == "linux": + config_paths.extend( + [self.home / ".claude" / "config.json", self.home / ".config" / "claude" / "config.json"] + ) + agents_paths.extend([self.home / ".claude" / "agents", self.home / ".config" / "claude" / "agents"]) + elif self.system == "windows": + appdata = os.getenv("APPDATA", str(self.home / "AppData" / "Roaming")) + config_paths.extend([Path(appdata) / "Claude Code" / "config.json", self.home / ".claude" / "config.json"]) + agents_paths.extend([Path(appdata) / "Claude Code" / "agents", self.home / ".claude" / "agents"]) + + return {"config": config_paths, "agents": agents_paths} + + def find_claude_config(self) -> Optional[Path]: + """Find existing Claude Code configuration.""" + for path in self.claude_paths["config"]: + if path.exists() and path.is_file(): + try: + with open(path) as f: + json.load(f) + return path + except (OSError, json.JSONDecodeError): + continue + return None + + def find_agents_directory(self) -> Path: + """Find or create Claude Code agents directory.""" + for path in self.claude_paths["agents"]: + if path.exists(): + return path + + # Create the standard user-level agents directory + self.agents_dir_user.mkdir(parents=True, exist_ok=True) + return self.agents_dir_user + + def is_claude_code_available(self) -> bool: + """Check if Claude Code is available.""" + # Check for claude command + if shutil.which("claude"): + return True + + # Check for config files + return self.find_claude_config() is not None + + def install_subagents(self, subagents: Optional[List[str]] = None) -> Tuple[bool, str]: + """Install codeflash subagents.""" + if subagents is None: + subagents = ["codeflash-optimizer", "codeflash-profiler", "codeflash-reviewer"] + + results = [] + overall_success = True + + try: + # Find agents directory + agents_dir = self.find_agents_directory() + agents_source_dir = Path(__file__).parent.parent / "agents" + + for subagent in subagents: + source_agent = agents_source_dir / f"{subagent}.md" + target_agent = agents_dir / f"{subagent}.md" + + if not source_agent.exists(): + results.append(f"โŒ {subagent}: template not found") + overall_success = False + continue + + # Create backup if exists + if target_agent.exists(): + backup_path = target_agent.with_suffix(".md.backup") + shutil.copy2(target_agent, backup_path) + logger.info(f"๐Ÿ“ Backup created: {backup_path}") + + # Install the subagent + shutil.copy2(source_agent, target_agent) + results.append(f"โœ… {subagent}: installed") + + return overall_success, "\n".join(results) + + except Exception as e: + return False, f"โŒ Failed to install subagents: {e}" + + def setup_mcp_tools(self) -> Tuple[bool, str]: + """Set up MCP tools for the subagent to use.""" + config_path = self.find_claude_config() + + if not config_path: + # Create new config + config_path = self.home / ".claude" / "config.json" + config_path.parent.mkdir(parents=True, exist_ok=True) + config = {"mcpServers": {}} + else: + # Load existing config + try: + with open(config_path) as f: + config = json.load(f) + except Exception as e: + return False, f"โŒ Failed to read config: {e}" + + # Ensure mcpServers section exists + if "mcpServers" not in config: + config["mcpServers"] = {} + + # Add codeflash MCP server for tools + config["mcpServers"]["codeflash"] = { + "command": f"{sys.executable}", + "args": ["-m", "codeflash.mcp.server"], + "env": {}, + "disabled": False, + } + + try: + # Create backup + if config_path.exists(): + backup_path = config_path.with_suffix(".json.backup") + shutil.copy2(config_path, backup_path) + + # Write updated config + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + return True, f"โœ… MCP tools configured: {config_path}" + + except Exception as e: + return False, f"โŒ Failed to configure MCP tools: {e}" + + def create_project_subagent(self, project_path: Path) -> Tuple[bool, str]: + """Create project-specific codeflash subagent.""" + try: + project_agents_dir = project_path / ".claude" / "agents" + project_agents_dir.mkdir(parents=True, exist_ok=True) + + # Copy the subagent with project-specific customizations + source_agent = Path(__file__).parent.parent / "agents" / "codeflash-optimizer.md" + target_agent = project_agents_dir / "codeflash-optimizer.md" + + if source_agent.exists(): + shutil.copy2(source_agent, target_agent) + return True, f"โœ… Project subagent created: {target_agent}" + return False, "โŒ Source subagent not found" + + except Exception as e: + return False, f"โŒ Failed to create project subagent: {e}" + + def get_integration_status(self) -> Dict[str, any]: + """Get current integration status.""" + config_path = self.find_claude_config() + agents_dir = self.find_agents_directory() + + # Check which subagents are installed + subagents = ["codeflash-optimizer", "codeflash-profiler", "codeflash-reviewer"] + installed_subagents = [] + for subagent in subagents: + if (agents_dir / f"{subagent}.md").exists(): + installed_subagents.append(subagent) + + mcp_configured = False + if config_path and config_path.exists(): + try: + with open(config_path) as f: + config = json.load(f) + mcp_configured = "mcpServers" in config and "codeflash" in config["mcpServers"] + except Exception: + pass + + return { + "claude_available": self.is_claude_code_available(), + "config_path": str(config_path) if config_path else None, + "agents_directory": str(agents_dir), + "installed_subagents": installed_subagents, + "all_subagents_installed": len(installed_subagents) == len(subagents), + "mcp_tools_configured": mcp_configured, + "system": self.system, + } + + def setup_complete_integration( + self, include_project: bool = False, project_path: Optional[Path] = None + ) -> Tuple[bool, str]: + """Set up complete Claude Code integration.""" + results = [] + overall_success = True + + # Step 1: Install subagents + success, message = self.install_subagents() + results.append("Subagent Installation:") + results.append(message) + if not success: + overall_success = False + + # Step 2: Set up MCP tools + success, message = self.setup_mcp_tools() + results.append(message) + if not success: + overall_success = False + + # Step 3: Create project subagent if requested + if include_project and project_path: + success, message = self.create_project_subagent(project_path) + results.append(message) + + return overall_success, "\n".join(results) + + def remove_integration(self) -> Tuple[bool, str]: + """Remove codeflash integration.""" + results = [] + + # Remove subagent + agents_dir = self.find_agents_directory() + subagent_path = agents_dir / "codeflash-optimizer.md" + + if subagent_path.exists(): + try: + subagent_path.unlink() + results.append("โœ… Removed codeflash optimizer subagent") + except Exception as e: + results.append(f"โŒ Failed to remove subagent: {e}") + + # Remove MCP configuration + config_path = self.find_claude_config() + if config_path and config_path.exists(): + try: + with open(config_path) as f: + config = json.load(f) + + if "mcpServers" in config and "codeflash" in config["mcpServers"]: + del config["mcpServers"]["codeflash"] + + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + results.append("โœ… Removed MCP tools configuration") + + except Exception as e: + results.append(f"โŒ Failed to remove MCP config: {e}") + + return True, "\n".join(results) + + +def setup_claude_integration( + include_project: bool = False, project_path: Optional[str] = None, force: bool = False +) -> bool: + """Main setup function for Claude Code integration.""" + integration = ClaudeCodeIntegration() + + if not force and not integration.is_claude_code_available(): + print("โš ๏ธ Claude Code not detected.") + print("Install Claude Code first: https://docs.anthropic.com/en/docs/claude-code") + print("Then run: codeflash integrate claude") + return False + + project_path_obj = Path(project_path) if project_path else None + success, message = integration.setup_complete_integration( + include_project=include_project, project_path=project_path_obj + ) + + print(message) + + if success: + print("\n๐ŸŽ‰ Codeflash is now integrated with Claude Code!") + print("\nHow to use:") + print("1. In Claude Code, type: '/agents' to see available subagents") + print("2. Invoke with: '@codeflash-optimizer optimize my bubble_sort function'") + print("3. Or let Claude auto-invoke when discussing performance") + print("\n๐Ÿ“ Restart Claude Code to ensure all changes take effect") + + return success diff --git a/codeflash/mcp/README.md b/codeflash/mcp/README.md new file mode 100644 index 000000000..c209004ea --- /dev/null +++ b/codeflash/mcp/README.md @@ -0,0 +1,116 @@ +# Codeflash MCP Server Integration + +This directory contains the Model Context Protocol (MCP) server implementation that makes Codeflash available as a Claude Code subagent. + +## Architecture Overview + +``` +codeflash/mcp/ +โ”œโ”€โ”€ __init__.py # Package initialization +โ”œโ”€โ”€ server.py # Main FastMCP server implementation +โ”œโ”€โ”€ integration.py # Claude Code integration utilities +โ””โ”€โ”€ README.md # This file +``` + +## Components + +### server.py +The main MCP server implementation using FastMCP. Provides: +- **12 MCP Tools** for all codeflash optimization workflows +- **Structured responses** using Pydantic models +- **Error handling** and timeout management +- **Resource providers** for project context + +### integration.py +Handles automatic setup and configuration: +- **Auto-detection** of Claude Code installation +- **Configuration management** for MCP servers +- **Cross-platform support** (macOS, Linux, Windows) +- **Status checking** and troubleshooting + +## Available MCP Tools + +| Tool | Description | Usage | +|------|-------------|--------| +| `optimize_function` | Optimize a specific function | `optimize_function("bubble_sort", "algorithms.py")` | +| `optimize_file` | Optimize all functions in a file | `optimize_file("algorithms.py")` | +| `trace_and_optimize` | Trace script and optimize | `trace_and_optimize("python main.py")` | +| `optimize_from_replay_tests` | Use tests for optimization | `optimize_from_replay_tests(["test_*.py"])` | +| `optimize_all_functions` | Optimize entire project | `optimize_all_functions()` โš ๏ธ | +| `initialize_project` | Set up codeflash | `initialize_project()` | +| `setup_github_actions` | Configure CI/CD | `setup_github_actions()` | +| `verify_installation` | Test setup | `verify_installation()` | +| `run_benchmarks` | Performance testing | `run_benchmarks()` | +| `get_codeflash_status` | Project status | `get_codeflash_status()` | +| `get_optimization_help` | Usage guide | `get_optimization_help()` | + +## Data Models + +### OptimizationResult +```python +{ + "success": bool, + "message": str, + "optimized_functions": List[str], + "performance_improvement": Optional[str], + "file_path": Optional[str], + "errors": List[str] +} +``` + +### ProjectStatus +```python +{ + "is_initialized": bool, + "config_file": Optional[str], + "module_root": Optional[str], + "tests_root": Optional[str], + "version": str +} +``` + +### TraceResult +```python +{ + "success": bool, + "trace_file": Optional[str], + "functions_traced": int, + "message": str +} +``` + +## Installation & Setup + +### 1. Install with MCP Support +```bash +pip install codeflash[mcp] +``` + +### 2. Auto-Configure Claude Code +```bash +codeflash setup claude-code +``` + +### 3. Manual Configuration (if needed) +Add to your Claude Code `config.json`: +```json +{ + "mcpServers": { + "codeflash": { + "command": "codeflash-mcp", + "args": [], + "env": {}, + "disabled": false + } + } +} +``` + +## Usage Examples + +### In Claude Code +``` +Human: Optimize the bubble_sort function in my algorithms.py file + +Claude: I'll use codeflash to optimize your bubble_sort function. +[Uses optimize_function tool] \ No newline at end of file diff --git a/codeflash/mcp/__init__.py b/codeflash/mcp/__init__.py new file mode 100644 index 000000000..9fef942fa --- /dev/null +++ b/codeflash/mcp/__init__.py @@ -0,0 +1,7 @@ +"""Codeflash Model Context Protocol (MCP) server integration. + +This module provides MCP server functionality to expose codeflash optimization +capabilities to Claude Code and other MCP-compatible clients. +""" + +__version__ = "1.0.0" diff --git a/codeflash/mcp/integration.py b/codeflash/mcp/integration.py new file mode 100644 index 000000000..2a41cef1e --- /dev/null +++ b/codeflash/mcp/integration.py @@ -0,0 +1,308 @@ +#!/usr/bin/env python3 +"""Claude Code integration utilities for automatic MCP server setup.""" + +import json +import os +import platform +import shutil +import sys +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from codeflash.cli_cmds.console import logger + + +class ClaudeCodeIntegration: + """Handles Claude Code MCP server integration.""" + + def __init__(self): + self.system = platform.system().lower() + self.home = Path.home() + self.claude_config_paths = self._get_claude_config_paths() + + def _get_claude_config_paths(self) -> List[Path]: + """Get possible Claude Code configuration file paths.""" + paths = [] + + if self.system == "darwin": # macOS + paths.extend( + [ + self.home / ".claude" / "config.json", + self.home / "Library" / "Application Support" / "Claude Code" / "config.json", + self.home / ".config" / "claude" / "config.json", + ] + ) + elif self.system == "linux": + paths.extend( + [ + self.home / ".claude" / "config.json", + self.home / ".config" / "claude" / "config.json", + Path("/etc/claude/config.json"), + ] + ) + elif self.system == "windows": + appdata = os.getenv("APPDATA", str(self.home / "AppData" / "Roaming")) + localappdata = os.getenv("LOCALAPPDATA", str(self.home / "AppData" / "Local")) + paths.extend( + [ + Path(appdata) / "Claude Code" / "config.json", + Path(localappdata) / "Claude Code" / "config.json", + self.home / ".claude" / "config.json", + ] + ) + + return paths + + def find_claude_config(self) -> Optional[Path]: + """Find the Claude Code configuration file.""" + for path in self.claude_config_paths: + if path.exists() and path.is_file(): + try: + with open(path) as f: + json.load(f) # Validate it's valid JSON + return path + except (OSError, json.JSONDecodeError): + continue + return None + + def is_claude_code_installed(self) -> bool: + """Check if Claude Code is installed and accessible.""" + try: + # Try to find claude command + claude_cmd = shutil.which("claude") + if claude_cmd: + return True + + # Check for config files + return self.find_claude_config() is not None + except Exception: + return False + + def get_codeflash_mcp_executable(self) -> str: + """Get the path to the codeflash MCP server executable.""" + # Check if we're in development mode + if hasattr(sys, "_called_from_test") or "pytest" in sys.modules: + return f"{sys.executable} -m codeflash.mcp.server" + + # Check if codeflash-mcp is in PATH (installed via pip) + mcp_cmd = shutil.which("codeflash-mcp") + if mcp_cmd: + return mcp_cmd + + # Fallback to python module execution + return f"{sys.executable} -m codeflash.mcp.server" + + def add_codeflash_to_claude_config(self, config_path: Path) -> bool: + """Add codeflash MCP server to Claude Code configuration.""" + try: + # Read existing config + with open(config_path) as f: + config = json.load(f) + + # Ensure mcpServers section exists + if "mcpServers" not in config: + config["mcpServers"] = {} + + # Add codeflash MCP server configuration + codeflash_config = { + "command": self.get_codeflash_mcp_executable(), + "args": [], + "env": {}, + "disabled": False, + } + + config["mcpServers"]["codeflash"] = codeflash_config + + # Create backup of original config + backup_path = config_path.with_suffix(".json.backup") + shutil.copy2(config_path, backup_path) + + # Write updated config + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + logger.info(f"โœ… Added codeflash MCP server to Claude Code config: {config_path}") + logger.info(f"๐Ÿ“ Backup created at: {backup_path}") + return True + + except Exception as e: + logger.error(f"โŒ Failed to update Claude Code config: {e}") + return False + + def create_claude_config(self) -> Optional[Path]: + """Create a new Claude Code configuration file.""" + # Try to create config in the most appropriate location + if self.system == "darwin": + config_path = self.home / ".claude" / "config.json" + elif self.system == "linux": + config_path = self.home / ".config" / "claude" / "config.json" + else: # windows + appdata = os.getenv("APPDATA", str(self.home / "AppData" / "Roaming")) + config_path = Path(appdata) / "Claude Code" / "config.json" + + try: + # Create directory if it doesn't exist + config_path.parent.mkdir(parents=True, exist_ok=True) + + # Create minimal config with codeflash MCP server + config = { + "mcpServers": { + "codeflash": { + "command": self.get_codeflash_mcp_executable(), + "args": [], + "env": {}, + "disabled": False, + } + } + } + + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + logger.info(f"โœ… Created Claude Code config with codeflash MCP server: {config_path}") + return config_path + + except Exception as e: + logger.error(f"โŒ Failed to create Claude Code config: {e}") + return None + + def setup_integration(self, force: bool = False) -> Tuple[bool, str]: + """Set up codeflash integration with Claude Code. + + Args: + force: Force setup even if already configured + + Returns: + Tuple of (success, message) + + """ + if not force and not self.is_claude_code_installed(): + return False, ( + "Claude Code not found. Please install Claude Code first:\n" + "Visit: https://docs.anthropic.com/en/docs/claude-code\n" + "Then run: codeflash setup claude-code" + ) + + config_path = self.find_claude_config() + + if config_path: + # Check if codeflash is already configured + try: + with open(config_path) as f: + config = json.load(f) + + if not force and "mcpServers" in config and "codeflash" in config["mcpServers"]: + return True, f"โœ… Codeflash MCP server already configured in: {config_path}" + + # Add codeflash to existing config + if self.add_codeflash_to_claude_config(config_path): + return True, f"โœ… Successfully added codeflash to Claude Code config: {config_path}" + return False, "โŒ Failed to update Claude Code configuration" + + except Exception as e: + return False, f"โŒ Error reading Claude Code config: {e}" + else: + # Create new config file + config_path = self.create_claude_config() + if config_path: + return True, f"โœ… Created new Claude Code config with codeflash: {config_path}" + return False, "โŒ Failed to create Claude Code configuration" + + def remove_integration(self) -> Tuple[bool, str]: + """Remove codeflash integration from Claude Code.""" + config_path = self.find_claude_config() + + if not config_path: + return True, "No Claude Code configuration found" + + try: + with open(config_path) as f: + config = json.load(f) + + if "mcpServers" not in config or "codeflash" not in config["mcpServers"]: + return True, "Codeflash not found in Claude Code configuration" + + # Remove codeflash MCP server + del config["mcpServers"]["codeflash"] + + # Create backup + backup_path = config_path.with_suffix(".json.backup") + shutil.copy2(config_path, backup_path) + + # Write updated config + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + return True, f"โœ… Removed codeflash from Claude Code config: {config_path}" + + except Exception as e: + return False, f"โŒ Failed to remove codeflash from Claude Code config: {e}" + + def get_status(self) -> Dict[str, any]: + """Get current integration status.""" + config_path = self.find_claude_config() + is_installed = self.is_claude_code_installed() + is_configured = False + + if config_path: + try: + with open(config_path) as f: + config = json.load(f) + is_configured = "mcpServers" in config and "codeflash" in config["mcpServers"] + except Exception: + pass + + return { + "claude_code_installed": is_installed, + "config_file": str(config_path) if config_path else None, + "codeflash_configured": is_configured, + "mcp_executable": self.get_codeflash_mcp_executable(), + "system": self.system, + } + + +def setup_claude_code_integration(force: bool = False) -> bool: + """Main function to set up Claude Code integration.""" + integration = ClaudeCodeIntegration() + success, message = integration.setup_integration(force=force) + + print(message) + + if success: + print("\n๐ŸŽ‰ Codeflash is now available as a Claude Code subagent!") + print("\nUsage in Claude Code:") + print("- 'Optimize the bubble_sort function using codeflash'") + print("- 'Use codeflash to trace and optimize my main.py script'") + print("- 'Initialize codeflash in this project'") + print("\nRestart Claude Code to load the new MCP server.") + + return success + + +def main(): + """Command-line interface for Claude Code integration.""" + import argparse + + parser = argparse.ArgumentParser(description="Set up codeflash integration with Claude Code") + parser.add_argument("--force", action="store_true", help="Force setup even if already configured") + parser.add_argument("--status", action="store_true", help="Show current integration status") + parser.add_argument("--remove", action="store_true", help="Remove codeflash integration") + + args = parser.parse_args() + + integration = ClaudeCodeIntegration() + + if args.status: + status = integration.get_status() + print(json.dumps(status, indent=2)) + elif args.remove: + success, message = integration.remove_integration() + print(message) + sys.exit(0 if success else 1) + else: + success = setup_claude_code_integration(force=args.force) + sys.exit(0 if success else 1) + + +if __name__ == "__main__": + main() diff --git a/codeflash/mcp/server.py b/codeflash/mcp/server.py new file mode 100644 index 000000000..9d061b3a2 --- /dev/null +++ b/codeflash/mcp/server.py @@ -0,0 +1,555 @@ +#!/usr/bin/env python3 +"""Codeflash MCP Server - Expose codeflash optimization capabilities to Claude Code. + +This server provides comprehensive access to all codeflash optimization features +through the Model Context Protocol (MCP), making it available as a Claude Code subagent. +""" + +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +try: + from fastmcp import FastMCP + from pydantic import BaseModel, Field +except ImportError as e: + print("Error: Missing MCP dependencies. Install with: pip install codeflash[mcp]") + print(f"Import error: {e}") + sys.exit(1) + +from codeflash.version import __version__ + +# Initialize FastMCP server +mcp = FastMCP("codeflash") + + +# Pydantic models for structured responses +class OptimizationResult(BaseModel): + """Result of a codeflash optimization operation.""" + + success: bool = Field(description="Whether the optimization was successful") + message: str = Field(description="Human-readable status message") + optimized_functions: List[str] = Field(default=[], description="List of functions that were optimized") + performance_improvement: Optional[str] = Field(default=None, description="Performance improvement summary") + file_path: Optional[str] = Field(default=None, description="Path to the optimized file") + errors: List[str] = Field(default=[], description="Any errors encountered during optimization") + + +class ProjectStatus(BaseModel): + """Status of codeflash setup in a project.""" + + is_initialized: bool = Field(description="Whether codeflash is initialized in the project") + config_file: Optional[str] = Field(default=None, description="Path to codeflash config file") + module_root: Optional[str] = Field(default=None, description="Project's Python module root") + tests_root: Optional[str] = Field(default=None, description="Project's tests root") + version: str = Field(description="Codeflash version") + + +class TraceResult(BaseModel): + """Result of a codeflash trace operation.""" + + success: bool = Field(description="Whether tracing was successful") + trace_file: Optional[str] = Field(default=None, description="Path to generated trace file") + functions_traced: int = Field(default=0, description="Number of functions traced") + message: str = Field(description="Status message") + + +def _run_codeflash_command(args: List[str], cwd: Optional[str] = None) -> Dict[str, Any]: + """Run a codeflash command and return structured results.""" + try: + cmd = [sys.executable, "-m", "codeflash"] + args + result = subprocess.run( + cmd, + cwd=cwd or os.getcwd(), + capture_output=True, + text=True, + timeout=300, + check=False, # 5 minute timeout + ) + + return { + "success": result.returncode == 0, + "stdout": result.stdout, + "stderr": result.stderr, + "returncode": result.returncode, + } + except subprocess.TimeoutExpired: + return {"success": False, "stdout": "", "stderr": "Command timed out after 5 minutes", "returncode": -1} + except Exception as e: + return {"success": False, "stdout": "", "stderr": str(e), "returncode": -1} + + +@mcp.tool() +def get_codeflash_status(project_path: str = ".") -> ProjectStatus: + """Get the current status of codeflash setup in a project. + + Args: + project_path: Path to the project directory (default: current directory) + + Returns: + ProjectStatus with current setup information + + """ + project_path = Path(project_path).resolve() + + # Check for pyproject.toml with codeflash config + pyproject_file = project_path / "pyproject.toml" + config_file = None + module_root = None + tests_root = None + is_initialized = False + + if pyproject_file.exists(): + try: + import tomlkit + + with open(pyproject_file) as f: + config = tomlkit.parse(f.read()) + + if "tool" in config and "codeflash" in config["tool"]: + is_initialized = True + config_file = str(pyproject_file) + codeflash_config = config["tool"]["codeflash"] + module_root = codeflash_config.get("module-root") + tests_root = codeflash_config.get("tests-root") + except Exception: + pass + + return ProjectStatus( + is_initialized=is_initialized, + config_file=config_file, + module_root=module_root, + tests_root=tests_root, + version=__version__, + ) + + +@mcp.tool() +def optimize_function( + function_name: str, file_path: Optional[str] = None, project_path: str = ".", create_pr: bool = True +) -> OptimizationResult: + """Optimize a specific Python function using Codeflash AI optimization. + + Args: + function_name: Name of the function to optimize + file_path: Path to the file containing the function (optional if auto-discoverable) + project_path: Path to the project root directory + create_pr: Whether to create a pull request with the optimization + + Returns: + OptimizationResult with optimization details and performance improvements + + """ + args = ["--function", function_name] + + if file_path: + args.extend(["--file", file_path]) + + if not create_pr: + args.append("--no-pr") + + result = _run_codeflash_command(args, cwd=project_path) + + # Parse the output to extract optimization details + optimized_functions = [] + performance_improvement = None + errors = [] + + if result["success"]: + # Extract optimization information from stdout + stdout = result["stdout"] + if "Optimized function" in stdout: + optimized_functions = [function_name] + if "performance improvement" in stdout.lower() or "speedup" in stdout.lower(): + lines = stdout.split("\n") + for line in lines: + if "speedup" in line.lower() or "faster" in line.lower(): + performance_improvement = line.strip() + break + elif result["stderr"]: + errors.append(result["stderr"]) + + return OptimizationResult( + success=result["success"], + message=result["stdout"] if result["success"] else result["stderr"], + optimized_functions=optimized_functions, + performance_improvement=performance_improvement, + file_path=file_path, + errors=errors, + ) + + +@mcp.tool() +def optimize_file(file_path: str, project_path: str = ".", create_pr: bool = True) -> OptimizationResult: + """Optimize all functions in a specific Python file. + + Args: + file_path: Path to the Python file to optimize + project_path: Path to the project root directory + create_pr: Whether to create a pull request with optimizations + + Returns: + OptimizationResult with details of all optimizations made + + """ + args = ["--file", file_path] + + if not create_pr: + args.append("--no-pr") + + result = _run_codeflash_command(args, cwd=project_path) + + optimized_functions = [] + errors = [] + + if result["success"]: + # Extract function names from output + stdout = result["stdout"] + lines = stdout.split("\n") + for line in lines: + if "Optimized function" in line: + # Extract function name from the line + parts = line.split("Optimized function") + if len(parts) > 1: + func_name = parts[1].strip().split()[0] + optimized_functions.append(func_name) + elif result["stderr"]: + errors.append(result["stderr"]) + + return OptimizationResult( + success=result["success"], + message=result["stdout"] if result["success"] else result["stderr"], + optimized_functions=optimized_functions, + file_path=file_path, + errors=errors, + ) + + +@mcp.tool() +def trace_and_optimize( + script_command: str, + project_path: str = ".", + output_file: str = "codeflash.trace", + max_functions: int = 100, + timeout: Optional[int] = None, +) -> TraceResult: + """Trace a Python script execution and optimize based on the trace. + + Args: + script_command: The Python command to trace (e.g., "python main.py") + project_path: Path to the project root directory + output_file: File to save the trace to + max_functions: Maximum number of function calls to trace + timeout: Maximum time in seconds to run the trace + + Returns: + TraceResult with tracing and optimization details + + """ + # Parse the script command to extract the Python file + cmd_parts = script_command.split() + if len(cmd_parts) < 2: + return TraceResult(success=False, message="Invalid script command. Expected format: 'python script.py [args]'") + + script_args = cmd_parts[1:] # Everything after 'python' + + args = ["optimize", "--output", output_file, "--max-function-count", str(max_functions)] + + if timeout: + args.extend(["--timeout", str(timeout)]) + + # Add the script arguments + args.extend(script_args) + + result = _run_codeflash_command(args, cwd=project_path) + + functions_traced = 0 + trace_file = None + + if result["success"]: + trace_file = str(Path(project_path) / output_file) + # Try to count functions from trace file + try: + if Path(trace_file).exists(): + with open(trace_file) as f: + content = f.read() + functions_traced = content.count('"function":') + except Exception: + pass + + return TraceResult( + success=result["success"], + trace_file=trace_file if result["success"] else None, + functions_traced=functions_traced, + message=result["stdout"] if result["success"] else result["stderr"], + ) + + +@mcp.tool() +def optimize_from_replay_tests( + test_files: List[str], project_path: str = ".", create_pr: bool = True +) -> OptimizationResult: + """Optimize functions using existing replay test files. + + Args: + test_files: List of paths to replay test files + project_path: Path to the project root directory + create_pr: Whether to create a pull request with optimizations + + Returns: + OptimizationResult with optimization details + + """ + args = ["--replay-test"] + test_files + + if not create_pr: + args.append("--no-pr") + + result = _run_codeflash_command(args, cwd=project_path) + + optimized_functions = [] + errors = [] + + if result["success"]: + stdout = result["stdout"] + lines = stdout.split("\n") + for line in lines: + if "Optimized function" in line: + parts = line.split("Optimized function") + if len(parts) > 1: + func_name = parts[1].strip().split()[0] + optimized_functions.append(func_name) + elif result["stderr"]: + errors.append(result["stderr"]) + + return OptimizationResult( + success=result["success"], + message=result["stdout"] if result["success"] else result["stderr"], + optimized_functions=optimized_functions, + errors=errors, + ) + + +@mcp.tool() +def optimize_all_functions( + start_directory: str = ".", project_path: str = ".", create_pr: bool = True +) -> OptimizationResult: + """Optimize all functions in a project or directory. + + WARNING: This can take a very long time for large projects. + + Args: + start_directory: Directory to start optimization from + project_path: Path to the project root directory + create_pr: Whether to create a pull request with optimizations + + Returns: + OptimizationResult with details of all optimizations + + """ + args = ["--all"] + + if start_directory and start_directory != ".": + args.append(start_directory) + + if not create_pr: + args.append("--no-pr") + + result = _run_codeflash_command(args, cwd=project_path) + + optimized_functions = [] + errors = [] + + if result["success"]: + stdout = result["stdout"] + lines = stdout.split("\n") + for line in lines: + if "Optimized function" in line: + parts = line.split("Optimized function") + if len(parts) > 1: + func_name = parts[1].strip().split()[0] + optimized_functions.append(func_name) + elif result["stderr"]: + errors.append(result["stderr"]) + + return OptimizationResult( + success=result["success"], + message=result["stdout"] if result["success"] else result["stderr"], + optimized_functions=optimized_functions, + errors=errors, + ) + + +@mcp.tool() +def initialize_project( + project_path: str = ".", + module_root: Optional[str] = None, + tests_root: Optional[str] = None, + test_framework: str = "pytest", +) -> Dict[str, Any]: + """Initialize codeflash in a Python project. + + Args: + project_path: Path to the project directory + module_root: Path to the Python module root (auto-detected if not provided) + tests_root: Path to the tests directory (auto-detected if not provided) + test_framework: Testing framework to use ('pytest' or 'unittest') + + Returns: + Dictionary with initialization results + + """ + args = ["init"] + + # Change to project directory and run init + result = _run_codeflash_command(args, cwd=project_path) + + return { + "success": result["success"], + "message": result["stdout"] if result["success"] else result["stderr"], + "config_created": result["success"], + } + + +@mcp.tool() +def setup_github_actions(project_path: str = ".") -> Dict[str, Any]: + """Set up GitHub Actions workflow for automatic codeflash optimization. + + Args: + project_path: Path to the project directory + + Returns: + Dictionary with setup results + + """ + args = ["init-actions"] + + result = _run_codeflash_command(args, cwd=project_path) + + return { + "success": result["success"], + "message": result["stdout"] if result["success"] else result["stderr"], + "workflow_created": result["success"], + } + + +@mcp.tool() +def verify_installation(project_path: str = ".") -> Dict[str, Any]: + """Verify that codeflash is working correctly by running a test optimization. + + Args: + project_path: Path to the project directory + + Returns: + Dictionary with verification results + + """ + args = ["--verify-setup"] + + result = _run_codeflash_command(args, cwd=project_path) + + return { + "success": result["success"], + "message": result["stdout"] if result["success"] else result["stderr"], + "verification_passed": result["success"], + } + + +@mcp.tool() +def run_benchmarks(project_path: str = ".", benchmarks_root: Optional[str] = None) -> Dict[str, Any]: + """Run benchmark tests and calculate optimization impact. + + Args: + project_path: Path to the project directory + benchmarks_root: Path to benchmarks directory (auto-detected if not provided) + + Returns: + Dictionary with benchmark results + + """ + args = ["--benchmark"] + + if benchmarks_root: + args.extend(["--benchmarks-root", benchmarks_root]) + + result = _run_codeflash_command(args, cwd=project_path) + + return { + "success": result["success"], + "message": result["stdout"] if result["success"] else result["stderr"], + "benchmarks_completed": result["success"], + } + + +@mcp.tool() +def get_optimization_help() -> Dict[str, Any]: + """Get comprehensive help and usage information for codeflash optimization. + + Returns: + Dictionary with help information and usage examples + + """ + return { + "version": __version__, + "description": "Codeflash is an AI-powered Python performance optimizer that automatically speeds up your code while verifying correctness.", + "common_workflows": { + "optimize_single_function": { + "description": "Optimize a specific function by name", + "example": "optimize_function('bubble_sort', file_path='algorithms.py')", + }, + "optimize_entire_file": { + "description": "Optimize all functions in a Python file", + "example": "optimize_file('algorithms.py')", + }, + "trace_and_optimize_script": { + "description": "Trace a script execution and optimize based on usage patterns", + "example": "trace_and_optimize('python main.py --data test_data.csv')", + }, + "optimize_from_tests": { + "description": "Use existing test files to guide optimization", + "example": "optimize_from_replay_tests(['tests/test_algorithms.py'])", + }, + }, + "best_practices": [ + "Always review optimizations before merging to ensure correctness", + "Use trace_and_optimize for end-to-end script optimization", + "Start with single functions before optimizing entire files", + "Set up GitHub Actions for continuous optimization", + "Run benchmarks to measure performance improvements", + ], + "supported_optimizations": [ + "Algorithm improvements (better time/space complexity)", + "Library method optimization (using more efficient methods)", + "Loop optimization and vectorization", + "Caching and memoization", + "Data structure optimization", + "Removing wasteful computations", + ], + } + + +# Resource for providing project context +@mcp.resource("codeflash://project-config") +def get_project_config() -> str: + """Get codeflash configuration from the current project.""" + try: + config_path = Path.cwd() / "pyproject.toml" + if config_path.exists(): + return config_path.read_text() + return "No codeflash configuration found in current project." + except Exception as e: + return f"Error reading project configuration: {e}" + + +def main(): + """Main entry point for the MCP server.""" + import uvloop + + uvloop.install() + mcp.run() + + +if __name__ == "__main__": + main() diff --git a/docs/claude-code-integration.mdx b/docs/claude-code-integration.mdx new file mode 100644 index 000000000..6497479d6 --- /dev/null +++ b/docs/claude-code-integration.mdx @@ -0,0 +1,252 @@ +--- +title: "Claude Code Integration" +description: "Use Codeflash as a Claude Code subagent for seamless AI-powered code optimization" +--- + +# Claude Code Integration + +Codeflash integrates natively with Claude Code through **specialized subagents** - dedicated AI assistants that provide comprehensive Python performance optimization capabilities directly within your development workflow. + +## ๐Ÿš€ Quick Setup + +### 1. Install Codeflash + +```bash +# Install codeflash (includes Claude Code integration) +pip install codeflash +``` + +### 2. Integrate with Claude Code + +```bash +# Automatically set up native subagents +codeflash integrate claude +``` + +### 3. Restart Claude Code + +After setup, restart Claude Code to load the new MCP server. + +## โœจ Three Specialized Subagents + +Codeflash provides **three dedicated subagents**, each with specialized expertise: + +### ๐Ÿš€ **@codeflash-optimizer** - AI-Powered Optimization +The primary optimization specialist that implements performance improvements: + +``` +"@codeflash-optimizer optimize the bubble_sort function" +"@codeflash-optimizer trace my script and optimize bottlenecks" +"@codeflash-optimizer set up GitHub Actions for continuous optimization" +``` + +**Capabilities:** +- Function and file-level optimization (2-55x speedups) +- End-to-end script tracing and optimization +- Algorithm improvements and efficiency gains +- Project setup and CI/CD integration + +### ๐Ÿ” **@codeflash-profiler** - Performance Analysis +Expert in identifying bottlenecks and analyzing performance patterns: + +``` +"@codeflash-profiler analyze why my script is slow" +"@codeflash-profiler profile this function and find bottlenecks" +"@codeflash-profiler benchmark my code and show performance metrics" +``` + +**Capabilities:** +- Execution tracing and hotspot identification +- Performance benchmarking and statistical analysis +- Memory usage analysis and optimization recommendations +- Scalability assessment and bottleneck prioritization + +### ๐Ÿ‘๏ธ **@codeflash-reviewer** - Performance Code Reviews +Specialized in reviewing code for performance issues during development: + +``` +"@codeflash-reviewer review this code for performance issues" +"@codeflash-reviewer check my pull request for optimization opportunities" +"@codeflash-reviewer analyze the algorithmic complexity of this function" +``` + +**Capabilities:** +- Algorithmic efficiency analysis (time/space complexity) +- Performance anti-pattern detection +- Library usage optimization recommendations +- Code review integration with performance focus + +## ๐Ÿ”ง Available MCP Tools + +The integration provides these tools that Claude Code can use: + +| Tool | Description | +|------|-------------| +| `optimize_function` | Optimize a specific function by name | +| `optimize_file` | Optimize all functions in a file | +| `trace_and_optimize` | Trace script execution and optimize | +| `optimize_from_replay_tests` | Use test files to guide optimization | +| `optimize_all_functions` | Optimize entire project (use with caution) | +| `initialize_project` | Set up codeflash in a project | +| `setup_github_actions` | Configure CI/CD optimization | +| `verify_installation` | Test that codeflash works correctly | +| `run_benchmarks` | Execute performance benchmarks | +| `get_codeflash_status` | Check current project setup | +| `get_optimization_help` | Get usage help and best practices | + +## ๐Ÿ“‹ Usage Examples + +### Basic Optimization +``` +Human: I have a slow sorting function in my project. Can you help optimize it? + +Claude: I'll help you optimize your sorting function using codeflash. Let me first check what functions are available in your project and then optimize the sorting function. + +[Claude uses codeflash tools to find and optimize the function] +``` + +### End-to-End Script Optimization +``` +Human: My data processing script takes too long. Can you trace it and find optimizations? + +Claude: I'll use codeflash to trace your script execution and identify optimization opportunities based on real usage patterns. + +[Claude uses trace_and_optimize tool] +``` + +### Project Setup +``` +Human: I want to set up codeflash in my new Python project + +Claude: I'll help you initialize codeflash in your project and set up the optimal configuration. + +[Claude uses initialize_project and setup tools] +``` + +## โš™๏ธ Configuration + +### Manual Setup + +If automatic setup doesn't work, you can manually add codeflash to your Claude Code configuration: + +1. **Find your Claude Code config file:** + - macOS: `~/.claude/config.json` or `~/Library/Application Support/Claude Code/config.json` + - Linux: `~/.config/claude/config.json` + - Windows: `%APPDATA%/Claude Code/config.json` + +2. **Add the MCP server configuration:** + +```json +{ + "mcpServers": { + "codeflash": { + "command": "codeflash-mcp", + "args": [], + "env": {}, + "disabled": false + } + } +} +``` + +### Verification + +Check if the integration is working: + +```bash +# Check integration status +codeflash setup status + +# Test in Claude Code +# Ask: "Show me codeflash help" or "What codeflash tools are available?" +``` + +## ๐Ÿ” Troubleshooting + +### Common Issues + +#### "codeflash-mcp command not found" +```bash +# Reinstall with MCP dependencies +pip uninstall codeflash +pip install codeflash[mcp] +``` + +#### "MCP server failed to start" +```bash +# Check if dependencies are installed +python -c "import mcp, fastmcp; print('MCP dependencies OK')" + +# Test the MCP server directly +codeflash-mcp +``` + +#### "Claude Code doesn't recognize codeflash" +1. Ensure Claude Code is restarted after setup +2. Check config file exists and is valid JSON +3. Verify MCP server path is correct + +### Debug Commands + +```bash +# Show detailed integration status +codeflash setup status + +# Remove and reconfigure +codeflash setup claude-code --remove +codeflash setup claude-code --force + +# Test MCP server manually +python -m codeflash.mcp.server +``` + +## ๐ŸŽฏ Best Practices + +### Optimization Workflow +1. **Start Small**: Begin with single function optimization +2. **Use Tracing**: For end-to-end scripts, use trace-based optimization +3. **Review Changes**: Always review optimizations before merging +4. **Benchmark**: Use benchmarking to measure improvements +5. **Iterative**: Optimize iteratively rather than all at once + +### Claude Code Tips +- Be specific about which functions/files to optimize +- Ask for explanations of optimizations made +- Request benchmarking to validate improvements +- Use natural language - Claude Code understands context + +### Performance Considerations +- Large project optimization can take significant time +- Use `--no-pr` flag for local-only testing +- Set up GitHub Actions for continuous optimization + +## ๐Ÿš€ Advanced Usage + +### Custom Workflows +``` +"Create a codeflash optimization workflow for my data pipeline project that: +1. Traces the main processing script +2. Optimizes the top 5 slowest functions +3. Runs benchmarks to verify improvements +4. Creates a PR with the optimizations" +``` + +### Integration with CI/CD +``` +"Help me set up automated codeflash optimization in my GitHub Actions workflow" +``` + +### Performance Analysis +``` +"Analyze my Python project's performance bottlenecks and suggest which functions to optimize first" +``` + +## ๐Ÿ“š Learn More + +- [Codeflash Documentation](/) +- [Model Context Protocol](https://modelcontextprotocol.io/) +- [Claude Code Documentation](https://docs.anthropic.com/en/docs/claude-code) + +--- + +*Codeflash Claude Code integration brings AI-powered code optimization directly into your development workflow, making performance improvements as easy as asking for them.* \ No newline at end of file diff --git a/docs/docs.json b/docs/docs.json index b3ccf9027..84786547f 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -27,7 +27,8 @@ { "group": "๐Ÿš€ Getting Started", "pages": [ - "getting-started/local-installation" + "getting-started/local-installation", + "claude-code-integration" ] }, { @@ -125,11 +126,11 @@ "href": "https://www.codeflash.ai/privacy-policy" }, { - "label": "Terms of Service", + "label": "Terms of Service", "href": "https://www.codeflash.ai/terms-of-service" } ] } ] } -} \ No newline at end of file +} diff --git a/pyproject.toml b/pyproject.toml index 3333a3b68..9b820f5a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,11 +45,20 @@ dependencies = [ "codeflash-benchmark", ] +[project.optional-dependencies] +mcp = [ + "mcp>=1.2.0", + "fastmcp>=2.0.0", + "uvloop>=0.20.0", + "click>=8.1.0", +] + [project.urls] Homepage = "https://codeflash.ai" [project.scripts] codeflash = "codeflash.main:main" +codeflash-mcp = "codeflash.mcp.server:main" [dependency-groups] dev = [ diff --git a/scripts/setup_mcp.py b/scripts/setup_mcp.py new file mode 100644 index 000000000..e98d1003d --- /dev/null +++ b/scripts/setup_mcp.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +"""Complete setup and verification script for Codeflash MCP integration.""" + +import subprocess +import sys + + +def check_dependencies(): + """Check if all required dependencies are installed.""" + print("๐Ÿ” Checking dependencies...") + + missing = [] + + # Check core codeflash + try: + import codeflash + + print(f" โœ… codeflash {codeflash.version.__version__}") + except ImportError: + missing.append("codeflash") + print(" โŒ codeflash not found") + + # Check MCP dependencies + try: + import mcp + + print(" โœ… mcp") + except ImportError: + missing.append("mcp") + print(" โŒ mcp not found") + + try: + import fastmcp + + print(" โœ… fastmcp") + except ImportError: + missing.append("fastmcp") + print(" โŒ fastmcp not found") + + try: + import uvloop + + print(" โœ… uvloop") + except ImportError: + missing.append("uvloop") + print(" โŒ uvloop not found") + + if missing: + print(f"\n๐Ÿ“ฆ Missing dependencies: {', '.join(missing)}") + print("Install with: pip install codeflash[mcp]") + return False + + return True + + +def test_mcp_server(): + """Test that the MCP server can start.""" + print("\n๐Ÿงช Testing MCP server...") + + try: + result = subprocess.run( + [sys.executable, "-c", "from codeflash.mcp.server import main; print('MCP server import OK')"], + capture_output=True, + text=True, + timeout=10, + check=False, + ) + + if result.returncode == 0: + print(" โœ… MCP server imports successfully") + return True + print(f" โŒ MCP server import failed: {result.stderr}") + return False + except subprocess.TimeoutExpired: + print(" โŒ MCP server test timed out") + return False + except Exception as e: + print(f" โŒ MCP server test error: {e}") + return False + + +def setup_claude_code(): + """Set up Claude Code integration.""" + print("\n๐Ÿ”ง Setting up Claude Code integration...") + + try: + from codeflash.mcp.integration import ClaudeCodeIntegration + + integration = ClaudeCodeIntegration() + success, message = integration.setup_integration() + + if success: + print(f" โœ… {message}") + return True + print(f" โš ๏ธ {message}") + return False + + except Exception as e: + print(f" โŒ Setup failed: {e}") + return False + + +def verify_integration(): + """Verify the complete integration.""" + print("\nโœ… Verifying integration...") + + try: + from codeflash.mcp.integration import ClaudeCodeIntegration + + integration = ClaudeCodeIntegration() + status = integration.get_status() + + print(f" Claude Code installed: {'โœ…' if status['claude_code_installed'] else 'โŒ'}") + print(f" Config file: {status['config_file'] or 'Not found'}") + print(f" Codeflash configured: {'โœ…' if status['codeflash_configured'] else 'โŒ'}") + print(f" MCP executable: {status['mcp_executable']}") + + if status["claude_code_installed"] and status["codeflash_configured"]: + print("\n๐ŸŽ‰ Integration is complete!") + print("\nNext steps:") + print("1. Restart Claude Code") + print("2. Try: 'Show me codeflash optimization help'") + print("3. Or: 'Optimize a function in my project using codeflash'") + return True + print("\nโš ๏ธ Integration incomplete") + return False + + except Exception as e: + print(f" โŒ Verification failed: {e}") + return False + + +def main(): + """Main setup process.""" + print("๐Ÿš€ Codeflash MCP Integration Setup") + print("=" * 40) + + success = True + + # Step 1: Check dependencies + if not check_dependencies(): + success = False + print("\nโŒ Please install missing dependencies first:") + print(" pip install codeflash[mcp]") + + # Step 2: Test MCP server + if success and not test_mcp_server(): + success = False + + # Step 3: Set up Claude Code + if success: + setup_claude_code() # This can succeed partially + + # Step 4: Verify integration + verify_integration() + + print("\n" + "=" * 40) + if success: + print("โœ… Setup completed successfully!") + else: + print("โš ๏ธ Setup completed with issues") + print("\nFor help:") + print("- Run: codeflash setup status") + print("- Check: https://docs.codeflash.ai/claude-code-integration") + + return success + + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) diff --git a/tests/test_mcp_integration.py b/tests/test_mcp_integration.py new file mode 100644 index 000000000..9bb2d8d14 --- /dev/null +++ b/tests/test_mcp_integration.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +"""Test suite for Codeflash MCP integration.""" + +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +# Add the codeflash module to the path for testing +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from codeflash.mcp.integration import ClaudeCodeIntegration + +class TestClaudeCodeIntegration(unittest.TestCase): + """Test Claude Code integration functionality.""" + + def setUp(self): + """Set up test environment.""" + self.integration = ClaudeCodeIntegration() + self.temp_dir = tempfile.mkdtemp() + self.temp_path = Path(self.temp_dir) + + def tearDown(self): + """Clean up test environment.""" + import shutil + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_get_claude_config_paths(self): + """Test that config paths are generated correctly.""" + paths = self.integration._get_claude_config_paths() + self.assertIsInstance(paths, list) + self.assertTrue(len(paths) > 0) + + # All paths should be Path objects + for path in paths: + self.assertIsInstance(path, Path) + + def test_get_codeflash_mcp_executable(self): + """Test MCP executable path detection.""" + executable = self.integration.get_codeflash_mcp_executable() + self.assertIsInstance(executable, str) + self.assertTrue(len(executable) > 0) + + def test_create_claude_config(self): + """Test creating a new Claude Code configuration.""" + # Mock the config path to use our temp directory + config_path = self.temp_path / "config.json" + + with patch.object(self.integration, '_get_claude_config_paths', return_value=[config_path]): + result = self.integration.create_claude_config() + + if result: # Only test if creation was successful + self.assertTrue(config_path.exists()) + + # Verify the config content + with open(config_path, 'r') as f: + config = json.load(f) + + self.assertIn("mcpServers", config) + self.assertIn("codeflash", config["mcpServers"]) + self.assertIn("command", config["mcpServers"]["codeflash"]) + + def test_add_codeflash_to_existing_config(self): + """Test adding codeflash to an existing Claude Code configuration.""" + config_path = self.temp_path / "config.json" + + # Create an existing config + existing_config = { + "someOtherSetting": "value", + "mcpServers": { + "existingServer": { + "command": "some-other-server" + } + } + } + + with open(config_path, 'w') as f: + json.dump(existing_config, f) + + # Add codeflash to the config + success = self.integration.add_codeflash_to_claude_config(config_path) + + if success: + # Verify the updated config + with open(config_path, 'r') as f: + config = json.load(f) + + self.assertIn("mcpServers", config) + self.assertIn("codeflash", config["mcpServers"]) + self.assertIn("existingServer", config["mcpServers"]) # Existing server should remain + self.assertEqual(config["someOtherSetting"], "value") # Other settings should remain + + def test_get_status(self): + """Test getting integration status.""" + status = self.integration.get_status() + + self.assertIsInstance(status, dict) + self.assertIn("claude_code_installed", status) + self.assertIn("config_file", status) + self.assertIn("codeflash_configured", status) + self.assertIn("mcp_executable", status) + self.assertIn("system", status) + + # Values should be of correct types + self.assertIsInstance(status["claude_code_installed"], bool) + self.assertIsInstance(status["codeflash_configured"], bool) + self.assertIsInstance(status["mcp_executable"], str) + self.assertIsInstance(status["system"], str) + + +class TestMCPServer(unittest.TestCase): + """Test MCP server functionality.""" + + @classmethod + def setUpClass(cls): + """Set up test environment.""" + # Check if MCP dependencies are available + try: + import mcp + import fastmcp + cls.mcp_available = True + except ImportError: + cls.mcp_available = False + + def setUp(self): + """Set up each test.""" + if not self.mcp_available: + self.skipTest("MCP dependencies not available") + + def test_mcp_server_import(self): + """Test that the MCP server can be imported.""" + try: + from codeflash.mcp.server import mcp + self.assertIsNotNone(mcp) + except ImportError as e: + self.fail(f"Failed to import MCP server: {e}") + + def test_mcp_tools_registration(self): + """Test that MCP tools are properly registered.""" + from codeflash.mcp.server import mcp + + # Check that key tools are registered + expected_tools = [ + "optimize_function", + "optimize_file", + "trace_and_optimize", + "get_codeflash_status", + "get_optimization_help" + ] + + # This is a basic test - in a real environment we'd need to introspect + # the FastMCP instance to check registered tools + self.assertTrue(hasattr(mcp, 'tool')) + + def test_project_status_model(self): + """Test ProjectStatus Pydantic model.""" + from codeflash.mcp.server import ProjectStatus + + status = ProjectStatus( + is_initialized=True, + version="1.0.0" + ) + + self.assertTrue(status.is_initialized) + self.assertEqual(status.version, "1.0.0") + self.assertIsNone(status.config_file) + + def test_optimization_result_model(self): + """Test OptimizationResult Pydantic model.""" + from codeflash.mcp.server import OptimizationResult + + result = OptimizationResult( + success=True, + message="Optimization completed", + optimized_functions=["bubble_sort", "quick_sort"], + performance_improvement="2x faster" + ) + + self.assertTrue(result.success) + self.assertEqual(len(result.optimized_functions), 2) + self.assertEqual(result.performance_improvement, "2x faster") + self.assertEqual(len(result.errors), 0) # Default empty list + + +class TestCLIIntegration(unittest.TestCase): + """Test CLI integration with MCP setup.""" + + def test_setup_command_help(self): + """Test that setup command provides proper help.""" + try: + result = subprocess.run( + [sys.executable, "-m", "codeflash", "setup", "--help"], + capture_output=True, + text=True, + timeout=10 + ) + + # Should not crash and should mention claude-code + self.assertIn("claude-code", result.stdout.lower()) + except (subprocess.TimeoutExpired, FileNotFoundError): + self.skipTest("Codeflash CLI not available or too slow") + + def test_mcp_executable_exists(self): + """Test that codeflash-mcp executable can be found.""" + try: + result = subprocess.run( + [sys.executable, "-m", "codeflash.mcp.server", "--help"], + capture_output=True, + text=True, + timeout=10 + ) + + # Should run without critical errors + # (might fail due to missing dependencies, but shouldn't crash Python) + self.assertNotEqual(result.returncode, -11) # Avoid segfaults + except (subprocess.TimeoutExpired, FileNotFoundError): + self.skipTest("MCP server not available") + + +class TestDocumentation(unittest.TestCase): + """Test that documentation is properly created.""" + + def test_claude_code_integration_doc_exists(self): + """Test that Claude Code integration documentation exists.""" + doc_path = Path(__file__).parent.parent / "docs" / "claude-code-integration.mdx" + + if doc_path.exists(): + content = doc_path.read_text() + self.assertIn("Claude Code Integration", content) + self.assertIn("MCP", content) + self.assertIn("codeflash setup claude-code", content) + + def test_docs_navigation_updated(self): + """Test that docs.json includes the new page.""" + docs_config_path = Path(__file__).parent.parent / "docs" / "docs.json" + + if docs_config_path.exists(): + with open(docs_config_path, 'r') as f: + config = json.load(f) + + # Find the getting started section and check for claude-code-integration + found = False + if "navigation" in config and "tabs" in config["navigation"]: + for tab in config["navigation"]["tabs"]: + if "groups" in tab: + for group in tab["groups"]: + if "pages" in group: + if "claude-code-integration" in group["pages"]: + found = True + break + + self.assertTrue(found, "claude-code-integration not found in docs navigation") + + +def run_integration_tests(): + """Run all integration tests.""" + print("๐Ÿงช Running Codeflash MCP Integration Tests\n") + + # Run the test suite + loader = unittest.TestLoader() + suite = loader.loadTestsFromModule(sys.modules[__name__]) + + runner = unittest.TextTestRunner(verbosity=2) + result = runner.run(suite) + + if result.wasSuccessful(): + print("\nโœ… All tests passed! MCP integration is working correctly.") + return True + else: + print(f"\nโŒ {len(result.failures)} test(s) failed, {len(result.errors)} error(s)") + return False + + +if __name__ == "__main__": + success = run_integration_tests() + sys.exit(0 if success else 1) \ No newline at end of file