96 lines
2.2 KiB
Python
Executable file
96 lines
2.2 KiB
Python
Executable file
from enum import Enum
|
|
from typing import Optional
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class RagStatus(str, Enum):
|
|
"""RAG status for agent reviews."""
|
|
RED = "Red"
|
|
AMBER = "Amber"
|
|
GREEN = "Green"
|
|
ERROR = "Error"
|
|
|
|
|
|
class OverallStatus(str, Enum):
|
|
"""Overall status for the proof review."""
|
|
PASSED = "Passed"
|
|
FAILED = "Failed"
|
|
ANALYSIS_ERROR = "Analysis Error"
|
|
REQUIRES_MANUAL_LEGAL_REVIEW = "Requires Manual Legal Review"
|
|
|
|
|
|
class SubReview(BaseModel):
|
|
"""Individual agent review result."""
|
|
ragStatus: RagStatus
|
|
feedback: str
|
|
issues: list[str]
|
|
isFinancialPromotion: Optional[bool] = None
|
|
financialPromotionReason: Optional[str] = None
|
|
|
|
class Config:
|
|
use_enum_values = True
|
|
|
|
|
|
class AgentReview(BaseModel):
|
|
"""Complete review from all agents."""
|
|
legalAgentReview: SubReview
|
|
brandAgentReview: SubReview
|
|
toneAgentReview: SubReview
|
|
channelAgentReview: SubReview
|
|
leadAgentSummary: str
|
|
overallStatus: OverallStatus
|
|
financialPromotionReason: Optional[str] = None
|
|
|
|
class Config:
|
|
use_enum_values = True
|
|
|
|
|
|
# WebSocket message types
|
|
class AnalyzeRequest(BaseModel):
|
|
"""Request to analyze a proof via WebSocket."""
|
|
type: str # Should be "analyze"
|
|
file_data: str # Base64 encoded file
|
|
file_type: str # MIME type
|
|
is_wip: bool = False
|
|
|
|
|
|
class AgentStartedMessage(BaseModel):
|
|
"""Message sent when an agent starts processing."""
|
|
type: str = "agent_started"
|
|
agent_name: str
|
|
|
|
|
|
class AgentCompletedMessage(BaseModel):
|
|
"""Message sent when an agent completes."""
|
|
type: str = "agent_completed"
|
|
agent_name: str
|
|
review: SubReview
|
|
|
|
class Config:
|
|
use_enum_values = True
|
|
|
|
|
|
class SummaryMessage(BaseModel):
|
|
"""Message sent when summary is ready."""
|
|
type: str = "summary"
|
|
lead_agent_summary: str
|
|
overall_status: OverallStatus
|
|
financial_promotion_reason: Optional[str] = None
|
|
|
|
class Config:
|
|
use_enum_values = True
|
|
|
|
|
|
class CompleteMessage(BaseModel):
|
|
"""Message sent when analysis is complete."""
|
|
type: str = "complete"
|
|
result: AgentReview
|
|
|
|
class Config:
|
|
use_enum_values = True
|
|
|
|
|
|
class ErrorMessage(BaseModel):
|
|
"""Message sent when an error occurs."""
|
|
type: str = "error"
|
|
message: str
|