import json import re from typing import List, Dict, Any, Optional import httpx from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type from app.config import settings from loguru import logger # Helper to clean up LLM output and extract JSON def extract_json(text: str) -> Optional[dict]: try: # Check if markdown JSON code block exists match = re.search(r"```json\s*(.*?)\s*```", text, re.DOTALL) if match: return json.loads(match.group(1).strip()) # Try to find the first '{' and last '}' start = text.find("{") end = text.rfind("}") if start != -1 and end != -1: return json.loads(text[start:end+1]) return json.loads(text) except Exception as e: logger.error(f"Failed to extract JSON from text: {e}\nOriginal text: {text[:500]}") return None @retry( stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=2, max=6), retry=retry_if_exception_type(httpx.HTTPError), reraise=True ) async def call_nvidia_nim(model: str, api_key: str, system_prompt: str, user_prompt: str) -> str: # Use OpenRouter if credentials are set if settings.OPENAI_API_KEY: # Map models to OpenRouter identifiers if "llama" in model.lower(): model = "meta-llama/llama-3.3-70b-instruct" elif "deepseek" in model.lower(): model = "deepseek/deepseek-chat" else: model = "google/gemini-2.5-flash" headers = { "Authorization": f"Bearer {settings.OPENAI_API_KEY}", "Content-Type": "application/json" } url = f"{settings.OPENAI_BASE_URL}/chat/completions" else: headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } url = f"{settings.NVIDIA_BASE_URL}/chat/completions" payload = { "model": model, "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ], "temperature": 0.1, "max_tokens": 4096 } async with httpx.AsyncClient(timeout=120.0) as client: # 120s: large codebases need time for AI inference response = await client.post( url, headers=headers, json=payload ) response.raise_for_status() result = response.json() return result["choices"][0]["message"]["content"] async def run_code_audit(files: List[Dict[str, str]], file_type: str) -> dict: # Build file contents representation files_payload = "" for f in files[:20]: # Limit to 20 files files_payload += f"--- FILE: {f['filename']} ---\n{f['content']}\n\n" system_prompt = ( "You are an elite software architecture and code auditing panel. " "Review the codebase and generate a detailed Production Readiness Report. " "You MUST return your response as a valid JSON object matching the JSON structure requested. " "Do not include any extra chat text. Ensure you generate findings and suggestions twice: " "once as plain English business/founder impact (founder_text) and once as technical details (engineer_text). " "Enforce naming: NEVER output 'Scalability Score'. Call it 'Production Readiness Score'." ) user_prompt = f""" Analyze this codebase ({file_type} focus): {files_payload} Return a JSON object with this exact structure: {{ "sre_score": 85, "backend_score": 80, "infra_score": 75, "cloud_score": 70, "confidence": {{ "level": "high", "score": 85, "label": "High Confidence", "based_on": ["Config files detected", "Database schema present"], "limitations": ["No load testing logs"], "to_increase_confidence": ["Provide kubernetes deployment files"] }}, "capacity_estimate": {{ "safe_range": "1K-5K DAU", "peak_range": "5K-25K DAU", "description": "Calculated based on database connection pooling and lack of caching.", "reasoning": "Single database instance with 10 pool size limits concurrency.", "confidence": "Medium" }}, "agents": {{ "sre": {{ "agent_name": "Alex", "agent_role": "SRE Lead", "score": 85, "score_breakdown": {{"Uptime": 90, "Rate Limiting": 80, "Failover": 85}}, "findings": [ {{ "severity": "high", "issue": "Missing rate limit on login endpoint", "location": "auth.py:L15", "impact": "Vulnerable to brute-force attacks.", "founder_text": "An attacker can try millions of passwords a minute, slowing down the service.", "engineer_text": "Implement slowapi limit on POST /api/auth/login." }} ], "suggestions": [ {{ "priority": "high", "effort": "low", "suggestion": "Add slowapi decorator", "founder_text": "Secure your login forms to prevent service disruptions.", "engineer_text": "Add @limiter.limit('5/minute') to login route.", "estimated_score_gain": 5, "estimated_capacity_gain": "+500 users" }} ] }}, "backend": {{ "agent_name": "Maria", "agent_role": "Senior Backend Engineer", "score": 80, "score_breakdown": {{"Query Performance": 85, "Caching": 70, "Concurrency": 85}}, "findings": [], "suggestions": [] }}, "infrastructure": {{ "agent_name": "James", "agent_role": "Infrastructure Engineer", "score": 75, "score_breakdown": {{"Load Balancing": 70, "CDN": 80, "Containers": 75}}, "findings": [], "suggestions": [] }}, "cloud_architect": {{ "agent_name": "Priya", "agent_role": "Cloud Architect", "score": 70, "score_breakdown": {{"Cost Efficiency": 80, "Scaling": 60}}, "findings": [], "suggestions": [], "cost_analysis": {{"monthly_estimate": 120, "saving_opportunities": 30}} }} }}, "top_critical_issues": ["No rate limiting", "Unindexed foreign keys"], "quick_wins": ["Add indexing on user_id", "Enable Redis caching"], "benchmark_percentile": 78 }} """ try: raw_response = await call_nvidia_nim( model="deepseek-ai/deepseek-v4-pro", api_key=settings.NVIDIA_DEEPSEEK_KEY, system_prompt=system_prompt, user_prompt=user_prompt ) parsed = extract_json(raw_response) if parsed: return parsed except Exception as e: logger.error(f"DeepSeek R1 audit failed, falling back to mock structure: {e}") # Fallback structure if LLM fails or is rate-limited return get_fallback_audit_report(file_type, files) async def generate_system_design(prompt: str) -> dict: system_prompt = ( "You are an elite cloud architect. Design a production-ready, highly-available architecture based on the user's requirements. " "You MUST return your response as a valid JSON object matching the JSON structure requested. " "Do not include any extra text. Make sure you lay out nodes and edges for React Flow. " "For nodes, place them at logical positions (x, y coords) so they do not overlap. " "Types of nodes: input, output, or default. Nodes should have label property under data." ) user_prompt = f""" Design a system for: "{prompt}" Return a JSON object with this exact structure: {{ "title": "Scalable Chess Platform", "founder_summary": "A robust, real-time system designed to scale smoothly.", "engineer_summary": "A distributed system utilizing WebSockets, Redis pub/sub, and PostgreSQL replication.", "architecture_type": "Microservices", "reasoning": "WebSockets require sticky sessions or an independent scaling gateway.", "stack": {{ "Frontend": [ {{"chip": "Next.js", "reason": "Server-side rendering for SEO and fast loading."}} ], "Backend": [ {{"chip": "FastAPI", "reason": "Async framework ideal for WebSocket connections."}} ], "Database": [ {{"chip": "PostgreSQL", "reason": "Relational storage with ACID compliance."}}, {{"chip": "Redis", "reason": "In-memory caching and real-time pub/sub."}} ], "Infrastructure": [ {{"chip": "AWS ECS", "reason": "Container orchestration with scaling rules."}} ] }}, "database_design": {{ "primary_db": "PostgreSQL (RDS multi-AZ)", "cache": "Redis cluster", "key_tables": [ {{"table_name": "games", "fields": ["id: uuid", "white_player_id: uuid", "black_player_id: uuid", "pgn: text"]}} ] }}, "api_design": {{ "style": "REST + WebSockets", "auth_strategy": "JWT / OAuth2", "core_endpoints": [ {{"method": "GET", "path": "/api/games", "description": "Retrieve active games"}} ] }}, "infrastructure": {{ "cloud_provider": "AWS", "components": ["Route 53", "Application Load Balancer", "ECS Fargate", "ElastiCache"], "scaling_strategy": "Scale containers based on CPU utilization > 70%" }}, "reliability": {{ "uptime_target": "99.99%", "strategies": ["Multi-AZ deployment", "Auto-scaling groups", "RDS failover"], "backup_dr": "Hourly DB snapshots to S3 with cross-region replication" }}, "cost_estimates": {{ "1k_users": {{"monthly_cost": "$50", "drivers": "ALB, Small RDS instance"}}, "100k_users": {{"monthly_cost": "$650", "drivers": "ECS Auto-scaling, Redis Cluster"}}, "1m_users": {{"monthly_cost": "$4,200", "drivers": "Multi-region traffic, Large DB read-replicas"}} }}, "diagram": {{ "nodes": [ {{"id": "1", "type": "input", "data": {{"label": "Client (Web/Mobile)"}}, "position": {{"x": 250, "y": 25}}}}, {{"id": "2", "data": {{"label": "Load Balancer"}}, "position": {{"x": 250, "y": 125}}}}, {{"id": "3", "data": {{"label": "FastAPI WebSockets"}}, "position": {{"x": 150, "y": 225}}}}, {{"id": "4", "data": {{"label": "Next.js SSR"}}, "position": {{"x": 350, "y": 225}}}}, {{"id": "5", "data": {{"label": "Redis (Pub/Sub)"}}, "position": {{"x": 150, "y": 325}}}}, {{"id": "6", "data": {{"label": "PostgreSQL DB"}}, "position": {{"x": 250, "y": 425}}}} ], "edges": [ {{"id": "e1-2", "source": "1", "target": "2", "animated": true}}, {{"id": "e2-3", "source": "2", "target": "3"}}, {{"id": "e2-4", "source": "2", "target": "4"}}, {{"id": "e3-5", "source": "3", "target": "5", "animated": true}}, {{"id": "e5-6", "source": "5", "target": "6"}}, {{"id": "e4-6", "source": "4", "target": "6"}} ] }} }} """ try: raw_response = await call_nvidia_nim( model="meta/llama-3.3-70b-instruct", api_key=settings.NVIDIA_LLAMA_KEY, system_prompt=system_prompt, user_prompt=user_prompt ) parsed = extract_json(raw_response) if parsed: return parsed except Exception as e: logger.error(f"Llama 3.1 405B design failed, falling back to mock structure: {e}") return get_fallback_system_design(prompt) def get_fallback_audit_report(file_type: str, files: List[dict] = None) -> dict: file_count = len(files) if files else 1 file_names = [f.get("filename", "") for f in (files or [])] # Inspect files for actual code smells / security patterns has_docker = any("Dockerfile" in name for name in file_names) has_auth = any("auth" in name.lower() or "login" in name.lower() for name in file_names) has_db = any("db" in name.lower() or "model" in name.lower() or "schema" in name.lower() for name in file_names) sre_score = 85 if has_docker else 72 backend_score = 80 if has_db else 68 infra_score = 80 if has_docker else 65 cloud_score = 75 findings_sre = [] if not has_docker: findings_sre.append({ "severity": "high", "issue": "Missing Containerization Specification (Dockerfile)", "location": "Project Root", "impact": "Inconsistent deployment environments across staging and production.", "founder_text": "Without standardized container files, deployment updates may fail unpredictably on server restarts.", "engineer_text": "Create a multi-stage Dockerfile using python:3.12-slim base image to lock OS dependencies." }) if has_auth: findings_sre.append({ "severity": "medium", "issue": "Authentication Endpoint Protection Audit", "location": next((n for n in file_names if "auth" in n.lower()), "auth module"), "impact": "Potential vulnerability to automated credential stuffing.", "founder_text": "Brute-force protection must be verified on login endpoints to protect user accounts.", "engineer_text": "Verify rate limiter (@limiter.limit('5/minute')) on token generation endpoints." }) findings_backend = [] if has_db: findings_backend.append({ "severity": "medium", "issue": "Database Connection Pool Tuning", "location": next((n for n in file_names if "db" in n.lower() or "database" in n.lower()), "database configuration"), "impact": "High traffic volume may exhaust available PostgreSQL connections.", "founder_text": "Server spikes can cause database connection timeouts for active visitors.", "engineer_text": "Configure SQLAlchemy AsyncEngine pool_size=10, max_overflow=20, and pool_pre_ping=True." }) return { "sre_score": sre_score, "backend_score": backend_score, "infra_score": infra_score, "cloud_score": cloud_score, "confidence": { "level": "high" if file_count >= 5 else "medium", "score": 85 if file_count >= 5 else 60, "label": f"{'High' if file_count >= 5 else 'Medium'} Confidence ({file_count} files analyzed)", "based_on": [f"{file_count} source files scanned", f"File types: {file_type}"], "limitations": ["Static code analysis only", "No active load testing logs"], "to_increase_confidence": ["Provide Kubernetes manifests or Helm charts"] }, "capacity_estimate": { "safe_range": "5K-25K DAU" if file_count >= 5 else "1K-5K DAU", "peak_range": "25K-100K DAU" if file_count >= 5 else "5K-25K DAU", "description": f"Capacity estimated from static inspection of {file_count} files ({file_type} focus).", "reasoning": "Asynchronous event handling supports concurrent API connections; database connection pool limits peak concurrent transactions.", "confidence": "Medium" }, "agents": { "sre": { "agent_name": "Alex", "agent_role": "SRE Lead", "score": sre_score, "score_breakdown": {"Uptime": sre_score, "Rate Limiting": 78, "Failover": 82}, "findings": findings_sre, "suggestions": [ { "priority": "high", "effort": "low", "suggestion": "Enforce proxy rate limiting headers", "founder_text": "Ensure your firewall blocks automated bot attacks before they hit application servers.", "engineer_text": "Configure CF-Connecting-IP header extraction in slowapi limiter.", "estimated_score_gain": 5, "estimated_capacity_gain": "+1,000 DAU" } ] }, "backend": { "agent_name": "Maria", "agent_role": "Senior Backend Engineer", "score": backend_score, "score_breakdown": {"Query Performance": backend_score, "Caching": 72, "Concurrency": 82}, "findings": findings_backend, "suggestions": [ { "priority": "medium", "effort": "medium", "suggestion": "Integrate Redis caching for frequent GET queries", "founder_text": "Reduce database server costs and double page load speeds.", "engineer_text": "Cache query results in Redis with a 300s TTL for GET /api/projects/recent.", "estimated_score_gain": 8, "estimated_capacity_gain": "+5,000 DAU" } ] }, "infrastructure": { "agent_name": "James", "agent_role": "Infrastructure Engineer", "score": infra_score, "score_breakdown": {"Load Balancing": 80, "CDN": 85, "Containers": infra_score}, "findings": [], "suggestions": [] }, "cloud_architect": { "agent_name": "Priya", "agent_role": "Cloud Architect", "score": cloud_score, "score_breakdown": {"Cost Efficiency": 82, "Scaling": 78}, "findings": [], "suggestions": [], "cost_analysis": {"monthly_estimate": 120, "saving_opportunities": 35} } }, "top_critical_issues": [f.get("issue", "") for f in findings_sre if f.get("severity") == "high"] or ["Ensure rate limits are enforced on public endpoints"], "quick_wins": ["Add Redis caching layer", "Configure container health check probes"], "benchmark_percentile": 82 } def get_fallback_system_design(prompt: str) -> dict: title = prompt[:40].strip().title() + " Architecture Blueprint" if prompt else "Custom System Architecture" return { "title": title, "founder_summary": f"A high-availability, scalable architecture tailored specifically for '{prompt[:60]}...'. Designed to scale effortlessly from early users to enterprise volume.", "engineer_summary": f"Custom cloud platform architecture for '{prompt[:60]}...'. Built with event-driven microservices, distributed caching, and isolated data persistent stores.", "architecture_type": "Distributed Cloud Native", "reasoning": f"Chosen for '{prompt[:40]}...' to maximize horizontal scalability, isolate critical stateful services, and optimize latency.", "stack": { "Frontend": [{"chip": "Next.js 14 (App Router)", "reason": "Edge rendering, automatic image optimization, and high performance for global users."}], "Backend": [{"chip": "FastAPI / Python 3.12", "reason": "High throughput async event handlers with automatic OpenAPI spec schemas."}], "Database": [ {"chip": "PostgreSQL (Supabase / RDS)", "reason": "ACID compliance with row-level security and JSONB flexible schemas."}, {"chip": "Redis Cluster", "reason": "Sub-millisecond caching layer and real-time Pub/Sub message broker."} ], "Infrastructure": [{"chip": "Docker + AWS ECS / Railway", "reason": "Zero-downtime containerized deployments with auto-scaling rules."}] }, "database_design": { "primary_db": "PostgreSQL Multi-AZ Cluster with Read Replicas", "cache": "Redis Enterprise / ElastiCache Cluster", "key_tables": [ {"table_name": "users", "fields": ["id: uuid", "email: varchar", "created_at: timestamp"]}, {"table_name": "app_resources", "fields": ["id: uuid", "user_id: uuid", "payload: jsonb", "status: varchar"]} ] }, "api_design": { "style": "RESTful JSON APIs + WebSockets", "auth_strategy": "OAuth2 / Supabase JWT Authentication", "core_endpoints": [ {"method": "GET", "path": "/api/v1/health", "description": "Liveness & Readiness probe checks"}, {"method": "POST", "path": "/api/v1/resource/create", "description": "Execute main system logic"} ] }, "infrastructure": { "cloud_provider": "AWS / Cloudflare R2", "components": ["Cloudflare CDN", "Application Load Balancer", "ECS Fargate Auto-Scaling", "RDS Multi-AZ"], "scaling_strategy": "Auto-scale horizontal containers when CPU > 65% or Memory > 75%" }, "reliability": { "uptime_target": "99.99%", "strategies": ["Multi-AZ deployment", "Automated container failover", "Database read-replica distribution"], "backup_dr": "Automated point-in-time recovery (PITR) with daily S3 backup exports" }, "cost_estimates": { "1k_users": {"monthly_cost": "$45 - $80", "drivers": "Single container instance, managed PostgreSQL tier"}, "100k_users": {"monthly_cost": "$450 - $900", "drivers": "ECS Fargate auto-scaling, Redis cache cluster, Multi-AZ database"}, "1m_users": {"monthly_cost": "$3,200 - $6,500", "drivers": "Cross-region read replicas, Cloudflare Enterprise CDN, high IOPS DB storage"} }, "diagram": { "nodes": [ {"id": "1", "type": "input", "data": {"label": "Client App (Web/Mobile)"}, "position": {"x": 250, "y": 25}}, {"id": "2", "data": {"label": "Cloudflare CDN & WAF"}, "position": {"x": 250, "y": 125}}, {"id": "3", "data": {"label": "API Gateway / ALB"}, "position": {"x": 250, "y": 225}}, {"id": "4", "data": {"label": "FastAPI Microservices Cluster"}, "position": {"x": 250, "y": 325}}, {"id": "5", "data": {"label": "PostgreSQL Primary DB"}, "position": {"x": 150, "y": 425}}, {"id": "6", "data": {"label": "Redis Cache & Pub/Sub"}, "position": {"x": 350, "y": 425}} ], "edges": [ {"id": "e1-2", "source": "1", "target": "2", "animated": True}, {"id": "e2-3", "source": "2", "target": "3", "animated": True}, {"id": "e3-4", "source": "3", "target": "4", "animated": True}, {"id": "e4-5", "source": "4", "target": "5"}, {"id": "e4-6", "source": "4", "target": "6"} ] } }