danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Satisfy ty by adding a typed to_pydantic helper

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-23 17:03:24 +03:00
parent b6ea07d2de
commit a74c3a6f9e
No known key found for this signature in database
4 changed files with 98 additions and 83 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -5,7 +5,7 @@ from datetime import datetime, timedelta
from enum import Enum from enum import Enum
from importlib import metadata from importlib import metadata
from pathlib import Path from pathlib import Path
from typing import Any from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4 from uuid import uuid4
import lancedb import lancedb
@ -18,9 +18,25 @@ from haiku.rag.config import AppConfig, Config
from haiku.rag.embeddings import get_embedder from haiku.rag.embeddings import get_embedder
from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
if TYPE_CHECKING:
from lancedb.query import AsyncQueryBase
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
async def query_to_pydantic[T: LanceModel](
query: "AsyncQueryBase", model: type[T]
) -> list[T]:
"""Typed wrapper around AsyncQueryBase.to_pydantic.
The upstream stub annotates `.to_pydantic()` as returning `list[LanceModel]`
regardless of the concrete model passed in. This helper narrows the return
type to the concrete model so attribute access on the results type-checks
at call sites without needing per-line cast / ignore comments.
"""
return cast("list[T]", await query.to_pydantic(model))
class ConnectionMode(Enum): class ConnectionMode(Enum):
LOCAL = "local" LOCAL = "local"
CLOUD = "cloud" CLOUD = "cloud"
@ -92,19 +108,26 @@ def get_documents_arrow_schema() -> pa.Schema:
return pa.schema(fields) return pa.schema(fields)
def create_chunk_model(vector_dim: int): class ChunkRecordBase(LanceModel):
"""Create a ChunkRecord model with the specified vector dimension. """Static base for ChunkRecord — declares the fields so attribute access
and constructor calls type-check. The concrete `vector` field is overridden
This creates a model with proper vector typing for LanceDB. by create_chunk_model() with a Vector(dim) whose fixed-size-list dimension
is only known at runtime.
""" """
class ChunkRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4())) id: str = Field(default_factory=lambda: str(uuid4()))
document_id: str document_id: str
content: str content: str
content_fts: str = Field(default="") content_fts: str = Field(default="")
metadata: str = Field(default="{}") metadata: str = Field(default="{}")
order: int = Field(default=0) order: int = Field(default=0)
vector: list[float] = Field(default_factory=list)
def create_chunk_model(vector_dim: int) -> type[ChunkRecordBase]:
"""Create a ChunkRecord model with the specified vector dimension."""
class ChunkRecord(ChunkRecordBase):
vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore
return ChunkRecord return ChunkRecord
@ -230,7 +253,7 @@ class Store:
# Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB) # Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB)
chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
self.ChunkRecord = create_chunk_model(chunk_vector_dim) self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim)
# Initialize tables (creates them if they don't exist) # Initialize tables (creates them if they don't exist)
await self._init_tables() await self._init_tables()
@ -506,8 +529,8 @@ class Store:
async def get_haiku_version(self) -> str: async def get_haiku_version(self) -> str:
"""Returns the user version stored in settings.""" """Returns the user version stored in settings."""
settings_records: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] settings_records = await query_to_pydantic(
self.settings_table.query().limit(1).to_pydantic(SettingsRecord) self.settings_table.query().limit(1), SettingsRecord
) )
if settings_records: if settings_records:
settings = ( settings = (
@ -525,8 +548,8 @@ class Store:
ReadOnlyError: If the store is in read-only mode. ReadOnlyError: If the store is in read-only mode.
""" """
self._assert_writable() self._assert_writable()
settings_records: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] settings_records = await query_to_pydantic(
self.settings_table.query().limit(1).to_pydantic(SettingsRecord) self.settings_table.query().limit(1), SettingsRecord
) )
if settings_records: if settings_records:
# Only write if version actually changes to avoid creating new table versions # Only write if version actually changes to avoid creating new table versions

61
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -10,7 +10,7 @@ if TYPE_CHECKING:
from lancedb.index import FTS from lancedb.index import FTS
from lancedb.rerankers import RRFReranker from lancedb.rerankers import RRFReranker
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store, query_to_pydantic
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -89,6 +89,7 @@ class ChunkRepository:
chunk_id = str(uuid4()) chunk_id = str(uuid4())
assert chunk.document_id is not None assert chunk.document_id is not None
assert chunk.embedding is not None
chunk_record = self.store.ChunkRecord( chunk_record = self.store.ChunkRecord(
id=chunk_id, id=chunk_id,
document_id=chunk.document_id, document_id=chunk.document_id,
@ -110,24 +111,22 @@ class ChunkRepository:
async def get_by_id(self, entity_id: str) -> Chunk | None: async def get_by_id(self, entity_id: str) -> Chunk | None:
"""Get a chunk by its ID.""" """Get a chunk by its ID."""
results = await ( results = await query_to_pydantic(
self.store.chunks_table.query() self.store.chunks_table.query().where(f"id = '{entity_id}'").limit(1),
.where(f"id = '{entity_id}'") self.store.ChunkRecord,
.limit(1)
.to_pydantic(self.store.ChunkRecord)
) )
if not results: if not results:
return None return None
chunk_record = results[0] chunk_record = results[0]
md = json.loads(chunk_record.metadata) # ty: ignore[unresolved-attribute] md = json.loads(chunk_record.metadata)
return Chunk( return Chunk(
id=chunk_record.id, # ty: ignore[unresolved-attribute] id=chunk_record.id,
document_id=chunk_record.document_id, # ty: ignore[unresolved-attribute] document_id=chunk_record.document_id,
content=chunk_record.content, # ty: ignore[unresolved-attribute] content=chunk_record.content,
metadata=md, metadata=md,
order=chunk_record.order, # ty: ignore[unresolved-attribute] order=chunk_record.order,
) )
async def update(self, entity: Chunk) -> Chunk: async def update(self, entity: Chunk) -> Chunk:
@ -175,18 +174,18 @@ class ChunkRepository:
if limit is not None: if limit is not None:
query = query.limit(limit) query = query.limit(limit)
results = await query.to_pydantic(self.store.ChunkRecord) results = await query_to_pydantic(query, self.store.ChunkRecord)
chunks: list[Chunk] = [] chunks: list[Chunk] = []
for rec in results: for rec in results:
md = json.loads(rec.metadata) # ty: ignore[unresolved-attribute] md = json.loads(rec.metadata)
chunks.append( chunks.append(
Chunk( Chunk(
id=rec.id, # ty: ignore[unresolved-attribute] id=rec.id,
document_id=rec.document_id, # ty: ignore[unresolved-attribute] document_id=rec.document_id,
content=rec.content, # ty: ignore[unresolved-attribute] content=rec.content,
metadata=md, metadata=md,
order=rec.order, # ty: ignore[unresolved-attribute] order=rec.order,
) )
) )
return chunks return chunks
@ -316,7 +315,7 @@ class ChunkRepository:
if limit is not None: if limit is not None:
query = query.limit(limit) query = query.limit(limit)
results = await query.to_pydantic(self.store.ChunkRecord) results = await query_to_pydantic(query, self.store.ChunkRecord)
# Get document info (only metadata columns, skip content/docling blobs) # Get document info (only metadata columns, skip content/docling blobs)
doc_rows = await ( doc_rows = await (
@ -333,14 +332,14 @@ class ChunkRepository:
chunks: list[Chunk] = [] chunks: list[Chunk] = []
for rec in results: for rec in results:
md = json.loads(rec.metadata) # ty: ignore[unresolved-attribute] md = json.loads(rec.metadata)
chunks.append( chunks.append(
Chunk( Chunk(
id=rec.id, # ty: ignore[unresolved-attribute] id=rec.id,
document_id=rec.document_id, # ty: ignore[unresolved-attribute] document_id=rec.document_id,
content=rec.content, # ty: ignore[unresolved-attribute] content=rec.content,
metadata=md, metadata=md,
order=rec.order, # ty: ignore[unresolved-attribute] order=rec.order,
document_uri=doc_uri, document_uri=doc_uri,
document_title=doc_title, document_title=doc_title,
document_meta=json.loads(doc_meta), document_meta=json.loads(doc_meta),
@ -378,18 +377,16 @@ class ChunkRepository:
f" AND `order` >= {min_order}" f" AND `order` >= {min_order}"
f" AND `order` <= {max_order}" f" AND `order` <= {max_order}"
) )
results = await ( results = await query_to_pydantic(
self.store.chunks_table.query() self.store.chunks_table.query().where(where), self.store.ChunkRecord
.where(where)
.to_pydantic(self.store.ChunkRecord)
) )
return [ return [
Chunk( Chunk(
id=rec.id, # ty: ignore[unresolved-attribute] id=rec.id,
document_id=rec.document_id, # ty: ignore[unresolved-attribute] document_id=rec.document_id,
content=rec.content, # ty: ignore[unresolved-attribute] content=rec.content,
metadata=json.loads(rec.metadata), # ty: ignore[unresolved-attribute] metadata=json.loads(rec.metadata),
order=rec.order, # ty: ignore[unresolved-attribute] order=rec.order,
) )
for rec in results for rec in results
] ]

31
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -4,7 +4,12 @@ from uuid import uuid4
from lancedb.index import BTree from lancedb.index import BTree
from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema from haiku.rag.store.engine import (
DocumentRecord,
Store,
get_documents_arrow_schema,
query_to_pydantic,
)
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
from haiku.rag.utils import escape_sql_string from haiku.rag.utils import escape_sql_string
@ -90,11 +95,9 @@ class DocumentRepository:
async def get_by_id(self, entity_id: str) -> Document | None: async def get_by_id(self, entity_id: str) -> Document | None:
"""Get a document by its ID.""" """Get a document by its ID."""
safe_id = escape_sql_string(entity_id) safe_id = escape_sql_string(entity_id)
results: list[DocumentRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] results = await query_to_pydantic(
self.store.documents_table.query() self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
.where(f"id = '{safe_id}'") DocumentRecord,
.limit(1)
.to_pydantic(DocumentRecord)
) )
if not results: if not results:
@ -240,7 +243,7 @@ class DocumentRepository:
query = query.limit(limit) query = query.limit(limit)
if include_content: if include_content:
results: list[DocumentRecord] = await query.to_pydantic(DocumentRecord) # type: ignore[assignment] # ty: ignore[invalid-assignment] results = await query_to_pydantic(query, DocumentRecord)
return [self._record_to_document(doc) for doc in results] return [self._record_to_document(doc) for doc in results]
return [ return [
@ -274,11 +277,9 @@ class DocumentRepository:
async def get_by_uri(self, uri: str) -> Document | None: async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI.""" """Get a document by its URI."""
escaped_uri = escape_sql_string(uri) escaped_uri = escape_sql_string(uri)
results: list[DocumentRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] results = await query_to_pydantic(
self.store.documents_table.query() self.store.documents_table.query().where(f"uri = '{escaped_uri}'").limit(1),
.where(f"uri = '{escaped_uri}'") DocumentRecord,
.limit(1)
.to_pydantic(DocumentRecord)
) )
if not results: if not results:
@ -309,9 +310,9 @@ class DocumentRepository:
# Get count before deletion # Get count before deletion
count = len( count = len(
await self.store.documents_table.query() # type: ignore[assignment] await query_to_pydantic(
.limit(1) self.store.documents_table.query().limit(1), DocumentRecord
.to_pydantic(DocumentRecord) )
) )
if count > 0: if count > 0:
# Drop and recreate table to clear all data # Drop and recreate table to clear all data

32
haiku_rag_slim/haiku/rag/store/repositories/settings.py
View file

@ -1,6 +1,6 @@
import json import json
from haiku.rag.store.engine import SettingsRecord, Store from haiku.rag.store.engine import SettingsRecord, Store, query_to_pydantic
class ConfigMismatchError(Exception): class ConfigMismatchError(Exception):
@ -23,11 +23,9 @@ class SettingsRepository:
async def get_by_id(self, entity_id: str) -> dict | None: async def get_by_id(self, entity_id: str) -> dict | None:
"""Get settings by ID.""" """Get settings by ID."""
results: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] results = await query_to_pydantic(
self.store.settings_table.query() self.store.settings_table.query().where(f"id = '{entity_id}'").limit(1),
.where(f"id = '{entity_id}'") SettingsRecord,
.limit(1)
.to_pydantic(SettingsRecord)
) )
if not results: if not results:
@ -51,20 +49,18 @@ class SettingsRepository:
self, limit: int | None = None, offset: int | None = None self, limit: int | None = None, offset: int | None = None
) -> list[dict]: ) -> list[dict]:
"""List all settings.""" """List all settings."""
results: list[ results = await query_to_pydantic(
SettingsRecord self.store.settings_table.query(), SettingsRecord
] = await self.store.settings_table.query().to_pydantic(SettingsRecord) # type: ignore[assignment] # ty: ignore[invalid-assignment] )
return [ return [
json.loads(record.settings) if record.settings else {} for record in results json.loads(record.settings) if record.settings else {} for record in results
] ]
async def get_current_settings(self) -> dict: async def get_current_settings(self) -> dict:
"""Get the current settings.""" """Get the current settings."""
results: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] results = await query_to_pydantic(
self.store.settings_table.query() self.store.settings_table.query().where("id = 'settings'").limit(1),
.where("id = 'settings'") SettingsRecord,
.limit(1)
.to_pydantic(SettingsRecord)
) )
if not results: if not results:
@ -78,11 +74,9 @@ class SettingsRepository:
current_config = self.store._config.model_dump(mode="json") current_config = self.store._config.model_dump(mode="json")
# Check if settings exist # Check if settings exist
existing: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] existing = await query_to_pydantic(
self.store.settings_table.query() self.store.settings_table.query().where("id = 'settings'").limit(1),
.where("id = 'settings'") SettingsRecord,
.limit(1)
.to_pydantic(SettingsRecord)
) )
if existing: if existing: