308 lines
14 KiB
Python
Executable file
308 lines
14 KiB
Python
Executable file
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, AuditRepository
|
|
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,
|
|
current_user_id: Optional[uuid.UUID] = None,
|
|
) -> 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
|
|
|
|
# Compute file hash for duplicate detection
|
|
file_hash = storage_service.get_checksum(file_data)
|
|
logger.info(f"[WEBSOCKET] Computed file hash: {file_hash}")
|
|
|
|
# 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,
|
|
"resolvedIssues": review.resolvedIssues,
|
|
"outstandingIssues": review.outstandingIssues,
|
|
"newIssues": review.newIssues,
|
|
}
|
|
})
|
|
|
|
# Extract brand selection for brand agent
|
|
brand = data.get("brand", "Barclaycard") # Default to Barclaycard if not specified
|
|
logger.info(f"[WEBSOCKET] Brand selection: {brand}")
|
|
|
|
# Fetch previous analysis if this is a revision
|
|
previous_analysis = None
|
|
previous_file_hash = None
|
|
is_identical_file = False
|
|
campaign_id = data.get("campaign_id")
|
|
proof_name = data.get("proof_name")
|
|
|
|
if campaign_id and proof_name:
|
|
try:
|
|
logger.info(f"[WEBSOCKET] Checking for previous analysis - campaign: {campaign_id}, proof: {proof_name}")
|
|
async with async_session_factory() as session:
|
|
proof_repo = ProofRepository(session)
|
|
existing_proof = await proof_repo.get_by_campaign_and_name(
|
|
uuid.UUID(campaign_id), proof_name
|
|
)
|
|
if existing_proof:
|
|
previous_analysis = await proof_repo.get_latest_version_review(existing_proof.id)
|
|
previous_file_hash = await proof_repo.get_latest_version_hash(existing_proof.id)
|
|
if previous_analysis:
|
|
logger.info(f"[WEBSOCKET] Found previous analysis version {previous_analysis.get('version')}")
|
|
else:
|
|
logger.info("[WEBSOCKET] No previous analysis found (new proof)")
|
|
# Check if file is identical to previous version
|
|
if previous_file_hash and previous_file_hash == file_hash:
|
|
is_identical_file = True
|
|
logger.info(f"[WEBSOCKET] Identical file detected - hash matches previous version: {file_hash}")
|
|
else:
|
|
logger.info("[WEBSOCKET] No existing proof found (new proof)")
|
|
except Exception as e:
|
|
logger.warning(f"[WEBSOCKET] Failed to fetch previous analysis: {str(e)}")
|
|
# Continue without previous analysis - still run the current analysis
|
|
|
|
# Extract proof metadata for agent context
|
|
channel = data.get("channel")
|
|
sub_channel = data.get("sub_channel")
|
|
proof_type = data.get("proof_type")
|
|
|
|
# Build a once-only callback that notifies the client when the primary
|
|
# Gemini model is unavailable and the fallback model is used instead.
|
|
fallback_notified = False
|
|
|
|
async def on_model_fallback() -> None:
|
|
nonlocal fallback_notified
|
|
if fallback_notified:
|
|
return
|
|
fallback_notified = True
|
|
logger.info(f"[WEBSOCKET] Sending model_fallback notification to client: {client_id}")
|
|
if manager.is_connected(client_id):
|
|
await manager.send_message(client_id, {"type": "model_fallback"})
|
|
else:
|
|
logger.warning(f"[WEBSOCKET] Cannot send model_fallback - client {client_id} not connected")
|
|
|
|
# Run the analysis
|
|
logger.info("[WEBSOCKET] Starting analysis...")
|
|
result, pdf_pages = await analysis_service.analyze_proof(
|
|
file_data=file_data,
|
|
file_type=file_type,
|
|
on_agent_update=on_agent_update,
|
|
is_wip=is_wip,
|
|
brand=brand,
|
|
previous_analysis=previous_analysis,
|
|
channel=channel,
|
|
sub_channel=sub_channel,
|
|
proof_type=proof_type,
|
|
on_fallback=on_model_fallback,
|
|
)
|
|
|
|
# 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,
|
|
"resolvedIssues": result.legalAgentReview.resolvedIssues,
|
|
"outstandingIssues": result.legalAgentReview.outstandingIssues,
|
|
"newIssues": result.legalAgentReview.newIssues,
|
|
},
|
|
"brandAgentReview": {
|
|
"ragStatus": result.brandAgentReview.ragStatus,
|
|
"feedback": result.brandAgentReview.feedback,
|
|
"issues": result.brandAgentReview.issues,
|
|
"resolvedIssues": result.brandAgentReview.resolvedIssues,
|
|
"outstandingIssues": result.brandAgentReview.outstandingIssues,
|
|
"newIssues": result.brandAgentReview.newIssues,
|
|
},
|
|
"channelBestPracticesAgentReview": {
|
|
"ragStatus": result.channelBestPracticesAgentReview.ragStatus,
|
|
"feedback": result.channelBestPracticesAgentReview.feedback,
|
|
"issues": result.channelBestPracticesAgentReview.issues,
|
|
"resolvedIssues": result.channelBestPracticesAgentReview.resolvedIssues,
|
|
"outstandingIssues": result.channelBestPracticesAgentReview.outstandingIssues,
|
|
"newIssues": result.channelBestPracticesAgentReview.newIssues,
|
|
},
|
|
"channelTechSpecsAgentReview": {
|
|
"ragStatus": result.channelTechSpecsAgentReview.ragStatus,
|
|
"feedback": result.channelTechSpecsAgentReview.feedback,
|
|
"issues": result.channelTechSpecsAgentReview.issues,
|
|
"resolvedIssues": result.channelTechSpecsAgentReview.resolvedIssues,
|
|
"outstandingIssues": result.channelTechSpecsAgentReview.outstandingIssues,
|
|
"newIssues": result.channelTechSpecsAgentReview.newIssues,
|
|
},
|
|
"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
|
|
|
|
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 images up to 10MB (data URLs are ~33% larger than binary)
|
|
thumbnail_url = None
|
|
if len(file_data) < 10000000:
|
|
if file_type.startswith('image/'):
|
|
thumbnail_url = await storage_service.generate_thumbnail_data_url(file_data, file_type)
|
|
elif file_type == 'application/pdf' and pdf_pages:
|
|
# Use first rasterized page as thumbnail
|
|
first_page_data, _, _ = pdf_pages[0]
|
|
thumbnail_url = await storage_service.generate_thumbnail_data_url(first_page_data, 'image/png')
|
|
|
|
# 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,
|
|
file_hash=file_hash,
|
|
is_identical_file=is_identical_file,
|
|
created_by=current_user_id,
|
|
)
|
|
|
|
# Auto-create ErrorItem when analysis results in an error
|
|
if result.overallStatus == "Analysis Error":
|
|
audit_repo = AuditRepository(session)
|
|
await audit_repo.create_error_item(
|
|
proof_version_id=version.id,
|
|
error_summary=result.leadAgentSummary,
|
|
)
|
|
logger.info(f"[WEBSOCKET] Created ErrorItem for analysis error on version {version.id}")
|
|
|
|
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,
|
|
"is_identical_file": is_identical_file,
|
|
}
|
|
# Include proof/version IDs if persisted
|
|
if proof_id:
|
|
response["proof_id"] = proof_id
|
|
if version_id:
|
|
response["version_id"] = version_id
|
|
|
|
# Include rasterized PDF pages if present
|
|
if pdf_pages:
|
|
import base64 as b64_module
|
|
response["pdf_pages"] = [
|
|
{
|
|
"page": i + 1,
|
|
"data_url": f"data:image/png;base64,{b64_module.b64encode(png_data).decode('utf-8')}",
|
|
"width": width,
|
|
"height": height,
|
|
}
|
|
for i, (png_data, width, height) in enumerate(pdf_pages)
|
|
]
|
|
logger.info(f"[WEBSOCKET] Including {len(pdf_pages)} rasterized PDF pages in response")
|
|
|
|
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)}"
|
|
})
|