"""Helpers: valid sample outputs per agent + an agent-detecting handler.""" from __future__ import annotations import json import re AGENT_MARKERS: list[tuple[str, str]] = [ ("discovery", "Discovery Agent"), ("requirements", "Requirements Engineer agent"), ("architecture", "Architecture agent"), ("database", "Database Design agent"), ("api", "API Design agent"), ("devops", "DevOps Engineer agent"), ("reviewer", "Review Agent"), ("summarizer", "Artifact Summarizer agent"), ] def detect_agent(system_prompt: str) -> str: for name, marker in AGENT_MARKERS: if marker in system_prompt: return name return "unknown" def discovery_output(status: str = "ready") -> dict: return { "status": status, "confidence": 0.94 if status == "ready" else 0.6, "summary": "A platform where users order food from local restaurants and track delivery.", "known_information": { "problem": "Ordering food takes too long", "target_users": ["Customers", "Restaurants"], "user_roles": ["Customer", "Restaurant Owner", "Admin"], "business_goals": ["Increase orders", "Fast delivery"], "core_features": [ "Restaurant discovery", "Order placement", "Online payment", "Delivery tracking", ], "payment_requirement": "Online payment", "notification_requirement": "Order status notifications", }, "missing_information": [] if status == "ready" else [ {"field": "user_roles", "importance": "critical", "reason": "Required to define authorization"} ], "questions": [] if status == "ready" else [ {"id": "q1", "question": "Who are the main users of the platform?", "reason": "Defines roles and permissions."} ], } def requirements_output() -> dict: return { "functional_requirements": [ "FR1: Users can browse restaurants and menus.", "FR2: Users can place orders and pay online.", "FR3: Users can track delivery status.", ], "non_functional_requirements": [ "NFR1: The system must respond within 500ms for core operations.", "NFR2: Payment data must be encrypted in transit and at rest.", ], "user_stories": [ "As a customer, I want to order food online, so that I save time.", "As a restaurant owner, I want to receive orders, so that I can fulfill them.", ], "acceptance_criteria": [ "AC1: A customer can complete an order end-to-end.", "AC2: Payment failures return a clear error.", ], "constraints": ["Must run on commodity cloud infrastructure."], "assumptions": ["Payment provider is Stripe."], } def architecture_output() -> dict: return { "system_components": [ {"name": "Web Frontend", "type": "frontend", "description": "React SPA", "technology": "React"}, {"name": "API Backend", "type": "backend", "description": "REST API", "technology": "FastAPI (Python)"}, {"name": "Database", "type": "database", "description": "Primary datastore", "technology": "PostgreSQL"}, {"name": "Payments", "type": "external", "description": "Payment processing", "technology": "Stripe"}, ], "communication": ["Frontend calls the API over HTTPS", "API persists to PostgreSQL"], "authentication": "JWT bearer tokens", "security": ["TLS everywhere", "Rate limiting on auth endpoints"], "scalability": ["Horizontal scaling of the API behind a load balancer"], "technology_stack": { "frontend": "React", "backend": "FastAPI (Python)", "database": "PostgreSQL", "payments": "Stripe", }, "deployment_architecture": "Containerized services on a single cloud VM, scaled out as needed.", "mermaid_diagram": "flowchart TD\n FE[Web Frontend] --> API[API Backend]\n API --> DB[(PostgreSQL)]\n API --> PS[Stripe]", } def database_output() -> dict: return { "database_technology": "PostgreSQL", "entities": [ { "name": "users", "description": "Platform users", "fields": [ {"name": "id", "type": "uuid", "primary_key": True, "nullable": False, "unique": False, "indexed": False, "foreign_key": None}, {"name": "email", "type": "varchar", "primary_key": False, "nullable": False, "unique": True, "indexed": True, "foreign_key": None}, {"name": "role", "type": "varchar", "primary_key": False, "nullable": False, "unique": False, "indexed": False, "foreign_key": None}, ], }, { "name": "orders", "description": "Food orders", "fields": [ {"name": "id", "type": "uuid", "primary_key": True, "nullable": False, "unique": False, "indexed": False, "foreign_key": None}, {"name": "user_id", "type": "uuid", "primary_key": False, "nullable": False, "unique": False, "indexed": True, "foreign_key": "users.id"}, {"name": "status", "type": "varchar", "primary_key": False, "nullable": False, "unique": False, "indexed": False, "foreign_key": None}, ], }, ], "relationships": ["orders.user_id references users.id"], "indexes": ["idx_orders_user_id ON orders(user_id)"], "constraints": ["orders.status CHECK in ('pending','paid','delivered')"], "sql_schema": "CREATE TABLE users (id uuid PRIMARY KEY, email varchar NOT NULL UNIQUE, role varchar NOT NULL);\nCREATE TABLE orders (id uuid PRIMARY KEY, user_id uuid NOT NULL REFERENCES users(id), status varchar NOT NULL);", "erd_mermaid": "erDiagram\n users ||--o{ orders : places", } def api_output() -> dict: return { "endpoints": [ { "method": "GET", "path": "/api/restaurants", "summary": "List restaurants", "auth": "none", "pagination": True, "filters": ["cuisine"], "request_schema": None, "response_schema": {"type": "object", "properties": {"items": {"type": "array"}}}, }, { "method": "POST", "path": "/api/orders", "summary": "Place an order", "auth": "jwt", "pagination": False, "filters": [], "request_schema": {"type": "object", "properties": {"restaurant_id": {"type": "string"}}}, "response_schema": {"type": "object", "properties": {"id": {"type": "string"}}}, }, { "method": "GET", "path": "/api/orders/{id}", "summary": "Get an order", "auth": "jwt", "pagination": False, "filters": [], "request_schema": None, "response_schema": {"type": "object", "properties": {"status": {"type": "string"}}}, }, ], "authentication": "JWT bearer tokens", "authorization": "Customers access their own orders; owners access their restaurant data", "error_handling": ["Standard error envelope with code/message/request_id", "404 for missing resources"], "pagination": "cursor-based", "filtering": "query parameters", "openapi_spec": {"openapi": "3.0.0", "info": {"title": "Food Delivery API", "version": "1.0.0"}, "paths": {}}, } def devops_output() -> dict: return { "dockerfile": "FROM python:3.11-slim\nWORKDIR /app\nCOPY requirements.txt .\nRUN pip install -r requirements.txt\nCOPY . .\nCMD [\"uvicorn\", \"app.main:app\", \"--host\", \"0.0.0.0\"]", "docker_compose": "services:\n api:\n build: .\n ports: ['8000:8000']\n db:\n image: postgres:16", "ci_cd_pipeline": "Lint, test, build image, push to registry, deploy to VM", "github_actions": "name: CI\non: [push]\njobs:\n test:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v4", "environment_variables": {"DATABASE_URL": "postgres://user:pass@db:5432/app", "STRIPE_SECRET_KEY": "${STRIPE_SECRET_KEY}"}, "deployment_strategy": "Docker Compose on a single VM with rolling container restarts", "health_checks": ["/health endpoint on API", "pg_isready for database"], "logging": ["Structured JSON logs to stdout"], "monitoring": ["Prometheus metrics on /metrics"], "secrets_management": "Secrets injected via environment variables from the CI/CD provider secret store", } def review_output(status: str = "approved") -> dict: if status == "approved": return {"status": "approved", "score": 0.96, "issues": [], "artifacts_to_regenerate": []} return { "status": "needs_revision", "score": 0.7, "issues": [ { "artifact": "database", "severity": "blocking", "problem": "Database technology conflicts with architecture.", "expected": "PostgreSQL", "actual": "MongoDB", "fix": "Switch database_technology to PostgreSQL.", "source_artifact": "architecture", "source_decision": "Database component technology: PostgreSQL", "conflicting_artifact": "database", "conflicting_decision": "database_technology: MongoDB", } ], "artifacts_to_regenerate": ["database"], } def review_output_targets(targets: list[str]) -> dict: """A needs_revision review flagging ``targets`` as blocking.""" issues = [ { "artifact": target, "severity": "blocking", "problem": f"{target} artifact is inconsistent.", "expected": "expected value", "actual": "actual value", "fix": "fix it", "source_artifact": "database", "source_decision": "decision", "conflicting_artifact": target, "conflicting_decision": "decision", } for target in targets ] return { "status": "needs_revision", "score": 0.6, "issues": issues, "artifacts_to_regenerate": targets, } def summary_text(name: str) -> str: """A plausible compact summary for each artifact (plain text, as the real summarizer produces).""" return { "requirements": ( "Requirements summary: FR1-FR3 and NFR1-NFR2 captured; " "key constraints preserved." ), "architecture": ( "Architecture summary: Web Frontend (React), API Backend (FastAPI), " "Database (PostgreSQL), Stripe; JWT auth." ), "database": ( "Database summary: PostgreSQL; entities users (id, email, role), " "orders (id, user_id, status); foreign key orders.user_id -> users.id." ), "api": ( "API summary: GET /api/restaurants, POST /api/orders, " "GET /api/orders/{id}; JWT auth; cursor pagination." ), "devops": ( "DevOps summary: Docker Compose on a VM; health checks /health and " "pg_isready; Prometheus metrics; env-var secrets." ), }.get(name, f"Summary of {name}: key facts preserved.") def build_handler(discovery_status: str = "ready", review_status: str = "approved", review_sequence: list[str] | None = None): """Return a handler that answers every agent with valid output. ``review_sequence`` overrides review answers call-by-call when provided. A revision call (user prompt contains ``REVISION TASK``) returns a *changed* artifact so the orchestrator's hash-based convergence check sees a real change (as the real LLM would produce). """ review_index = 0 reviews = list(review_sequence) if review_sequence else None def handler(system_prompt: str, user_prompt: str) -> str: nonlocal review_index agent = detect_agent(system_prompt) if agent == "discovery": return json.dumps(discovery_output(discovery_status)) if agent == "requirements": base = requirements_output() return json.dumps(revised_output(agent, base) if "REVISION TASK" in user_prompt else base) if agent == "architecture": base = architecture_output() return json.dumps(revised_output(agent, base) if "REVISION TASK" in user_prompt else base) if agent == "database": base = database_output() return json.dumps(revised_output(agent, base) if "REVISION TASK" in user_prompt else base) if agent == "api": base = api_output() return json.dumps(revised_output(agent, base) if "REVISION TASK" in user_prompt else base) if agent == "devops": base = devops_output() return json.dumps(revised_output(agent, base) if "REVISION TASK" in user_prompt else base) if agent == "reviewer": if reviews is not None: answer = reviews[min(review_index, len(reviews) - 1)] review_index += 1 return json.dumps(review_output(answer)) return json.dumps(review_output(review_status)) if agent == "summarizer": match = re.search(r"Artifact: (\w+)", user_prompt) name = match.group(1) if match else "unknown" return summary_text(name) raise AssertionError(f"Unknown agent marker in: {system_prompt[:80]}") return handler def revised_output(agent: str, base: dict) -> dict: """Return a *changed* artifact for a revision run, so the orchestrator's hash-based convergence check sees the artifact actually changed.""" out = dict(base) if agent == "database": out["constraints"] = list(base.get("constraints") or []) + ["review_revision_applied"] elif agent == "api": out["authentication"] = (base.get("authentication") or "") + " (revised)" elif agent == "devops": out["deployment_strategy"] = (base.get("deployment_strategy") or "") + " (revised)" elif agent == "architecture": out["deployment_architecture"] = (base.get("deployment_architecture") or "") + " (revised)" elif agent == "requirements": out["constraints"] = list(base.get("constraints") or []) + ["review_revision_applied"] return out