c1030ee292
PDFs are now converted to PNG images at 200 DPI before being sent to Gemini for analysis. This fixes the unreliable iframe-based PDF preview and ensures all pages are properly analyzed. - Add PyMuPDF dependency for PDF rasterization - Create pdf_service.py with rasterize() and get_page_count() - Update agent interfaces to accept list of images for multi-page support - Add analyze_with_images() to Gemini service for multi-image analysis - Return rasterized PDF pages via WebSocket for frontend display - Add page navigation UI for multi-page PDFs in preview components Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
210 lines
8.7 KiB
Python
Executable file
210 lines
8.7 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
|
|
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, 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,
|
|
)
|
|
|
|
# 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
|
|
|
|
# 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)}"
|
|
})
|