Basic store for chunks and documents

src/haiku/rag/store/__init__.py

@@ -0,0 +1,4 @@
+from .engine import Store
+from .models import Chunk, Document
+
+__all__ = ["Store", "Chunk", "Document"]

src/haiku/rag/store/engine.py

@@ -0,0 +1,70 @@
+import sqlite3
+import struct
+from pathlib import Path
+from typing import Literal
+
+import sqlite_vec
+
+from haiku.rag.embeddings.ollama import Embedder
+
+
+class Store:
+    def __init__(self, db_path: Path | Literal[":memory:"]):
+        self.db_path: Path | Literal[":memory:"] = db_path
+        self._connection = self.create_db()
+
+    def create_db(self) -> sqlite3.Connection:
+        """Create the database and tables with sqlite-vec support for embeddings."""
+        db = sqlite3.connect(self.db_path)
+        db.enable_load_extension(True)
+        sqlite_vec.load(db)
+
+        # Create documents table
+        db.execute("""
+            CREATE TABLE IF NOT EXISTS documents (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                content TEXT NOT NULL,
+                metadata TEXT DEFAULT '{}',
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+
+        # Create chunks table
+        db.execute("""
+            CREATE TABLE IF NOT EXISTS chunks (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                document_id INTEGER NOT NULL,
+                content TEXT NOT NULL,
+                metadata TEXT DEFAULT '{}',
+                FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE
+            )
+        """)
+
+        # Create vector table for chunk embeddings
+        embedder = Embedder()
+        db.execute(f"""
+            CREATE VIRTUAL TABLE IF NOT EXISTS chunk_embeddings USING vec0(
+                chunk_id INTEGER PRIMARY KEY,
+                embedding FLOAT[{embedder._vector_dim}]
+            )
+        """)
+
+        # Create indexes for better performance
+        db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)"
+        )
+
+        db.commit()
+        return db
+
+    @staticmethod
+    def serialize_embedding(embedding: list[float]) -> bytes:
+        """Serialize a list of floats to bytes for sqlite-vec storage."""
+        return struct.pack(f"{len(embedding)}f", *embedding)
+
+    def close(self):
+        """Close the database connection if it's an in-memory database."""
+        if self._connection is not None:
+            self._connection.close()
+            self._connection = None

src/haiku/rag/store/models/__init__.py

@@ -0,0 +1,4 @@
+from .chunk import Chunk
+from .document import Document
+
+__all__ = ["Chunk", "Document"]

src/haiku/rag/store/models/chunk.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+
+
+class Chunk(BaseModel):
+    """
+    Represents a document with an ID, content, and metadata.
+    """
+
+    id: int | None = None
+    document_id: int
+    content: str
+    metadata: dict = {}

src/haiku/rag/store/models/document.py

@@ -0,0 +1,145 @@
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from haiku.rag.chunker import chunker
+from haiku.rag.embeddings.ollama import Embedder
+from haiku.rag.store.models.chunk import Chunk
+
+if TYPE_CHECKING:
+    from haiku.rag.store.engine import Store
+
+
+class Document(BaseModel):
+    """
+    Represents a document with an ID, content, and metadata.
+    """
+
+    id: int | None = None
+    content: str
+    metadata: dict = {}
+    created_at: datetime = Field(default_factory=datetime.now)
+    updated_at: datetime = Field(default_factory=datetime.now)
+
+    async def create_with_chunks(self, store: "Store") -> "Document":
+        """
+        Create a document in the database along with its chunks and embeddings.
+
+        Args:
+            store: The Store instance to use for database operations
+
+        Returns:
+            Document: The created document with updated id
+        """
+        if store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = store._connection.cursor()
+        embedder = Embedder()
+
+        # Insert the document
+        cursor.execute(
+            """
+            INSERT INTO documents (content, metadata, created_at, updated_at)
+            VALUES (?, ?, ?, ?)
+            """,
+            (
+                self.content,
+                json.dumps(self.metadata),
+                self.created_at,
+                self.updated_at,
+            ),
+        )
+
+        document_id = cursor.lastrowid
+        assert document_id is not None, "Failed to create document in database"
+        self.id = document_id
+
+        # Chunk the document content
+        chunk_texts = await chunker.chunk(self.content)
+
+        # Create chunks with embeddings
+        for order, chunk_text in enumerate(chunk_texts):
+            # Create chunk with order in metadata
+            chunk = Chunk(
+                document_id=document_id, content=chunk_text, metadata={"order": order}
+            )
+
+            cursor.execute(
+                """
+                INSERT INTO chunks (document_id, content, metadata)
+                VALUES (?, ?, ?)
+                """,
+                (chunk.document_id, chunk.content, json.dumps(chunk.metadata)),
+            )
+            chunk_id = cursor.lastrowid
+
+            # Generate and store embedding
+            embedding = await embedder.embed(chunk_text)
+            serialized_embedding = store.serialize_embedding(embedding)
+            cursor.execute(
+                """
+                INSERT INTO chunk_embeddings (chunk_id, embedding)
+                VALUES (?, ?)
+                """,
+                (chunk_id, serialized_embedding),
+            )
+
+        store._connection.commit()
+        return self
+
+    @classmethod
+    async def search_chunks(
+        cls, store: "Store", query: str, limit: int = 5
+    ) -> list[Chunk]:
+        """
+        Search for relevant chunks using vector similarity with sqlite-vec.
+
+        Args:
+            store: The Store instance to use for database operations
+            query: The text query to search for
+            limit: Maximum number of chunks to return
+
+        Returns:
+            List of relevant Chunk objects ordered by similarity
+        """
+        if store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        embedder = Embedder()
+        cursor = store._connection.cursor()
+
+        # Generate embedding for the query
+        query_embedding = await embedder.embed(query)
+        serialized_query_embedding = store.serialize_embedding(query_embedding)
+
+        # Search for similar chunks using sqlite-vec
+        cursor.execute(
+            """
+            SELECT c.id, c.document_id, c.content, c.metadata, distance
+            FROM chunk_embeddings
+            JOIN chunks c ON c.id = chunk_embeddings.chunk_id
+            WHERE embedding MATCH ? AND k = ?
+            ORDER BY distance
+            """,
+            (serialized_query_embedding, limit),
+        )
+
+        results = cursor.fetchall()
+        chunks = []
+
+        for row in results:
+            chunk_id, document_id, content, metadata_json, distance = row
+            metadata = json.loads(metadata_json) if metadata_json else {}
+            chunks.append(
+                Chunk(
+                    id=chunk_id,
+                    document_id=document_id,
+                    content=content,
+                    metadata=metadata,
+                )
+            )
+
+        return chunks

tests/test_document.py

@@ -0,0 +1,89 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.models.document import Document
+
+
+@pytest.mark.asyncio
+async def test_create_document_with_chunks(qa_corpus: Dataset):
+    """Test creating a document with chunks from the qa_corpus."""
+    # Create an in-memory store
+    store = Store(":memory:")
+    
+    # Get the first document from the corpus
+    first_doc = qa_corpus[0]
+    document_text = first_doc["document_extracted"]
+    
+    # Create a Document instance
+    document = Document(
+        content=document_text,
+        metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")}
+    )
+    
+    # Create the document with chunks in the database
+    created_document = await document.create_with_chunks(store)
+    
+    # Verify the document was created
+    assert created_document.id is not None
+    assert created_document.content == document_text
+    
+    # Check that chunks were created in the database
+    if store._connection is not None:
+        cursor = store._connection.cursor()
+        cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (created_document.id,))
+        chunk_count = cursor.fetchone()[0]
+        
+        assert chunk_count > 0
+        
+        # Check that embeddings were created
+        cursor.execute("""
+            SELECT COUNT(*) FROM chunk_embeddings ce
+            JOIN chunks c ON c.id = ce.chunk_id
+            WHERE c.document_id = ?
+        """, (created_document.id,))
+        embedding_count = cursor.fetchone()[0]
+        
+        assert embedding_count == chunk_count
+        
+        # Verify chunk metadata contains order information
+        cursor.execute("SELECT metadata FROM chunks WHERE document_id = ? ORDER BY id", (created_document.id,))
+        chunk_metadata = cursor.fetchall()
+        
+        for i, (metadata_json,) in enumerate(chunk_metadata):
+            import json
+            metadata = json.loads(metadata_json)
+            assert "order" in metadata
+            assert metadata["order"] == i
+    
+    store.close()
+
+
+@pytest.mark.asyncio
+async def test_search_chunks(qa_corpus: Dataset):
+    """Test vector search functionality."""
+    # Create an in-memory store
+    store = Store(":memory:")
+    
+    # Get the first document from the corpus
+    first_doc = qa_corpus[0]
+    document_text = first_doc["document_extracted"]
+    
+    # Create and store a document
+    document = Document(
+        content=document_text,
+        metadata={"source": "qa_corpus"}
+    )
+    await document.create_with_chunks(store)
+    
+    # Perform a search
+    search_query = "news"  # Simple query
+    results = await Document.search_chunks(store, search_query, limit=3)
+    
+    # Verify search results
+    assert len(results) <= 3
+    assert all(hasattr(chunk, "content") for chunk in results)
+    assert all(hasattr(chunk, "document_id") for chunk in results)
+    assert all(chunk.document_id == document.id for chunk in results)
+    
+    store.close()