diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 08f3279a..bd1eeb5a 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta from enum import Enum from importlib import metadata from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any, cast from uuid import uuid4 import lancedb @@ -18,9 +18,25 @@ from haiku.rag.config import AppConfig, Config from haiku.rag.embeddings import get_embedder from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError +if TYPE_CHECKING: + from lancedb.query import AsyncQueryBase + 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): LOCAL = "local" CLOUD = "cloud" @@ -92,19 +108,26 @@ def get_documents_arrow_schema() -> pa.Schema: return pa.schema(fields) -def create_chunk_model(vector_dim: int): - """Create a ChunkRecord model with the specified vector dimension. - - This creates a model with proper vector typing for LanceDB. +class ChunkRecordBase(LanceModel): + """Static base for ChunkRecord — declares the fields so attribute access + and constructor calls type-check. The concrete `vector` field is overridden + 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())) - document_id: str - content: str - content_fts: str = Field(default="") - metadata: str = Field(default="{}") - order: int = Field(default=0) + id: str = Field(default_factory=lambda: str(uuid4())) + document_id: str + content: str + content_fts: str = Field(default="") + metadata: str = Field(default="{}") + 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 return ChunkRecord @@ -230,7 +253,7 @@ class Store: # Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB) 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) await self._init_tables() @@ -506,8 +529,8 @@ class Store: async def get_haiku_version(self) -> str: """Returns the user version stored in settings.""" - settings_records: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.settings_table.query().limit(1).to_pydantic(SettingsRecord) + settings_records = await query_to_pydantic( + self.settings_table.query().limit(1), SettingsRecord ) if settings_records: settings = ( @@ -525,8 +548,8 @@ class Store: ReadOnlyError: If the store is in read-only mode. """ self._assert_writable() - settings_records: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.settings_table.query().limit(1).to_pydantic(SettingsRecord) + settings_records = await query_to_pydantic( + self.settings_table.query().limit(1), SettingsRecord ) if settings_records: # Only write if version actually changes to avoid creating new table versions diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 81362be9..d0be740a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: from lancedb.index import FTS 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 logger = logging.getLogger(__name__) @@ -89,6 +89,7 @@ class ChunkRepository: chunk_id = str(uuid4()) assert chunk.document_id is not None + assert chunk.embedding is not None chunk_record = self.store.ChunkRecord( id=chunk_id, document_id=chunk.document_id, @@ -110,24 +111,22 @@ class ChunkRepository: async def get_by_id(self, entity_id: str) -> Chunk | None: """Get a chunk by its ID.""" - results = await ( - self.store.chunks_table.query() - .where(f"id = '{entity_id}'") - .limit(1) - .to_pydantic(self.store.ChunkRecord) + results = await query_to_pydantic( + self.store.chunks_table.query().where(f"id = '{entity_id}'").limit(1), + self.store.ChunkRecord, ) if not results: return None chunk_record = results[0] - md = json.loads(chunk_record.metadata) # ty: ignore[unresolved-attribute] + md = json.loads(chunk_record.metadata) return Chunk( - id=chunk_record.id, # ty: ignore[unresolved-attribute] - document_id=chunk_record.document_id, # ty: ignore[unresolved-attribute] - content=chunk_record.content, # ty: ignore[unresolved-attribute] + id=chunk_record.id, + document_id=chunk_record.document_id, + content=chunk_record.content, metadata=md, - order=chunk_record.order, # ty: ignore[unresolved-attribute] + order=chunk_record.order, ) async def update(self, entity: Chunk) -> Chunk: @@ -175,18 +174,18 @@ class ChunkRepository: if limit is not None: query = query.limit(limit) - results = await query.to_pydantic(self.store.ChunkRecord) + results = await query_to_pydantic(query, self.store.ChunkRecord) chunks: list[Chunk] = [] for rec in results: - md = json.loads(rec.metadata) # ty: ignore[unresolved-attribute] + md = json.loads(rec.metadata) chunks.append( Chunk( - id=rec.id, # ty: ignore[unresolved-attribute] - document_id=rec.document_id, # ty: ignore[unresolved-attribute] - content=rec.content, # ty: ignore[unresolved-attribute] + id=rec.id, + document_id=rec.document_id, + content=rec.content, metadata=md, - order=rec.order, # ty: ignore[unresolved-attribute] + order=rec.order, ) ) return chunks @@ -316,7 +315,7 @@ class ChunkRepository: if limit is not None: 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) doc_rows = await ( @@ -333,14 +332,14 @@ class ChunkRepository: chunks: list[Chunk] = [] for rec in results: - md = json.loads(rec.metadata) # ty: ignore[unresolved-attribute] + md = json.loads(rec.metadata) chunks.append( Chunk( - id=rec.id, # ty: ignore[unresolved-attribute] - document_id=rec.document_id, # ty: ignore[unresolved-attribute] - content=rec.content, # ty: ignore[unresolved-attribute] + id=rec.id, + document_id=rec.document_id, + content=rec.content, metadata=md, - order=rec.order, # ty: ignore[unresolved-attribute] + order=rec.order, document_uri=doc_uri, document_title=doc_title, document_meta=json.loads(doc_meta), @@ -378,18 +377,16 @@ class ChunkRepository: f" AND `order` >= {min_order}" f" AND `order` <= {max_order}" ) - results = await ( - self.store.chunks_table.query() - .where(where) - .to_pydantic(self.store.ChunkRecord) + results = await query_to_pydantic( + self.store.chunks_table.query().where(where), self.store.ChunkRecord ) return [ Chunk( - id=rec.id, # ty: ignore[unresolved-attribute] - document_id=rec.document_id, # ty: ignore[unresolved-attribute] - content=rec.content, # ty: ignore[unresolved-attribute] - metadata=json.loads(rec.metadata), # ty: ignore[unresolved-attribute] - order=rec.order, # ty: ignore[unresolved-attribute] + id=rec.id, + document_id=rec.document_id, + content=rec.content, + metadata=json.loads(rec.metadata), + order=rec.order, ) for rec in results ] diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 731e5e7b..d2d0b234 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -4,7 +4,12 @@ from uuid import uuid4 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.utils import escape_sql_string @@ -90,11 +95,9 @@ class DocumentRepository: async def get_by_id(self, entity_id: str) -> Document | None: """Get a document by its ID.""" safe_id = escape_sql_string(entity_id) - results: list[DocumentRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.store.documents_table.query() - .where(f"id = '{safe_id}'") - .limit(1) - .to_pydantic(DocumentRecord) + results = await query_to_pydantic( + self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1), + DocumentRecord, ) if not results: @@ -240,7 +243,7 @@ class DocumentRepository: query = query.limit(limit) 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 [ @@ -274,11 +277,9 @@ class DocumentRepository: async def get_by_uri(self, uri: str) -> Document | None: """Get a document by its URI.""" escaped_uri = escape_sql_string(uri) - results: list[DocumentRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.store.documents_table.query() - .where(f"uri = '{escaped_uri}'") - .limit(1) - .to_pydantic(DocumentRecord) + results = await query_to_pydantic( + self.store.documents_table.query().where(f"uri = '{escaped_uri}'").limit(1), + DocumentRecord, ) if not results: @@ -309,9 +310,9 @@ class DocumentRepository: # Get count before deletion count = len( - await self.store.documents_table.query() # type: ignore[assignment] - .limit(1) - .to_pydantic(DocumentRecord) + await query_to_pydantic( + self.store.documents_table.query().limit(1), DocumentRecord + ) ) if count > 0: # Drop and recreate table to clear all data diff --git a/haiku_rag_slim/haiku/rag/store/repositories/settings.py b/haiku_rag_slim/haiku/rag/store/repositories/settings.py index 5ef126cf..40e206f4 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/settings.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/settings.py @@ -1,6 +1,6 @@ import json -from haiku.rag.store.engine import SettingsRecord, Store +from haiku.rag.store.engine import SettingsRecord, Store, query_to_pydantic class ConfigMismatchError(Exception): @@ -23,11 +23,9 @@ class SettingsRepository: async def get_by_id(self, entity_id: str) -> dict | None: """Get settings by ID.""" - results: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.store.settings_table.query() - .where(f"id = '{entity_id}'") - .limit(1) - .to_pydantic(SettingsRecord) + results = await query_to_pydantic( + self.store.settings_table.query().where(f"id = '{entity_id}'").limit(1), + SettingsRecord, ) if not results: @@ -51,20 +49,18 @@ class SettingsRepository: self, limit: int | None = None, offset: int | None = None ) -> list[dict]: """List all settings.""" - results: list[ - SettingsRecord - ] = await self.store.settings_table.query().to_pydantic(SettingsRecord) # type: ignore[assignment] # ty: ignore[invalid-assignment] + results = await query_to_pydantic( + self.store.settings_table.query(), SettingsRecord + ) return [ json.loads(record.settings) if record.settings else {} for record in results ] async def get_current_settings(self) -> dict: """Get the current settings.""" - results: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.store.settings_table.query() - .where("id = 'settings'") - .limit(1) - .to_pydantic(SettingsRecord) + results = await query_to_pydantic( + self.store.settings_table.query().where("id = 'settings'").limit(1), + SettingsRecord, ) if not results: @@ -78,11 +74,9 @@ class SettingsRepository: current_config = self.store._config.model_dump(mode="json") # Check if settings exist - existing: list[SettingsRecord] = await ( # type: ignore[assignment] # ty: ignore[invalid-assignment] - self.store.settings_table.query() - .where("id = 'settings'") - .limit(1) - .to_pydantic(SettingsRecord) + existing = await query_to_pydantic( + self.store.settings_table.query().where("id = 'settings'").limit(1), + SettingsRecord, ) if existing: