orchestrator.py

import ollama
import re
from rich.console import Console
from rich.panel import Panel

console = Console()

def extract_python_code(text: str) -> str:
    """Extracts python code from markdown code blocks."""
    match = re.search(r"```python\s*(.*?)\s*```", text, re.DOTALL)
    if match:
        return match.group(1).strip()
    
    # Fallback to plain code block
    match = re.search(r"```\s*(.*?)\s*```", text, re.DOTALL)
    if match:
        return match.group(1).strip()
    
    return text.strip()

def main():
    console.print(Panel.fit("[bold cyan]Local-Claude-Brain Orchestrator Starting...[/bold cyan]"))
    
    # Read the messy code
    try:
        with open("messy_calculator.py", "r", encoding="utf-8") as f:
            messy_code = f.read()
    except FileNotFoundError:
        console.print("[bold red]Error: messy_calculator.py not found![/bold red]")
        return

    MODEL_NAME = "qwen3.5:4b"

    # =========================================================
    # Step 1: The Planner Agent
    # =========================================================
    console.print("\n[bold magenta]Step 1: Planner Agent analyzing code...[/bold magenta]")
    
    planner_prompt = (
        "You are an expert software architect. Analyze the provided messy Python code. "
        "Create a strict, bulleted, step-by-step refactoring plan to improve it. "
        "You must address missing functions, type hints, bad variable names, and repeated code/error handling. "
        "DO NOT write any actual Python code in your response, ONLY the bulleted plan."
    )
    
    planner_response = ollama.chat(model=MODEL_NAME, messages=[
        {'role': 'system', 'content': planner_prompt},
        {'role': 'user', 'content': messy_code}
    ], think=False)
    
    plan = planner_response['message']['content']
    # Print plan to console in yellow
    console.print(Panel(plan, title="Planner: Refactoring Plan", border_style="yellow", style="yellow"))

    # =========================================================
    # Step 2: The Executor Agent
    # =========================================================
    console.print("\n[bold magenta]Step 2: Executor Agent writing clean code...[/bold magenta]")
    
    executor_prompt = (
        "You are an elite Python developer. You will receive messy Python code and a refactoring plan. "
        "Rewrite the code perfectly based on the structural plan provided. "
        "Output ONLY the final refactored python code inside a ```python ``` code block. "
        "Do not include any other text, explanations, or conversation."
    )
    
    executor_input = (
        f"--- MESSY CODE ---\n{messy_code}\n\n"
        f"--- REFACTORING PLAN ---\n{plan}"
    )
    
    executor_response = ollama.chat(model=MODEL_NAME, messages=[
        {'role': 'system', 'content': executor_prompt},
        {'role': 'user', 'content': executor_input}
    ], think=False)
    
    raw_executed_code = executor_response['message']['content']
    clean_code = extract_python_code(raw_executed_code)

    # =========================================================
    # Step 3: The Output
    # =========================================================
    with open("clean_calculator.py", "w", encoding="utf-8") as f:
        f.write(clean_code)
        
    console.print(Panel.fit(
        "[bold green]Success! Refactored code executed and saved to `clean_calculator.py`[/bold green]",
        border_style="green"
    ))

if __name__ == "__main__":
    main()