#!/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 content""" summary: Optional[str] = None actors: Optional[str] = None key_details: Optional[str] = None class TreeNode(BaseModel): """Node in the document tree (folder or file)""" name: str type: str = Field(..., description="Type: 'folder' or 'file'") file_path: Optional[str] = None file_size: Optional[int] = None mime_type: Optional[str] = None created_at: Optional[str] = None analysis: Optional[DocumentAnalysis] = None children: Optional[List['TreeNode']] = None class DocumentsTree(BaseModel): """Root of the document tree""" name: str = "root" type: str = "folder" children: Optional[List[TreeNode]] = None class ChatRequest(BaseModel): """Chat request model""" userId: str = Field(..., description="Unique user identifier (UUID)") 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") documents_tree: Optional[DocumentsTree] = Field(default=None, description="Hierarchical tree of lawyer's documents") 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") userId: str = Field(..., description="Unique user identifier (UUID) for actors integration") 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") extracted_text: Optional[str] = Field(None, description="Extracted text content") 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") class DocCreatorRequest(BaseModel): """Document creator request model""" instruction: str = Field(..., description="User's instruction for document editing") documentContent: str = Field(..., description="HTML document content") contentFormat: str = Field(default="html", description="Format of document content (always 'html')") documents_tree: Optional[DocumentsTree] = Field(default=None, description="Hierarchical tree of documents for context") conversationHistory: Optional[List[Message]] = Field(default=[], description="Previous conversation messages") userId: str = Field(..., description="Unique user identifier (UUID)") documentId: str = Field(None, description="Document ID for live updates (UUID)") class DocCreatorResponse(BaseModel): """Document creator response model""" response: str = Field(..., description="Assistant's response") modifiedDocument: Optional[str] = Field(None, description="Modified HTML document (if changes were made)") processing_time: float = Field(..., description="Processing time in seconds") timestamp: str = Field(..., description="Response timestamp") error: Optional[str] = Field(None, description="Error message if any") class Actor(BaseModel): """Actor model representing a person or organization""" name: str = Field(..., description="Actor name") aliases: List[str] = Field(default=[], description="Alternative names for this actor") description: str = Field(default="", description="Description of the actor") implication: str = Field(default="", description="Legal or business implication of this actor")