nvtitan/graphRAG

1"""2Data models for GraphLLM system following the manual specifications3"""4from pydantic import BaseModel, Field5from typing import Optional, List, Dict, Any, Literal6from datetime import datetime7from enum import Enum8import uuid9 10 11# Enums12class ChunkType(str, Enum):13    """Types of chunks extracted from PDF"""14    PARAGRAPH = "paragraph"15    CODE = "code"16    TABLE = "table"17    IMAGE = "image"18    IMAGE_TEXT = "image_text"19 20 21class NodeType(str, Enum):22    """Types of graph nodes"""23    CONCEPT = "concept"24    PERSON = "person"25    METHOD = "method"26    TERM = "term"27    CLASS = "class"28    FUNCTION = "function"29    ENTITY = "entity"30 31 32class RelationType(str, Enum):33    """Canonical relation types for edges"""34    IS_A = "is_a"35    PART_OF = "part_of"36    METHOD_OF = "method_of"37    CAUSES = "causes"38    USES = "uses"39    RELATED_TO = "related_to"40    DEFINED_AS = "defined_as"41    DEPENDS_ON = "depends_on"42    IMPLEMENTS = "implements"43    SIMILAR_TO = "similar_to"44    OBSERVES = "observes"45    MEASURES = "measures"46    PRODUCES = "produces"47    CONTAINS = "contains"48    AFFECTS = "affects"49    ENABLES = "enables"50    REQUIRES = "requires"51    INTERACTS_WITH = "interacts_with"52    ENRICHES = "enriches"53    ENHANCES = "enhances"54    SUPPORTS = "supports"55    DESCRIBES = "describes"56    EXPLAINS = "explains"57    REFERS_TO = "refers_to"58    ASSOCIATED_WITH = "associated_with"59 60 61# Core Data Models62 63class Chunk(BaseModel):64    """Individual chunk of text/content from PDF"""65    chunk_id: str = Field(default_factory=lambda: str(uuid.uuid4()))66    pdf_id: str67    page_number: int68    char_range: tuple[int, int]69    type: ChunkType70    text: str71    table_json: Optional[Dict[str, Any]] = None72    image_id: Optional[str] = None73    metadata: Dict[str, Any] = Field(default_factory=dict)74    created_at: datetime = Field(default_factory=datetime.utcnow)75 76 77class EmbeddingEntry(BaseModel):78    """Vector embedding for a chunk"""79    chunk_id: str80    embedding: List[float]81    created_at: datetime = Field(default_factory=datetime.utcnow)82    metadata: Dict[str, Any] = Field(default_factory=dict)83 84 85class SupportingChunk(BaseModel):86    """Reference to a chunk supporting a node or edge"""87    chunk_id: str88    score: float89    page_number: Optional[int] = None90    snippet: Optional[str] = None91 92 93class GraphNode(BaseModel):94    """Node in the knowledge graph"""95    node_id: str = Field(default_factory=lambda: str(uuid.uuid4()))96    label: str97    type: NodeType98    aliases: List[str] = Field(default_factory=list)99    supporting_chunks: List[SupportingChunk] = Field(default_factory=list)100    importance_score: float = 0.0101    metadata: Dict[str, Any] = Field(default_factory=dict)102    created_at: datetime = Field(default_factory=datetime.utcnow)103 104 105class GraphEdge(BaseModel):106    """Edge in the knowledge graph"""107    edge_id: str = Field(default_factory=lambda: str(uuid.uuid4()))108    from_node: str = Field(alias="from")109    to_node: str = Field(alias="to")110    relation: RelationType111    confidence: float112    supporting_chunks: List[SupportingChunk] = Field(default_factory=list)113    metadata: Dict[str, Any] = Field(default_factory=dict)114    created_at: datetime = Field(default_factory=datetime.utcnow)115 116    class Config:117        populate_by_name = True118        # FastAPI automatically serializes enums as their string values in JSON119 120 121class Triple(BaseModel):122    """Extracted triple from text"""123    subject: str124    predicate: str125    object: str126    confidence: float = 1.0127    source_chunk_id: Optional[str] = None128    page_number: Optional[int] = None129    justification: Optional[str] = None130 131 132class CanonicalTriple(BaseModel):133    """LLM-canonicalized triple"""134    subject_label: str135    object_label: str136    relation: RelationType137    confidence: float138    justification: str139    page_number: int140 141 142# API Request/Response Models143 144class UploadResponse(BaseModel):145    """Response from PDF upload"""146    pdf_id: str147    filename: str148    status: str149    message: str150    num_pages: Optional[int] = None151    num_chunks: Optional[int] = None152 153 154class GraphResponse(BaseModel):155    """Response containing graph data"""156    nodes: List[GraphNode]157    edges: List[GraphEdge]158    metadata: Dict[str, Any] = Field(default_factory=dict)159 160 161class SourceCitation(BaseModel):162    """Source citation with page number and snippet"""163    page_number: int164    snippet: str165    chunk_id: str166    score: Optional[float] = None167 168 169class NodeDetailResponse(BaseModel):170    """Response for node detail request"""171    node_id: str172    label: str173    type: NodeType174    summary: str175    sources: List[SourceCitation]176    related_nodes: List[Dict[str, Any]] = Field(default_factory=list)177    raw_chunks: Optional[List[Chunk]] = None178 179 180class ChatMessage(BaseModel):181    """Chat message"""182    role: Literal["user", "assistant", "system"]183    content: str184    sources: Optional[List[SourceCitation]] = None185    timestamp: datetime = Field(default_factory=datetime.utcnow)186 187 188class ChatRequest(BaseModel):189    """Chat request"""190    query: str191    pdf_id: str192    include_citations: bool = True193    max_sources: int = 5194 195 196class ChatResponse(BaseModel):197    """Chat response with answer and citations"""198    answer: str199    sources: List[SourceCitation]200    context_chunks: Optional[List[str]] = None201 202 203class PDFMetadata(BaseModel):204    """Metadata for uploaded PDF"""205    pdf_id: str = Field(default_factory=lambda: str(uuid.uuid4()))206    filename: str207    filepath: str208    num_pages: int209    file_size_bytes: int210    upload_timestamp: datetime = Field(default_factory=datetime.utcnow)211    processing_status: str = "pending"212    num_chunks: int = 0213    num_nodes: int = 0214    num_edges: int = 0215    metadata: Dict[str, Any] = Field(default_factory=dict)216 217 218class IngestionLog(BaseModel):219    """Log entry for ingestion process"""220    log_id: str = Field(default_factory=lambda: str(uuid.uuid4()))221    pdf_id: str222    timestamp: datetime = Field(default_factory=datetime.utcnow)223    stage: str224    status: str225    message: str226    details: Optional[Dict[str, Any]] = None227 228 229class AdminStatus(BaseModel):230    """Admin status response"""231    total_pdfs: int232    total_chunks: int233    total_nodes: int234    total_edges: int235    vector_index_size: int236    recent_logs: List[IngestionLog]237