structured_outputs/api_models.py

#!/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")