import asyncio import json from datetime import datetime, timedelta import redis from sqlalchemy import update from sqlalchemy.future import select from app.config import settings from contextlib import asynccontextmanager _task_engine = None _task_session_factory = None def get_task_session_factory(): global _task_engine, _task_session_factory if _task_session_factory is None: from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from uuid import uuid4 database_url = settings.DATABASE_URL if database_url.startswith("postgresql://"): database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1) _task_engine = create_async_engine( database_url, pool_size=5, max_overflow=10, pool_pre_ping=True, connect_args={ "statement_cache_size": 0, "prepared_statement_cache_size": 0, "prepared_statement_name_func": lambda: f"__asyncpg_{uuid4().hex}__" } ) _task_session_factory = async_sessionmaker(bind=_task_engine, class_=AsyncSession, expire_on_commit=False) return _task_session_factory @asynccontextmanager async def local_session(): factory = get_task_session_factory() async with factory() as session: yield session from app import models from app.utils.ai import run_code_audit, generate_system_design from loguru import logger # Redis client for publishing streaming progress events redis_client = redis.Redis.from_url(settings.REDIS_URL) def publish_event(job_id: str, event_type: str, data: dict = None): payload = {"event": event_type} if data: payload.update(data) raw_payload = json.dumps(payload) redis_client.publish(f"job:{job_id}", raw_payload) try: redis_client.rpush(f"job:{job_id}:events", raw_payload) redis_client.expire(f"job:{job_id}:events", 3600) except Exception as e: logger.error(f"Failed to cache event in Redis: {e}") async def run_audit_task(project_id_str: str, job_id: str, files: list, file_type: str): logger.info(f"Starting audit task for project {project_id_str}, job {job_id}") # Publish initial agent start events publish_event(job_id, "agent_started", {"agent": "connecting"}) publish_event(job_id, "agent_complete", {"agent": "connecting"}) publish_event(job_id, "agent_started", {"agent": "sre"}) publish_event(job_id, "agent_started", {"agent": "backend"}) publish_event(job_id, "agent_started", {"agent": "infrastructure"}) publish_event(job_id, "agent_started", {"agent": "cloud_architect"}) try: # Run AI audit analysis report_data = await run_code_audit(files, file_type) # Calculate stats for confidence score files_count = len(files) total_chars = sum(len(f["content"]) for f in files) has_config_files = any( any(f["filename"].endswith(ext) for ext in [".json", ".yaml", ".yml", ".toml", ".config", ".env"]) for f in files ) has_docker = any("Dockerfile" in f["filename"] for f in files) # Confidence score weighting: # files>=10(+20), config_files(+15), docker(+15), chars>=50K(+20), not_truncated(+20) confidence_score = 0 if files_count >= 10: confidence_score += 20 if has_config_files: confidence_score += 15 if has_docker: confidence_score += 15 if total_chars >= 50000: confidence_score += 20 confidence_score += 20 # Assuming not truncated by default if confidence_score >= 70: confidence_level = "high" elif confidence_score >= 40: confidence_level = "medium" else: confidence_level = "low" confidence_dict = { "level": confidence_level, "score": confidence_score, "label": f"{confidence_level.capitalize()} Confidence", "based_on": report_data.get("confidence", {}).get("based_on", ["Initial file scan completed"]), "limitations": report_data.get("confidence", {}).get("limitations", ["No runtime metrics available"]), "to_increase_confidence": report_data.get("confidence", {}).get("to_increase_confidence", ["Upload load test scripts"]) } # Enforce server-side score calculation # overall_score = round(sre*0.30 + backend*0.30 + infra*0.20 + cloud*0.20) sre_score = report_data.get("sre_score", 80) backend_score = report_data.get("backend_score", 80) infra_score = report_data.get("infra_score", 80) cloud_score = report_data.get("cloud_score", 80) overall_score = round( sre_score * 0.30 + backend_score * 0.30 + infra_score * 0.20 + cloud_score * 0.20 ) # Enforce capacity tiers # 0-30→50-200 · 31-50→200-1K · 51-65→1K-5K · 66-79→5K-25K · 80-89→25K-100K · 90-100→100K-500K if overall_score <= 30: safe, peak = "50-200 DAU", "200-1K DAU" elif overall_score <= 50: safe, peak = "200-1K DAU", "1K-5K DAU" elif overall_score <= 65: safe, peak = "1K-5K DAU", "5K-25K DAU" elif overall_score <= 79: safe, peak = "5K-25K DAU", "25K-100K DAU" elif overall_score <= 89: safe, peak = "25K-100K DAU", "100K-500K DAU" else: safe, peak = "100K-500K DAU", "500K-2M DAU" capacity_dict = { "safe_range": safe, "peak_range": peak, "description": report_data.get("capacity_estimate", {}).get("description", "DAU estimate based on static architecture scanning."), "reasoning": report_data.get("capacity_estimate", {}).get("reasoning", "Database concurrency constraints set maximum peaks."), "confidence": report_data.get("capacity_estimate", {}).get("confidence", "Medium") } # Prepare Agent reports mapping agents_data = report_data.get("agents", {}) # Re-set agents overall scores if modified if "sre" in agents_data: agents_data["sre"]["score"] = sre_score if "backend" in agents_data: agents_data["backend"]["score"] = backend_score if "infrastructure" in agents_data: agents_data["infrastructure"]["score"] = infra_score if "cloud_architect" in agents_data: agents_data["cloud_architect"]["score"] = cloud_score disclaimer = ( "This score is based on static code analysis only. It does " "not guarantee runtime performance. Validate with load testing " "before production launch." ) async with local_session() as session: # Load project import uuid project_uuid = uuid.UUID(project_id_str) proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid)) project = proj_result.scalars().first() if not project: raise Exception("Project not found") # Create Audit Report audit_report = models.AuditReport( project_id=project.id, overall_score=overall_score, confidence=confidence_dict, capacity_estimate=capacity_dict, agents=agents_data, top_critical_issues=report_data.get("top_critical_issues", []), quick_wins=report_data.get("quick_wins", []), benchmark_percentile=report_data.get("benchmark_percentile", 75.0), score_disclaimer=disclaimer, files_analyzed=files_count, files_skipped=0, was_truncated=False ) session.add(audit_report) project.status = "complete" # Update user stats user_result = await session.execute(select(models.User).where(models.User.id == project.user_id)) user = user_result.scalars().first() if user: user.total_audits += 1 session.add(user) await session.commit() # Complete agents progress steps publish_event(job_id, "agent_complete", {"agent": "sre"}) publish_event(job_id, "agent_complete", {"agent": "backend"}) publish_event(job_id, "agent_complete", {"agent": "infrastructure"}) publish_event(job_id, "agent_complete", {"agent": "cloud_architect"}) publish_event(job_id, "job_complete", {"audit_id": str(project.id)}) logger.info(f"Audit job {job_id} completed successfully") except Exception as e: logger.error(f"Audit task failed: {e}") async with local_session() as session: import uuid project_uuid = uuid.UUID(project_id_str) proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid)) project = proj_result.scalars().first() if project: project.status = "failed" project.error_msg = str(e) await session.commit() publish_event(job_id, "job_failed", {"error": str(e)}) async def run_design_task(project_id_str: str, job_id: str, idea_prompt: str): logger.info(f"Starting design task for project {project_id_str}, job {job_id}") # Emit progress events without artificial delay publish_event(job_id, "agent_started", {"agent": "connecting"}) publish_event(job_id, "agent_complete", {"agent": "connecting"}) publish_event(job_id, "agent_started", {"agent": "sre"}) publish_event(job_id, "agent_started", {"agent": "backend"}) publish_event(job_id, "agent_started", {"agent": "infrastructure"}) publish_event(job_id, "agent_started", {"agent": "cloud_architect"}) try: # Run AI System Design design_data = await generate_system_design(idea_prompt) async with local_session() as session: # Load project import uuid project_uuid = uuid.UUID(project_id_str) proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid)) project = proj_result.scalars().first() if not project: raise Exception("Project not found") # Create System Design system_design = models.SystemDesign( project_id=project.id, idea_prompt=idea_prompt, title=design_data.get("title", "System Architecture Blueprint"), founder_summary=design_data.get("founder_summary", ""), engineer_summary=design_data.get("engineer_summary", ""), architecture_type=design_data.get("architecture_type", "Hybrid"), reasoning=design_data.get("reasoning", ""), stack=design_data.get("stack", {}), database_design=design_data.get("database_design", {}), api_design=design_data.get("api_design", {}), infrastructure=design_data.get("infrastructure", {}), reliability=design_data.get("reliability", {}), cost_estimates=design_data.get("cost_estimates", {}), diagram=design_data.get("diagram", {"nodes": [], "edges": []}) ) session.add(system_design) project.status = "complete" # Update user stats user_result = await session.execute(select(models.User).where(models.User.id == project.user_id)) user = user_result.scalars().first() if user: user.total_designs += 1 session.add(user) await session.commit() publish_event(job_id, "agent_complete", {"agent": "cloud_architect"}) publish_event(job_id, "job_complete", {"design_id": str(project.id)}) logger.info(f"Design job {job_id} completed successfully") except Exception as e: logger.error(f"Design task failed: {e}") async with local_session() as session: import uuid project_uuid = uuid.UUID(project_id_str) proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid)) project = proj_result.scalars().first() if project: project.status = "failed" project.error_msg = str(e) await session.commit() publish_event(job_id, "job_failed", {"error": str(e)}) # RQ requires functions to be importable at module level (no inner/local functions) # This is the sync wrapper for GitHub audit jobs — called by RQ worker def run_github_audit_job(proj_id: str, job_id: str, repo: str, branch: str, access_token: str, file_type: str): """Module-level RQ job for GitHub-sourced audits. Picklable by RQ.""" from app.utils.github_client import get_repo_files files = get_repo_files(access_token, repo, branch) loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: loop.run_until_complete(run_audit_task(proj_id, job_id, files, file_type)) finally: loop.close() # Sync wrappers for RQ — each creates its own event loop to avoid cross-thread conflicts def audit_job_wrapper(project_id_str: str, job_id: str, files: list, file_type: str): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: loop.run_until_complete(run_audit_task(project_id_str, job_id, files, file_type)) finally: loop.close() def design_job_wrapper(project_id_str: str, job_id: str, idea_prompt: str): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: loop.run_until_complete(run_design_task(project_id_str, job_id, idea_prompt)) finally: loop.close()