agents/multiagent_workflow.py

#!/usr/bin/env python3
"""
Advanced multi-agent workflow for complex programming projects.

This orchestrates multiple specialized agents:
- architect: Designs system architecture
- coder: Writes implementation code  
- tester: Writes and runs tests
- reviewer: Reviews code quality and standards compliance

Usage:
    python agents/multiagent_workflow.py "Build a JWT authentication system"
"""

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from smolagents import CodeAgent, OpenAIModel, HfApiModel, ToolCollection
from smolagents.default_tools import WebSearchTool


class ProjectWorkflow:
    """Orchestrates a team of AI agents to build a software project."""

    def __init__(self, model, project_name: str = "Untitled"):
        self.model = model
        self.project_name = project_name
        self.agents = self._create_agents()

    def _create_agents(self):
        """Create specialized agents for different roles."""

        architect = CodeAgent(
            tools=[WebSearchTool()],
            model=self.model,
            name="architect",
            description="""Designs system architecture. Reads PRD.md and creates:
- File structure plan
- Data models
- API contract design
- Technology choices with justification""",
            max_steps=8,
        )

        coder = CodeAgent(
            tools=[],
            model=self.model,
            name="coder",
            description="""Writes production code. Requirements:
- Type hints on ALL functions
- Google docstrings
- Error handling with specific exceptions
- Follows architecture from architect agent
- No business logic in route handlers""",
            max_steps=12,
        )

        tester = CodeAgent(
            tools=[],
            model=self.model,
            name="tester",
            description="""Writes comprehensive tests. Must:
- Test happy paths and edge cases
- Test error conditions
- Use pytest fixtures from conftest.py
- Target 80%+ coverage
- Mock external APIs""",
            max_steps=8,
        )

        reviewer = CodeAgent(
            tools=[],
            model=self.model,
            name="reviewer",
            description="""Reviews code against standards. Checks:
- Type hints present?
- Docstrings complete?
- Error handling specific?
- No business logic in routes?
- Tests exist for all functions?
- No forbidden patterns (print, bare except, hardcoded values)?""",
            max_steps=6,
        )

        return {
            "architect": architect,
            "coder": coder,
            "tester": tester,
            "reviewer": reviewer,
        }

    def run(self, task: str) -> dict:
        """Execute the full workflow: architect → coder → tester → reviewer."""

        results = {}

        # Phase 1: Architecture
        print("\n" + "="*60)
        print("PHASE 1: Architecture Design")
        print("="*60)
        arch_prompt = f"""Design the architecture for this task: {task}

Read docs/PRD.md if it exists for requirements.
Create a detailed plan including:
1. File structure (which files in which directories)
2. Data models with field types
3. API endpoints with request/response schemas
4. Technology choices and why

Return your plan as structured markdown."""

        results["architecture"] = self.agents["architect"].run(arch_prompt)
        print(results["architecture"])

        # Phase 2: Implementation
        print("\n" + "="*60)
        print("PHASE 2: Code Implementation")
        print("="*60)
        code_prompt = f"""Implement the code based on this architecture:

{results['architecture']}

Task: {task}

Write production-ready code following these rules:
- Full type annotations on all functions
- Google-style docstrings
- Specific exception handling
- Business logic in services/, thin routes
- Read docs/CONTEXT.md for coding standards

Create or modify the necessary files."""

        results["code"] = self.agents["coder"].run(code_prompt)
        print(results["code"])

        # Phase 3: Testing
        print("\n" + "="*60)
        print("PHASE 3: Test Writing")
        print("="*60)
        test_prompt = f"""Write comprehensive tests for the code that was just written.

Task: {task}

Requirements:
- Create tests in tests/ mirroring src/ structure
- Test all functions including edge cases
- Use pytest fixtures
- Mock external dependencies
- Target 80%+ line coverage
- Include both unit and integration tests

Read existing tests for style consistency."""

        results["tests"] = self.agents["tester"].run(test_prompt)
        print(results["tests"])

        # Phase 4: Review
        print("\n" + "="*60)
        print("PHASE 4: Code Review")
        print("="*60)
        review_prompt = f"""Review all the code and tests produced for this task.

Task: {task}

Check against these standards from docs/CONTEXT.md:
1. All functions have type hints?
2. All public functions have Google docstrings?
3. Business logic is in services/ not routes?
4. No print() statements — only logging?
5. No bare except: clauses?
6. No hardcoded values?
7. Tests exist for all new functions?
8. Error handling is specific?

Report any issues found and suggest fixes."""

        results["review"] = self.agents["reviewer"].run(review_prompt)
        print(results["review"])

        return results


def main():
    parser = argparse.ArgumentParser(description="Multi-agent project workflow")
    parser.add_argument("task", help="The project task to execute")
    parser.add_argument("--model", default="gemma4:4b")
    parser.add_argument("--api-base", default="http://localhost:11434/v1")
    parser.add_argument("--hf", action="store_true", help="Use Hugging Face model")
    args = parser.parse_args()

    if args.hf:
        model = HfApiModel(model_id=args.model)
    else:
        model = OpenAIModel(
            model_id=args.model,
            api_base=args.api_base,
            api_key="ollama",
        )

    workflow = ProjectWorkflow(model, project_name=args.task[:50])
    results = workflow.run(args.task)

    # Save results
    output_dir = Path("agent_outputs")
    output_dir.mkdir(exist_ok=True)

    for phase, content in results.items():
        (output_dir / f"{phase}.md").write_text(str(content))

    print(f"\n✓ Results saved to {output_dir}/")
    print("  - architecture.md")
    print("  - code.md")
    print("  - tests.md")
    print("  - review.md")


if __name__ == "__main__":
    main()