Satisfy ty by adding a typed to_pydantic helper

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

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
         ]

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

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: