doc_rag / models.py
TytonTerrapin's picture
Upload 13 files
68fe04b verified
Raw
History Blame Contribute Delete
4.49 kB
"""
models.py — Pydantic request/response schemas for Space 2.
Changes vs previous version:
- ChunkSchema: added `outline_path` and `source_block_id` fields so the
enriched payload produced by the new serialiser.py round-trips cleanly.
Both are Optional with defaults so older payloads (fields absent) still
validate without error.
- IngestPayload: added all enriched fields that serialiser.py now emits:
block_to_chunks { block_id: [chunk_id, ...] }
outline nested section tree
chunk_outline_path { chunk_id: [h1_title, h2_title, ...] }
content_type_index { "paragraph": [...], "table": [...], "figure": [...] }
page_index { "1": [block_id, ...], "2": [...] }
All are Optional[...] = None so old single-field payloads keep working.
`graph` is also Optional — serialiser.py supports include_graph=False for
lightweight payloads.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
# ── Ingest (Space 1 → Space 2) ────────────────────────────────────────────────
class ChunkSchema(BaseModel):
chunk_id: str
doc_name: str
page_num: int
region_type: str
text: str
section_title: str
bbox: List[float]
confidence: float
char_count: int
table_html: Optional[str] = None
figure_path: Optional[str] = None
source_block_id: Optional[str] = None
outline_path: Optional[List[str]] = None # [h1_title, h2_title, ...]
# Space 1 FIX #1: carry-forward context is now a separate field, NOT in text.
# chunk.text is always clean; context is the previous sentence for
# display / re-ranking only. Never embed this field.
context: Optional[str] = None
class IngestPayload(BaseModel):
doc_name: str
total_pages: int
reading_order: List[str]
chunks: List[ChunkSchema]
# Node-link graph — Optional because serialiser supports include_graph=False
graph: Optional[Dict[str, Any]] = None
# Enriched indexes added by the new serialiser.py (Optional for back-compat)
block_to_chunks: Optional[Dict[str, List[str]]] = None
outline: Optional[List[Dict[str, Any]]] = None
chunk_outline_path: Optional[Dict[str, List[str]]] = None
content_type_index: Optional[Dict[str, List[str]]] = None
page_index: Optional[Dict[str, List[str]]] = None
class IngestResponse(BaseModel):
doc_name: str
chunks_indexed: int
message: str
# ── Query (caller → Space 2) ──────────────────────────────────────────────────
class QueryRequest(BaseModel):
question: str = Field(..., min_length=3)
doc_name: Optional[str] = Field(
None,
description="Restrict retrieval to a single ingested document. "
"If omitted, all documents are searched.",
)
top_k: Optional[int] = Field(
None,
description="Override the default RETRIEVAL_TOP_K for this request.",
ge=1, le=20,
)
class RetrievedChunk(BaseModel):
chunk_id: str
doc_name: str
page_num: int
region_type: str
section_title: str
text: str
score: float # final score after graph re-ranking
table_html: Optional[str] = None
outline_path: Optional[List[str]] = None # full breadcrumb for generator context headers
class RAGResponse(BaseModel):
question: str
answer: str
retrieved_chunks: List[RetrievedChunk]
doc_names_searched: List[str]
# ── Health ────────────────────────────────────────────────────────────────────
class HealthResponse(BaseModel):
status: str
docs_ingested: int
embedder_ready: bool
llm_ready: bool