Charles Grandjean
update api_models.py
892ca64
#!/usr/bin/env python3
"""
Structured outputs for Agent API
"""
from typing import List, Optional
from pydantic import BaseModel, Field
class Message(BaseModel):
"""Chat message"""
role: str = Field(..., description="Role: 'user' or 'assistant'")
content: str = Field(..., description="Message content")
class LawyerProfile(BaseModel):
"""Lawyer's professional profile"""
full_name: Optional[str] = Field(None, description="Lawyer's full name")
primary_specialty: Optional[str] = Field(None, description="Primary legal specialty")
legal_specialties: Optional[List[str]] = Field(None, description="List of legal specialties")
jurisdiction: Optional[str] = Field(None, description="Primary jurisdiction of practice")
experience_level: Optional[str] = Field(None, description="Experience level (e.g., Junior, Mid-level, Senior, Partner)")
languages: Optional[List[str]] = Field(None, description="Languages spoken")
lawyer_description: Optional[str] = Field(None, description="Professional description or bio")
class DocumentAnalysis(BaseModel):
"""Document analysis result"""
file_name: str
summary: Optional[str] = None
actors: Optional[str] = None
key_details: Optional[str] = None
class ChatRequest(BaseModel):
"""Chat request model"""
clientId: str = Field(..., description="Unique client identifier")
message: str = Field(..., description="User's question")
conversationHistory: Optional[List[Message]] = Field(default=[], description="Previous conversation messages")
userType: Optional[str] = Field(default="client", description="User type: 'client' for general users or 'lawyer' for legal professionals")
jurisdiction: Optional[str] = Field(default="Romania", description="Jurisdiction of the user")
documentAnalyses: Optional[List[DocumentAnalysis]] = Field(default=None, description="Lawyer's document analyses")
lawyerProfile: Optional[LawyerProfile] = Field(default=None, description="Lawyer's professional profile")
class ChatResponse(BaseModel):
"""Chat response model"""
response: str = Field(..., description="Assistant's response")
processing_time: float = Field(..., description="Processing time in seconds")
references: List[str] = Field(default=[], description="Referenced documents")
timestamp: str = Field(..., description="Response timestamp")
error: Optional[str] = Field(None, description="Error message if any")
class HealthResponse(BaseModel):
"""Health check response model"""
status: str = Field(..., description="Health status")
agent_ready: bool = Field(..., description="Whether agent is ready")
lightrag_healthy: bool = Field(..., description="Whether LightRAG is healthy")
timestamp: str = Field(..., description="Health check timestamp")
class AnalyzePDFRequest(BaseModel):
"""PDF analysis request model"""
pdf_content: str = Field(..., description="Base64 encoded document content (PDF or image)")
filename: Optional[str] = Field(default="document.pdf", description="Original filename")
class AnalyzePDFResponse(BaseModel):
"""PDF analysis response model"""
actors: str = Field(..., description="Extracted actors")
key_details: str = Field(..., description="Key details extracted")
summary: str = Field(..., description="High-level summary")
processing_status: str = Field(..., description="Processing status")
processing_time: float = Field(..., description="Processing time in seconds")
timestamp: str = Field(..., description="Analysis timestamp")
error: Optional[str] = Field(None, description="Error message if any")