import base64 import logging import uuid from typing import Optional from fastapi import WebSocket from app.websocket.manager import ConnectionManager from app.services.analysis_service import AnalysisService from app.models.schemas import SubReview from app.models.database import async_session_factory from app.repositories import ProofRepository, CampaignRepository, UserRepository from app.services.storage_service import storage_service logger = logging.getLogger(__name__) async def handle_analyze_message( websocket: WebSocket, client_id: str, data: dict, manager: ConnectionManager, analysis_service: AnalysisService, ) -> None: """ Handle an 'analyze' message from the client. Runs the proof analysis and sends real-time updates via WebSocket. Args: websocket: The WebSocket connection client_id: Unique client identifier data: The message data containing file_data, file_type, is_wip manager: Connection manager for sending messages analysis_service: Service to run the analysis """ try: logger.info(f"[WEBSOCKET] Received analyze request from client: {client_id}") # Extract and decode the file data file_data_b64 = data.get("file_data", "") file_type = data.get("file_type", "image/png") is_wip = data.get("is_wip", False) logger.info(f"[WEBSOCKET] File type: {file_type}, is_wip: {is_wip}, base64 length: {len(file_data_b64)}") # Decode base64 file data try: file_data = base64.b64decode(file_data_b64) logger.info(f"[WEBSOCKET] Decoded file size: {len(file_data)} bytes") except Exception as e: logger.error(f"[WEBSOCKET] Failed to decode file data: {str(e)}") await manager.send_message(client_id, { "type": "error", "message": f"Failed to decode file data: {str(e)}" }) return # Create callback for real-time updates async def on_agent_update(agent_name: str, review: SubReview | None) -> None: if not manager.is_connected(client_id): logger.warning(f"[WEBSOCKET] Client {client_id} disconnected, skipping update") return if review is None: # Agent is starting logger.info(f"[WEBSOCKET] Sending agent_started: {agent_name}") await manager.send_message(client_id, { "type": "agent_started", "agent_name": agent_name }) else: # Agent completed logger.info(f"[WEBSOCKET] Sending agent_completed: {agent_name} - ragStatus: {review.ragStatus}") await manager.send_message(client_id, { "type": "agent_completed", "agent_name": agent_name, "review": { "ragStatus": review.ragStatus, "feedback": review.feedback, "issues": review.issues, "isFinancialPromotion": review.isFinancialPromotion, "financialPromotionReason": review.financialPromotionReason, } }) # Run the analysis logger.info("[WEBSOCKET] Starting analysis...") result = await analysis_service.analyze_proof( file_data=file_data, file_type=file_type, on_agent_update=on_agent_update, is_wip=is_wip, ) # Build the result dict result_dict = { "legalAgentReview": { "ragStatus": result.legalAgentReview.ragStatus, "feedback": result.legalAgentReview.feedback, "issues": result.legalAgentReview.issues, "isFinancialPromotion": result.legalAgentReview.isFinancialPromotion, "financialPromotionReason": result.legalAgentReview.financialPromotionReason, }, "brandAgentReview": { "ragStatus": result.brandAgentReview.ragStatus, "feedback": result.brandAgentReview.feedback, "issues": result.brandAgentReview.issues, }, "toneAgentReview": { "ragStatus": result.toneAgentReview.ragStatus, "feedback": result.toneAgentReview.feedback, "issues": result.toneAgentReview.issues, }, "channelAgentReview": { "ragStatus": result.channelAgentReview.ragStatus, "feedback": result.channelAgentReview.feedback, "issues": result.channelAgentReview.issues, }, "leadAgentSummary": result.leadAgentSummary, "overallStatus": result.overallStatus, "financialPromotionReason": result.financialPromotionReason, } # Persist to database if campaign info provided proof_id: Optional[str] = None version_id: Optional[str] = None campaign_id = data.get("campaign_id") proof_name = data.get("proof_name") if campaign_id and proof_name: try: logger.info(f"[WEBSOCKET] Persisting result for campaign {campaign_id}, proof {proof_name}") async with async_session_factory() as session: proof_repo = ProofRepository(session) # Store the file file_storage_key = await storage_service.store_file( file_data=file_data, campaign_id=uuid.UUID(campaign_id), proof_name=proof_name, version=1, # Will be updated by add_version_with_review file_type=file_type, ) # Generate thumbnail for small files thumbnail_url = None if len(file_data) < 500000: # < 500KB thumbnail_url = await storage_service.generate_thumbnail_data_url(file_data, file_type) # Save proof and version proof, version = await proof_repo.add_version_with_review( campaign_id=uuid.UUID(campaign_id), proof_name=proof_name, channel=data.get("channel"), sub_channel=data.get("sub_channel"), proof_type=data.get("proof_type"), file_storage_key=file_storage_key, thumbnail_url=thumbnail_url, agent_review=result_dict, overall_status=result.overallStatus, ) await session.commit() proof_id = str(proof.id) version_id = str(version.id) logger.info(f"[WEBSOCKET] Persisted proof {proof_id} version {version.version}") except Exception as e: logger.error(f"[WEBSOCKET] Failed to persist result: {str(e)}") # Continue - still send result to client even if persistence fails # Send the complete result logger.info(f"[WEBSOCKET] Analysis complete, sending result - overallStatus: {result.overallStatus}") if manager.is_connected(client_id): response = { "type": "complete", "result": result_dict, } # Include proof/version IDs if persisted if proof_id: response["proof_id"] = proof_id if version_id: response["version_id"] = version_id await manager.send_message(client_id, response) logger.info(f"[WEBSOCKET] Result sent to client: {client_id}") except Exception as e: # Send error message logger.error(f"[WEBSOCKET] Analysis failed for client {client_id}: {str(e)}") if manager.is_connected(client_id): await manager.send_message(client_id, { "type": "error", "message": f"Analysis failed: {str(e)}" })