Phase 2: domain models

openrag/core/models/__init__.py

@@ -0,0 +1,39 @@
+"""Domain models — pure Pydantic, no infrastructure imports."""
+
+from openrag.core.models.catalog import DocumentRecord, DocumentStatus, IndexationJob, JobStatus
+from openrag.core.models.chunk import Chunk, ChunkType
+from openrag.core.models.contextualization import ContextualizedQuery
+from openrag.core.models.conversation import Conversation, Message
+from openrag.core.models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock
+from openrag.core.models.prompt import Prompt, PromptType
+from openrag.core.models.query import RetrievalQuery
+from openrag.core.models.retrieval_response import RetrievalResponse
+from openrag.core.models.retrieval_result import RetrievalResult, ScoredChunk
+from openrag.core.models.user import OIDCSession, PartitionRole, User, UserPartition
+
+__all__ = [
+    "Chunk",
+    "ChunkType",
+    "ContextualizedQuery",
+    "Conversation",
+    "Document",
+    "DocumentRecord",
+    "DocumentStatus",
+    "DocumentType",
+    "ImageBlock",
+    "IndexationJob",
+    "JobStatus",
+    "Message",
+    "OIDCSession",
+    "PartitionRole",
+    "ProcessedDocument",
+    "Prompt",
+    "PromptType",
+    "RetrievalQuery",
+    "RetrievalResponse",
+    "RetrievalResult",
+    "ScoredChunk",
+    "TextBlock",
+    "User",
+    "UserPartition",
+]

openrag/core/models/catalog.py

@@ -0,0 +1,56 @@
+"""Catalog domain models — document records, indexation jobs, status tracking."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class DocumentStatus(str, Enum):
+    QUEUED = "QUEUED"
+    SERIALIZING = "SERIALIZING"
+    CHUNKING = "CHUNKING"
+    INSERTING = "INSERTING"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    CANCELLED = "CANCELLED"
+
+
+class JobStatus(str, Enum):
+    QUEUED = "QUEUED"
+    RUNNING = "RUNNING"
+    SUCCESS = "SUCCESS"
+    FAILED = "FAILED"
+    PARTIAL = "PARTIAL"
+
+
+class DocumentRecord(BaseModel):
+    """A document entry in the catalog (PostgreSQL)."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    file_id: str = ""
+    filename: str = ""
+    partition: str = "default"
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    status: DocumentStatus = DocumentStatus.QUEUED
+    error_message: str | None = None
+    created_by: int | None = None
+    relationship_id: str | None = None
+    parent_id: str | None = None
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+
+class IndexationJob(BaseModel):
+    """An indexation job tracking batch document processing."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    status: JobStatus = JobStatus.QUEUED
+    total_documents: int = 0
+    partition: str = "default"
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    started_at: datetime | None = None
+    completed_at: datetime | None = None

openrag/core/models/chunk.py

@@ -0,0 +1,72 @@
+"""Chunk — the unit of indexable and retrievable text."""
+
+from __future__ import annotations
+
+import uuid
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class ChunkType(str, Enum):
+    TEXT = "text"
+    IMAGE_CAPTION = "image_caption"
+    TABLE = "table"
+    CONTEXTUALIZED = "contextualized"
+
+
+class Chunk(BaseModel):
+    """A chunk of text extracted from a document, optionally embedded."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    document_id: str = ""
+    text: str = ""
+    chunk_index: int = 0
+    chunk_type: ChunkType = ChunkType.TEXT
+    embedding: list[float] | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    partition: str = "default"
+    page_number: int | None = None
+    token_count: int | None = None
+    header: str | None = None
+    context: str | None = None
+    content: str | None = None
+
+    def with_embedding(self, embedding: list[float]) -> Chunk:
+        """Return a copy with the embedding set."""
+        return self.model_copy(update={"embedding": embedding})
+
+    @classmethod
+    def from_langchain(cls, doc: Any) -> Chunk:
+        """Convert a LangChain Document to a Chunk.
+
+        Import is deferred to method body so core/ stays pure at import time.
+        """
+        metadata = dict(doc.metadata) if doc.metadata else {}
+        return cls(
+            id=metadata.pop("_id", str(uuid.uuid4())),
+            document_id=metadata.pop("file_id", ""),
+            text=doc.page_content,
+            partition=metadata.pop("partition", "default"),
+            page_number=metadata.pop("page", None),
+            chunk_type=ChunkType(metadata.pop("chunk_type", "text")),
+            metadata=metadata,
+        )
+
+    def to_langchain(self) -> Any:
+        """Convert back to a LangChain Document.
+
+        Import is deferred to method body so core/ stays pure at import time.
+        """
+        from langchain_core.documents.base import Document
+
+        metadata = {
+            **self.metadata,
+            "_id": self.id,
+            "file_id": self.document_id,
+            "partition": self.partition,
+            "page": self.page_number,
+            "chunk_type": self.chunk_type.value,
+        }
+        return Document(page_content=self.text, metadata=metadata)

openrag/core/models/contextualization.py

@@ -0,0 +1,19 @@
+"""Contextualized query models — query rewriting, HyDE, reasoning."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class ContextualizedQuery(BaseModel):
+    """A user query after contextualization by the LLM.
+
+    May include rewritten variants, a hypothetical document (HyDE),
+    and sub-queries for multi-step retrieval.
+    """
+
+    original: str = ""
+    query_list: list[str] = Field(default_factory=list)
+    intent: str = "qa"
+    hypothetical_doc: str | None = None
+    sub_queries: list[str] = Field(default_factory=list)

openrag/core/models/conversation.py

@@ -0,0 +1,32 @@
+"""Conversation and message domain models."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class Message(BaseModel):
+    """A single message within a conversation."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    conversation_id: str = ""
+    role: str = "user"
+    content: str = ""
+    sources_json: list[dict[str, Any]] | None = None
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+
+class Conversation(BaseModel):
+    """A persistent conversation between a user and the RAG system."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    user_id: int = 0
+    partition_scope: str = "default"
+    title: str = ""
+    messages: list[Message] = Field(default_factory=list)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

openrag/core/models/document.py

@@ -0,0 +1,112 @@
+"""Document — the input to the indexing pipeline."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class DocumentType(str, Enum):
+    PDF = "pdf"
+    TEXT = "text"
+    HTML = "html"
+    MARKDOWN = "markdown"
+    IMAGE = "image"
+    AUDIO = "audio"
+    VIDEO = "video"
+    DOCX = "docx"
+    PPTX = "pptx"
+    DOC = "doc"
+    EML = "eml"
+
+
+class TextBlock(BaseModel):
+    """A block of text extracted from a document."""
+
+    text: str
+    page_number: int | None = None
+    block_type: str = "paragraph"
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class ImageBlock(BaseModel):
+    """An image extracted from a document."""
+
+    image_bytes: bytes = Field(exclude=True, repr=False)
+    page_number: int | None = None
+    caption: str | None = None
+    mime_type: str = "image/png"
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class Document(BaseModel):
+    """A document before or during indexing."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    filename: str = ""
+    content_type: DocumentType = DocumentType.TEXT
+    text: str | None = None
+    raw_bytes: bytes | None = Field(None, exclude=True)
+    partition: str = "default"
+    tags: list[str] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    @staticmethod
+    def detect_content_type(filename: str) -> DocumentType:
+        """Detect DocumentType from filename extension."""
+        ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
+        mapping = {
+            "pdf": DocumentType.PDF,
+            "txt": DocumentType.TEXT,
+            "md": DocumentType.MARKDOWN,
+            "html": DocumentType.HTML,
+            "htm": DocumentType.HTML,
+            "png": DocumentType.IMAGE,
+            "jpg": DocumentType.IMAGE,
+            "jpeg": DocumentType.IMAGE,
+            "mp3": DocumentType.AUDIO,
+            "wav": DocumentType.AUDIO,
+            "mp4": DocumentType.VIDEO,
+            "docx": DocumentType.DOCX,
+            "pptx": DocumentType.PPTX,
+            "doc": DocumentType.DOC,
+            "eml": DocumentType.EML,
+        }
+        return mapping.get(ext, DocumentType.TEXT)
+
+    @classmethod
+    def from_langchain(cls, doc: Any) -> Document:
+        """Convert a LangChain Document to a domain Document."""
+        metadata = dict(doc.metadata) if doc.metadata else {}
+        return cls(
+            filename=metadata.pop("source", ""),
+            text=doc.page_content,
+            partition=metadata.pop("partition", "default"),
+            metadata=metadata,
+        )
+
+    def to_langchain(self) -> Any:
+        """Convert back to a LangChain Document."""
+        from langchain_core.documents.base import Document as LCDocument
+
+        metadata = {
+            **self.metadata,
+            "source": self.filename,
+            "partition": self.partition,
+        }
+        return LCDocument(page_content=self.text or "", metadata=metadata)
+
+
+class ProcessedDocument(BaseModel):
+    """Document after parsing/extraction — contains text blocks and images."""
+
+    document_id: str = ""
+    text_blocks: list[TextBlock] = Field(default_factory=list)
+    images: list[ImageBlock] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    page_count: int = 0

openrag/core/models/prompt.py

@@ -0,0 +1,31 @@
+"""Prompt domain models."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class PromptType(str, Enum):
+    SYS_PROMPT = "sys_prompt"
+    QUERY_CONTEXTUALIZER = "query_contextualizer"
+    CHUNK_CONTEXTUALIZER = "chunk_contextualizer"
+    IMAGE_CAPTIONING = "image_captioning"
+    HYDE = "hyde"
+    MULTI_QUERY = "multi_query"
+    SPOKEN_STYLE_ANSWER = "spoken_style_answer"
+
+
+class Prompt(BaseModel):
+    """A prompt template stored in the library."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    prompt_type: str = ""
+    name: str = ""
+    content: str = ""
+    is_default: bool = False
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

openrag/core/models/query.py

@@ -0,0 +1,23 @@
+"""Retrieval query domain model."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class RetrievalQuery(BaseModel):
+    """A user query with retrieval parameters."""
+
+    text: str
+    partition: str = "default"
+    top_k: int = 10
+    similarity_threshold: float = 0.95
+    filters: dict[str, Any] = Field(default_factory=dict)
+    include_related: bool = False
+    include_ancestors: bool = False
+    related_limit: int = 10
+    max_ancestor_depth: int | None = None
+    with_surrounding_chunks: bool = True
+    rerank: bool = True