diff --git a/src/haiku/rag/store/__init__.py b/src/haiku/rag/store/__init__.py new file mode 100644 index 00000000..ad1e8338 --- /dev/null +++ b/src/haiku/rag/store/__init__.py @@ -0,0 +1,4 @@ +from .engine import Store +from .models import Chunk, Document + +__all__ = ["Store", "Chunk", "Document"] \ No newline at end of file diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py new file mode 100644 index 00000000..b0f53893 --- /dev/null +++ b/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 diff --git a/src/haiku/rag/store/models/__init__.py b/src/haiku/rag/store/models/__init__.py new file mode 100644 index 00000000..3e9d83be --- /dev/null +++ b/src/haiku/rag/store/models/__init__.py @@ -0,0 +1,4 @@ +from .chunk import Chunk +from .document import Document + +__all__ = ["Chunk", "Document"] \ No newline at end of file diff --git a/src/haiku/rag/store/models/chunk.py b/src/haiku/rag/store/models/chunk.py new file mode 100644 index 00000000..5b17bea4 --- /dev/null +++ b/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 = {} diff --git a/src/haiku/rag/store/models/document.py b/src/haiku/rag/store/models/document.py new file mode 100644 index 00000000..84ce81ea --- /dev/null +++ b/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 diff --git a/tests/test_document.py b/tests/test_document.py new file mode 100644 index 00000000..529188b4 --- /dev/null +++ b/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() \ No newline at end of file