Merge pull request #44 from ggozad/feat/lance-db

Replace sqlite & sqlite-vac with lancedb.
This commit is contained in:
Yiorgis Gozadinos 2025-09-03 11:13:41 +03:00 committed by GitHub
commit 6990a05ad7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
47 changed files with 1788 additions and 1451 deletions

View file

@ -21,12 +21,12 @@ repos:
hooks: hooks:
- id: pyright - id: pyright
- repo: https://github.com/RodrigoGonzalez/check-mkdocs # - repo: https://github.com/RodrigoGonzalez/check-mkdocs
rev: v1.2.0 # rev: v1.2.0
hooks: # hooks:
- id: check-mkdocs # - id: check-mkdocs
name: check-mkdocs # name: check-mkdocs
args: ["--config", "mkdocs.yml"] # Optional, mkdocs.yml is the default # args: ["--config", "mkdocs.yml"] # Optional, mkdocs.yml is the default
# If you have additional plugins or libraries that are not included in # # If you have additional plugins or libraries that are not included in
# check-mkdocs, add them here # # check-mkdocs, add them here
additional_dependencies: ["mkdocs-material"] # additional_dependencies: ["mkdocs-material"]

View file

@ -1,15 +1,17 @@
# Haiku SQLite RAG # Haiku RAG
Retrieval-Augmented Generation (RAG) library on SQLite. Retrieval-Augmented Generation (RAG) library built on LanceDB.
`haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work on SQLite alone without the need for external vector databases. It uses [sqlite-vec](https://github.com/asg017/sqlite-vec) for storing the embeddings and performs semantic (vector) search as well as full-text search combined through Reciprocal Rank Fusion. Both open-source (Ollama) as well as commercial (OpenAI, VoyageAI) embedding providers are supported. `haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work with LanceDB as a local vector database. It uses LanceDB for storing embeddings and performs semantic (vector) search as well as full-text search combined through native hybrid search with Reciprocal Rank Fusion. Both open-source (Ollama) as well as commercial (OpenAI, VoyageAI) embedding providers are supported.
> **Note**: Starting with version 0.7.0, haiku.rag uses LanceDB instead of SQLite. If you have an existing SQLite database, use `haiku-rag migrate old_database.sqlite` to migrate your data safely.
## Features ## Features
- **Local SQLite**: No external servers required - **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI
- **Multiple QA providers**: Any provider/model supported by Pydantic AI - **Multiple QA providers**: Any provider/model supported by Pydantic AI
- **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
- **Reranking**: Default search result reranking with MixedBread AI or Cohere - **Reranking**: Default search result reranking with MixedBread AI or Cohere
- **Question answering**: Built-in QA agents on your documents - **Question answering**: Built-in QA agents on your documents
- **File monitoring**: Auto-index files when run as server - **File monitoring**: Auto-index files when run as server
@ -39,6 +41,9 @@ haiku-rag ask "Who is the author of haiku.rag?" --cite
# Rebuild database (re-chunk and re-embed all documents) # Rebuild database (re-chunk and re-embed all documents)
haiku-rag rebuild haiku-rag rebuild
# Migrate from SQLite to LanceDB
haiku-rag migrate old_database.sqlite
# Start server with file monitoring # Start server with file monitoring
export MONITOR_DIRECTORIES="/path/to/docs" export MONITOR_DIRECTORIES="/path/to/docs"
haiku-rag serve haiku-rag serve
@ -49,7 +54,7 @@ haiku-rag serve
```python ```python
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
async with HaikuRAG("database.db") as client: async with HaikuRAG("database.lancedb") as client:
# Add document # Add document
doc = await client.create_document("Your content") doc = await client.create_document("Your content")

View file

@ -7,19 +7,19 @@ You can perform your own evaluations using as example the script found at
## Recall ## Recall
In order to calculate recall, we load the `News Stories` from `repliqa_3` which is 1035 documents and index them in a sqlite db. Subsequently, we run a search over the `question` field for each row of the dataset and check whether we match the document that answers the question. In order to calculate recall, we load the `News Stories` from `repliqa_3` (1035 documents) and index them. Subsequently, we run a search over the `question` field for each row of the dataset and check whether we match the document that answers the question. Questions for which the answer cannot be found in the documents are ignored.
The recall obtained is ~0.73 for matching in the top result, raising to ~0.75 for the top 3 results. The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 for the top 3 results with the "bare" default settings (Ollama `qwen3`, `mxbai-embed-large` embeddings, no reranking).
| Embedding Model | Document in top 1 | Document in top 3 | Reranker | | Embedding Model | Document in top 1 | Document in top 3 | Reranker |
|---------------------------------------|-------------------|-------------------|------------------------| |---------------------------------------|-------------------|-------------------|------------------------|
| Ollama / `mxbai-embed-large` | 0.77 | 0.89 | None | | Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None |
| Ollama / `mxbai-embed-large` | 0.81 | 0.91 | `mxbai-rerank-base-v2` | | Ollama / `mxbai-embed-large` | 0.90 | 0.95 | `mxbai-rerank-base-v2` |
| Ollama / `nomic-embed-text` | 0.74 | 0.88 | None | <!-- | Ollama / `nomic-embed-text` | 0.74 | 0.88 | None |
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None | | OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None | | OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
| OpenAI / `text-embeddings-3-small` | 0.83 | 0.90 | Cohere / `rerank-v3.5` | | OpenAI / `text-embeddings-3-small` | 0.83 | 0.90 | Cohere / `rerank-v3.5` | -->
## Question/Answer evaluation ## Question/Answer evaluation
@ -27,7 +27,10 @@ Again using the same dataset, we use a QA agent to answer the question. In addit
| Embedding Model | QA Model | Accuracy | Reranker | | Embedding Model | QA Model | Accuracy | Reranker |
|------------------------------------|-----------------------------------|-----------|------------------------| |------------------------------------|-----------------------------------|-----------|------------------------|
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.64 | None | | Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.85 | None |
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.72 | `mxbai-rerank-base-v2` | | Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.87 | `mxbai-rerank-base-v2` |
| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None | | Ollama / `mxbai-embed-large` | Ollama / `qwen3:0.6b` | 0.28 | None |
| OpenAI / `text-embeddings-3-small` | OpenAI / `gpt-4-turbo` | 0.62 | None |
Note the significant degradation when very small models are used such as `qwen3:0.6b`.
<!-- | Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None |
| OpenAI / `text-embeddings-3-small` | OpenAI / `gpt-4-turbo` | 0.62 | None | -->

View file

@ -45,6 +45,23 @@ haiku-rag rebuild
Use this when you want to change things like the embedding model or chunk size for example. Use this when you want to change things like the embedding model or chunk size for example.
## Migration
### Migrate from SQLite to LanceDB
Migrate an existing SQLite database to LanceDB:
```bash
haiku-rag migrate /path/to/old_database.sqlite
```
This will:
- Read all documents, chunks, embeddings, and settings from the SQLite database
- Create a new LanceDB database with the same data in the same directory
- Optimize the new database for best performance
The original SQLite database remains unchanged, so you can safely migrate without risk of data loss.
## Search ## Search
Basic search: Basic search:
@ -54,7 +71,7 @@ haiku-rag search "machine learning"
With options: With options:
```bash ```bash
haiku-rag search "python programming" --limit 10 --k 100 haiku-rag search "python programming" --limit 10
``` ```
## Question Answering ## Question Answering

View file

@ -109,25 +109,7 @@ See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the com
Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results.
Reranking is **automatically enabled** by default using Ollama, or if you install the appropriate reranking provider package. Reranking is **disabled by default** (`RERANK_PROVIDER=""`) for faster searches. You can enable it by configuring one of the providers below.
### Disabling Reranking
To disable reranking completely for faster searches:
```bash
RERANK_PROVIDER=""
```
### Ollama (Default)
Ollama reranking uses LLMs with structured output to rank documents by relevance:
```bash
RERANK_PROVIDER="ollama"
RERANK_MODEL="qwen3:1.7b" # or any model that supports structured output
OLLAMA_BASE_URL="http://localhost:11434"
```
### MixedBread AI ### MixedBread AI
@ -158,11 +140,41 @@ COHERE_API_KEY="your-api-key"
### Database and Storage ### Database and Storage
By default, `haiku.rag` uses a local LanceDB database:
```bash ```bash
# Default data directory (where SQLite database is stored) # Default data directory (where local LanceDB is stored)
DEFAULT_DATA_DIR="/path/to/data" DEFAULT_DATA_DIR="/path/to/data"
``` ```
For remote storage, use the `LANCEDB_URI` setting with various backends:
```bash
# LanceDB Cloud
LANCEDB_URI="db://your-database-name"
LANCEDB_API_KEY="your-api-key"
LANCEDB_REGION="us-west-2" # optional
# Amazon S3
LANCEDB_URI="s3://my-bucket/my-table"
# Use AWS credentials or IAM roles
# Azure Blob Storage
LANCEDB_URI="az://my-container/my-table"
# Use Azure credentials
# Google Cloud Storage
LANCEDB_URI="gs://my-bucket/my-table"
# Use GCP credentials
# HDFS
LANCEDB_URI="hdfs://namenode:port/path/to/table"
```
Authentication is handled through standard cloud provider credentials (AWS CLI, Azure CLI, gcloud, etc.) or by setting `LANCEDB_API_KEY` for LanceDB Cloud.
**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization is still performed locally.
### Document Processing ### Document Processing
```bash ```bash

View file

@ -1,12 +1,14 @@
# haiku.rag # haiku.rag
`haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work on SQLite alone without the need for external vector databases. It uses [sqlite-vec](https://github.com/asg017/sqlite-vec) for storing the embeddings and performs semantic (vector) search as well as full-text search combined through Reciprocal Rank Fusion. Both open-source (Ollama, MixedBread AI) as well as commercial (OpenAI, VoyageAI) embedding providers are supported. `haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work with LanceDB as a local vector database. It uses LanceDB for storing embeddings and performs semantic (vector) search as well as full-text search combined through native hybrid search with Reciprocal Rank Fusion. Both open-source (Ollama, MixedBread AI) as well as commercial (OpenAI, VoyageAI) embedding providers are supported.
> **Note**: Starting with version 0.7.0, haiku.rag uses LanceDB instead of SQLite. If you have an existing SQLite database, use `haiku-rag migrate old_database.sqlite` to migrate your data safely.
## Features ## Features
- **Local SQLite**: No need to run additional servers - **Local LanceDB**: No need to run additional servers
- **Support for various embedding providers**: Ollama, VoyageAI, OpenAI or add your own - **Support for various embedding providers**: Ollama, VoyageAI, OpenAI or add your own
- **Hybrid Search**: Vector search using `sqlite-vec` combined with full-text search `FTS5`, using Reciprocal Rank Fusion - **Native Hybrid Search**: Vector search combined with full-text search using native LanceDB RRF reranking
- **Reranking**: Optional result reranking with MixedBread AI or Cohere - **Reranking**: Optional result reranking with MixedBread AI or Cohere
- **Question Answering**: Built-in QA agents using Ollama, OpenAI, or Anthropic. - **Question Answering**: Built-in QA agents using Ollama, OpenAI, or Anthropic.
- **File monitoring**: Automatically index files when run as a server - **File monitoring**: Automatically index files when run as a server
@ -26,7 +28,7 @@ Use from Python:
```python ```python
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
async with HaikuRAG("database.db") as client: async with HaikuRAG("database.lancedb") as client:
# Add a document # Add a document
doc = await client.create_document("Your content here") doc = await client.create_document("Your content here")
@ -34,7 +36,7 @@ async with HaikuRAG("database.db") as client:
results = await client.search("query") results = await client.search("query")
# Ask questions # Ask questions
answer = await client.ask("Who is the author of haiku.rag?", rerank=False) answer = await client.ask("Who is the author of haiku.rag?")
``` ```
Or use the CLI: Or use the CLI:
@ -42,6 +44,7 @@ Or use the CLI:
haiku-rag add "Your document content" haiku-rag add "Your document content"
haiku-rag search "query" haiku-rag search "query"
haiku-rag ask "Who is the author of haiku.rag?" haiku-rag ask "Who is the author of haiku.rag?"
haiku-rag migrate old_database.sqlite # Migrate from SQLite
``` ```
## Documentation ## Documentation

View file

@ -31,5 +31,4 @@ uv pip install haiku.rag[mxbai]
## Requirements ## Requirements
- Python 3.10+ - Python 3.10+
- SQLite 3.38+
- Ollama (for default embeddings) - Ollama (for default embeddings)

View file

@ -19,10 +19,10 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients.
## Starting MCP Server ## Starting MCP Server
The MCP server starts automatically with the serve command and supports `Streamable HTTP`, `stdio` and `SSE` transports: The MCP server starts automatically with the serve command and supports Streamable HTTP, stdio and SSE transports:
```bash ```bash
# Default HTTP transport # Default streamable HTTP transport
haiku-rag serve haiku-rag serve
# stdio transport (for Claude Desktop) # stdio transport (for Claude Desktop)

View file

@ -9,7 +9,7 @@ from pathlib import Path
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
# Use as async context manager (recommended) # Use as async context manager (recommended)
async with HaikuRAG("path/to/database.db") as client: async with HaikuRAG("path/to/database.lancedb") as client:
# Your code here # Your code here
pass pass
``` ```
@ -101,9 +101,9 @@ async for doc_id in client.rebuild_database():
## Searching Documents ## Searching Documents
The search method performs hybrid search (vector + full-text) with **reranking enabled by default** for improved relevance: The search method performs native hybrid search (vector + full-text) using LanceDB with optional reranking for improved relevance:
Basic search (with reranking): Basic hybrid search (default):
```python ```python
results = await client.search("machine learning algorithms", limit=5) results = await client.search("machine learning algorithms", limit=5)
for chunk, score in results: for chunk, score in results:
@ -112,13 +112,27 @@ for chunk, score in results:
print(f"Document ID: {chunk.document_id}") print(f"Document ID: {chunk.document_id}")
``` ```
With options: Search with different search types:
```python ```python
# Vector search only
results = await client.search( results = await client.search(
query="machine learning", query="machine learning",
limit=5, # Maximum results to return limit=5,
k=60, # RRF parameter for reciprocal rank fusion search_type="vector"
rerank=False # Disable reranking for faster search )
# Full-text search only
results = await client.search(
query="machine learning",
limit=5,
search_type="fts"
)
# Hybrid search (default - combines vector + fts with native LanceDB RRF)
results = await client.search(
query="machine learning",
limit=5,
search_type="hybrid"
) )
# Process results # Process results

View file

@ -9,7 +9,7 @@ haiku-rag serve
``` ```
Transport options: Transport options:
- `--http` (default) - Streamable HTTP transport - Default - Streamable HTTP transport
- `--stdio` - Standard input/output transport - `--stdio` - Standard input/output transport
- `--sse` - Server-sent events transport - `--sse` - Server-sent events transport

View file

@ -1,5 +1,5 @@
site_name: haiku.rag site_name: haiku.rag
site_description: Retrieval-Augmented Generation (RAG) library on SQLite. site_description: Retrieval-Augmented Generation (RAG) library on LanceDB.
site_url: https://ggozad.github.io/haiku.rag/ site_url: https://ggozad.github.io/haiku.rag/
theme: theme:
name: material name: material

View file

@ -1,12 +1,12 @@
[project] [project]
name = "haiku.rag" name = "haiku.rag"
version = "0.6.0" version = "0.6.0"
description = "Retrieval Augmented Generation (RAG) with SQLite" description = "Retrieval Augmented Generation (RAG) with LanceDB"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" } license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" } readme = { file = "README.md", content-type = "text/markdown" }
requires-python = ">=3.11" requires-python = ">=3.11"
keywords = ["RAG", "sqlite", "sqlite-vec", "ml", "mcp"] keywords = ["RAG", "lancedb", "vector-database", "ml", "mcp"]
classifiers = [ classifiers = [
"Development Status :: 4 - Beta", "Development Status :: 4 - Beta",
"Environment :: Console", "Environment :: Console",
@ -25,12 +25,12 @@ dependencies = [
"docling>=2.15.0", "docling>=2.15.0",
"fastmcp>=2.8.1", "fastmcp>=2.8.1",
"httpx>=0.28.1", "httpx>=0.28.1",
"lancedb>=0.17.0",
"ollama>=0.5.3", "ollama>=0.5.3",
"pydantic>=2.11.7", "pydantic>=2.11.7",
"pydantic-ai>=0.7.2", "pydantic-ai>=0.7.2",
"python-dotenv>=1.1.0", "python-dotenv>=1.1.0",
"rich>=14.0.0", "rich>=14.0.0",
"sqlite-vec>=0.1.6",
"tiktoken>=0.9.0", "tiktoken>=0.9.0",
"typer>=0.16.0", "typer>=0.16.0",
"watchfiles>=1.1.0", "watchfiles>=1.1.0",
@ -56,7 +56,7 @@ dev = [
"mkdocs>=1.6.1", "mkdocs>=1.6.1",
"mkdocs-material>=9.6.14", "mkdocs-material>=9.6.14",
"pre-commit>=4.2.0", "pre-commit>=4.2.0",
"pyright>=1.1.403", "pyright>=1.1.404",
"pytest>=8.4.0", "pytest>=8.4.0",
"pytest-asyncio>=1.0.0", "pytest-asyncio>=1.0.0",
"pytest-cov>=6.2.1", "pytest-cov>=6.2.1",

View file

@ -40,7 +40,7 @@ class HaikuRAGApp:
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]" f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
) )
async def get_document(self, doc_id: int): async def get_document(self, doc_id: str):
async with HaikuRAG(db_path=self.db_path) as self.client: async with HaikuRAG(db_path=self.db_path) as self.client:
doc = await self.client.get_document_by_id(doc_id) doc = await self.client.get_document_by_id(doc_id)
if doc is None: if doc is None:
@ -48,14 +48,14 @@ class HaikuRAGApp:
return return
self._rich_print_document(doc, truncate=False) self._rich_print_document(doc, truncate=False)
async def delete_document(self, doc_id: int): async def delete_document(self, doc_id: str):
async with HaikuRAG(db_path=self.db_path) as self.client: async with HaikuRAG(db_path=self.db_path) as self.client:
await self.client.delete_document(doc_id) await self.client.delete_document(doc_id)
self.console.print(f"[b]Document {doc_id} deleted successfully.[/b]") self.console.print(f"[b]Document {doc_id} deleted successfully.[/b]")
async def search(self, query: str, limit: int = 5, k: int = 60): async def search(self, query: str, limit: int = 5):
async with HaikuRAG(db_path=self.db_path) as self.client: async with HaikuRAG(db_path=self.db_path) as self.client:
results = await self.client.search(query, limit=limit, k=k) results = await self.client.search(query, limit=limit)
if not results: if not results:
self.console.print("[red]No results found.[/red]") self.console.print("[red]No results found.[/red]")
return return

View file

@ -8,6 +8,7 @@ from rich.console import Console
from haiku.rag.app import HaikuRAGApp from haiku.rag.app import HaikuRAGApp
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.migration import migrate_sqlite_to_lancedb
from haiku.rag.utils import is_up_to_date from haiku.rag.utils import is_up_to_date
if not Config.ENV == "development": if not Config.ENV == "development":
@ -47,7 +48,7 @@ def main(
help="Show version and exit", help="Show version and exit",
), ),
): ):
"""haiku.rag CLI - SQLite-based RAG system""" """haiku.rag CLI - Vector database RAG system"""
# Run version check before any command # Run version check before any command
asyncio.run(check_version()) asyncio.run(check_version())
@ -55,9 +56,9 @@ def main(
@cli.command("list", help="List all stored documents") @cli.command("list", help="List all stored documents")
def list_documents( def list_documents(
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
@ -70,9 +71,9 @@ def add_document_text(
help="The text content of the document to add", help="The text content of the document to add",
), ),
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
@ -85,9 +86,9 @@ def add_document_src(
help="The file path or URL of the document to add", help="The file path or URL of the document to add",
), ),
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
@ -96,13 +97,13 @@ def add_document_src(
@cli.command("get", help="Get and display a document by its ID") @cli.command("get", help="Get and display a document by its ID")
def get_document( def get_document(
doc_id: int = typer.Argument( doc_id: str = typer.Argument(
help="The ID of the document to get", help="The ID of the document to get",
), ),
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
@ -111,13 +112,13 @@ def get_document(
@cli.command("delete", help="Delete a document by its ID") @cli.command("delete", help="Delete a document by its ID")
def delete_document( def delete_document(
doc_id: int = typer.Argument( doc_id: str = typer.Argument(
help="The ID of the document to delete", help="The ID of the document to delete",
), ),
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
@ -135,19 +136,14 @@ def search(
"-l", "-l",
help="Maximum number of results to return", help="Maximum number of results to return",
), ),
k: int = typer.Option(
60,
"--k",
help="Reciprocal Rank Fusion k parameter",
),
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
asyncio.run(app.search(query=query, limit=limit, k=k)) asyncio.run(app.search(query=query, limit=limit))
@cli.command("ask", help="Ask a question using the QA agent") @cli.command("ask", help="Ask a question using the QA agent")
@ -156,9 +152,9 @@ def ask(
help="The question to ask", help="The question to ask",
), ),
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
cite: bool = typer.Option( cite: bool = typer.Option(
False, False,
@ -182,9 +178,9 @@ def settings():
) )
def rebuild( def rebuild(
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
): ):
app = HaikuRAGApp(db_path=db) app = HaikuRAGApp(db_path=db)
@ -196,9 +192,9 @@ def rebuild(
) )
def serve( def serve(
db: Path = typer.Option( db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite", Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db", "--db",
help="Path to the SQLite database file", help="Path to the LanceDB database file",
), ),
stdio: bool = typer.Option( stdio: bool = typer.Option(
False, False,
@ -227,5 +223,20 @@ def serve(
asyncio.run(app.serve(transport=transport)) asyncio.run(app.serve(transport=transport))
@cli.command("migrate", help="Migrate an SQLite database to LanceDB")
def migrate(
sqlite_path: Path = typer.Argument(
help="Path to the SQLite database file to migrate",
),
):
# Generate LanceDB path in same parent directory
lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path))
if not success:
raise typer.Exit(1)
if __name__ == "__main__": if __name__ == "__main__":
cli() cli()

View file

@ -3,7 +3,6 @@ import mimetypes
import tempfile import tempfile
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from pathlib import Path from pathlib import Path
from typing import Literal
from urllib.parse import urlparse from urllib.parse import urlparse
import httpx import httpx
@ -16,6 +15,7 @@ from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.settings import SettingsRepository
from haiku.rag.utils import text_to_docling_document from haiku.rag.utils import text_to_docling_document
@ -24,17 +24,15 @@ class HaikuRAG:
def __init__( def __init__(
self, self,
db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR db_path: Path = Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
/ "haiku.rag.sqlite",
skip_validation: bool = False, skip_validation: bool = False,
): ):
"""Initialize the RAG client with a database path. """Initialize the RAG client with a database path.
Args: Args:
db_path: Path to the SQLite database file or ":memory:" for in-memory database. db_path: Path to the database file.
skip_validation: Whether to skip configuration validation on database load. skip_validation: Whether to skip configuration validation on database load.
""" """
if isinstance(db_path, Path):
if not db_path.parent.exists(): if not db_path.parent.exists():
Path.mkdir(db_path.parent, parents=True) Path.mkdir(db_path.parent, parents=True)
self.store = Store(db_path, skip_validation=skip_validation) self.store = Store(db_path, skip_validation=skip_validation)
@ -269,7 +267,7 @@ class HaikuRAG:
# Default to .html for web content # Default to .html for web content
return ".html" return ".html"
async def get_document_by_id(self, document_id: int) -> Document | None: async def get_document_by_id(self, document_id: str) -> Document | None:
"""Get a document by its ID. """Get a document by its ID.
Args: Args:
@ -300,7 +298,7 @@ class HaikuRAG:
document, docling_document document, docling_document
) )
async def delete_document(self, document_id: int) -> bool: async def delete_document(self, document_id: str) -> bool:
"""Delete a document by its ID.""" """Delete a document by its ID."""
return await self.document_repository.delete(document_id) return await self.document_repository.delete(document_id)
@ -319,14 +317,14 @@ class HaikuRAG:
return await self.document_repository.list_all(limit=limit, offset=offset) return await self.document_repository.list_all(limit=limit, offset=offset)
async def search( async def search(
self, query: str, limit: int = 5, k: int = 60 self, query: str, limit: int = 5, search_type: str = "hybrid"
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking. """Search for relevant chunks using the specified search method with optional reranking.
Args: Args:
query: The search query string. query: The search query string.
limit: Maximum number of results to return. limit: Maximum number of results to return.
k: Parameter for Reciprocal Rank Fusion (default: 60). search_type: Type of search - "vector", "fts", or "hybrid" (default).
Returns: Returns:
List of (chunk, score) tuples ordered by relevance. List of (chunk, score) tuples ordered by relevance.
@ -335,12 +333,15 @@ class HaikuRAG:
reranker = get_reranker() reranker = get_reranker()
if reranker is None: if reranker is None:
return await self.chunk_repository.search_chunks_hybrid(query, limit, k) # No reranking - return direct search results
return await self.chunk_repository.search(query, limit, search_type)
# Get more initial results (3X) for reranking # Get more initial results (3X) for reranking
search_results = await self.chunk_repository.search_chunks_hybrid( search_limit = limit * 3
query, limit * 3, k search_results = await self.chunk_repository.search(
query, search_limit, search_type
) )
# Apply reranking # Apply reranking
chunks = [chunk for chunk, _ in search_results] chunks = [chunk for chunk, _ in search_results]
reranked_results = await reranker.rerank(query, chunks, top_n=limit) reranked_results = await reranker.rerank(query, chunks, top_n=limit)
@ -493,7 +494,7 @@ class HaikuRAG:
qa_agent = get_qa_agent(self, use_citations=cite) qa_agent = get_qa_agent(self, use_citations=cite)
return await qa_agent.answer(question) return await qa_agent.answer(question)
async def rebuild_database(self) -> AsyncGenerator[int, None]: async def rebuild_database(self) -> AsyncGenerator[str, None]:
"""Rebuild the database by deleting all chunks and re-indexing all documents. """Rebuild the database by deleting all chunks and re-indexing all documents.
For documents with URIs: For documents with URIs:
@ -510,10 +511,8 @@ class HaikuRAG:
self.store.recreate_embeddings_table() self.store.recreate_embeddings_table()
# Update settings to current config # Update settings to current config
from haiku.rag.store.repositories.settings import SettingsRepository
settings_repo = SettingsRepository(self.store) settings_repo = SettingsRepository(self.store)
settings_repo.save() settings_repo.save_current_settings()
documents = await self.list_documents() documents = await self.list_documents()
@ -547,13 +546,10 @@ class HaikuRAG:
# Document without URI - re-create chunks from existing content # Document without URI - re-create chunks from existing content
docling_document = text_to_docling_document(doc.content) docling_document = text_to_docling_document(doc.content)
await self.chunk_repository.create_chunks_for_document( await self.chunk_repository.create_chunks_for_document(
doc.id, docling_document, commit=False doc.id, docling_document
) )
yield doc.id yield doc.id
if self.store._connection:
self.store._connection.commit()
def close(self): def close(self):
"""Close the underlying store connection.""" """Close the underlying store connection."""
self.store.close() self.store.close()

View file

@ -12,6 +12,10 @@ load_dotenv()
class AppConfig(BaseModel): class AppConfig(BaseModel):
ENV: str = "production" ENV: str = "production"
LANCEDB_API_KEY: str = ""
LANCEDB_URI: str = ""
LANCEDB_REGION: str = ""
DEFAULT_DATA_DIR: Path = get_default_data_dir() DEFAULT_DATA_DIR: Path = get_default_data_dir()
MONITOR_DIRECTORIES: list[Path] = [] MONITOR_DIRECTORIES: list[Path] = []
@ -19,8 +23,8 @@ class AppConfig(BaseModel):
EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_MODEL: str = "mxbai-embed-large"
EMBEDDINGS_VECTOR_DIM: int = 1024 EMBEDDINGS_VECTOR_DIM: int = 1024
RERANK_PROVIDER: str = "ollama" RERANK_PROVIDER: str = ""
RERANK_MODEL: str = "qwen3" RERANK_MODEL: str = ""
QA_PROVIDER: str = "ollama" QA_PROVIDER: str = "ollama"
QA_MODEL: str = "qwen3" QA_MODEL: str = "qwen3"

View file

@ -1,5 +1,5 @@
from pathlib import Path from pathlib import Path
from typing import Any, Literal from typing import Any
from fastmcp import FastMCP from fastmcp import FastMCP
from pydantic import BaseModel from pydantic import BaseModel
@ -8,13 +8,13 @@ from haiku.rag.client import HaikuRAG
class SearchResult(BaseModel): class SearchResult(BaseModel):
document_id: int document_id: str
content: str content: str
score: float score: float
class DocumentResult(BaseModel): class DocumentResult(BaseModel):
id: int | None id: str | None
content: str content: str
uri: str | None = None uri: str | None = None
metadata: dict[str, Any] = {} metadata: dict[str, Any] = {}
@ -22,14 +22,14 @@ class DocumentResult(BaseModel):
updated_at: str updated_at: str
def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP: def create_mcp_server(db_path: Path) -> FastMCP:
"""Create an MCP server with the specified database path.""" """Create an MCP server with the specified database path."""
mcp = FastMCP("haiku-rag") mcp = FastMCP("haiku-rag")
@mcp.tool() @mcp.tool()
async def add_document_from_file( async def add_document_from_file(
file_path: str, metadata: dict[str, Any] | None = None file_path: str, metadata: dict[str, Any] | None = None
) -> int | None: ) -> str | None:
"""Add a document to the RAG system from a file path.""" """Add a document to the RAG system from a file path."""
try: try:
async with HaikuRAG(db_path) as rag: async with HaikuRAG(db_path) as rag:
@ -43,7 +43,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
@mcp.tool() @mcp.tool()
async def add_document_from_url( async def add_document_from_url(
url: str, metadata: dict[str, Any] | None = None url: str, metadata: dict[str, Any] | None = None
) -> int | None: ) -> str | None:
"""Add a document to the RAG system from a URL.""" """Add a document to the RAG system from a URL."""
try: try:
async with HaikuRAG(db_path) as rag: async with HaikuRAG(db_path) as rag:
@ -55,7 +55,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
@mcp.tool() @mcp.tool()
async def add_document_from_text( async def add_document_from_text(
content: str, uri: str | None = None, metadata: dict[str, Any] | None = None content: str, uri: str | None = None, metadata: dict[str, Any] | None = None
) -> int | None: ) -> str | None:
"""Add a document to the RAG system from text content.""" """Add a document to the RAG system from text content."""
try: try:
async with HaikuRAG(db_path) as rag: async with HaikuRAG(db_path) as rag:
@ -73,6 +73,9 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
search_results = [] search_results = []
for chunk, score in results: for chunk, score in results:
assert chunk.document_id is not None, (
"Chunk document_id should not be None in search results"
)
search_results.append( search_results.append(
SearchResult( SearchResult(
document_id=chunk.document_id, document_id=chunk.document_id,
@ -86,7 +89,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
return [] return []
@mcp.tool() @mcp.tool()
async def get_document(document_id: int) -> DocumentResult | None: async def get_document(document_id: str) -> DocumentResult | None:
"""Get a document by its ID.""" """Get a document by its ID."""
try: try:
async with HaikuRAG(db_path) as rag: async with HaikuRAG(db_path) as rag:
@ -130,7 +133,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
return [] return []
@mcp.tool() @mcp.tool()
async def delete_document(document_id: int) -> bool: async def delete_document(document_id: str) -> bool:
"""Delete a document by its ID.""" """Delete a document by its ID."""
try: try:
async with HaikuRAG(db_path) as rag: async with HaikuRAG(db_path) as rag:

316
src/haiku/rag/migration.py Normal file
View file

@ -0,0 +1,316 @@
#!/usr/bin/env python3
"""
Migration script to migrate from SQLite to LanceDB.
This script will:
1. Read data from an existing SQLite database
2. Create a new LanceDB database with the same data
3. Preserve all documents, chunks, embeddings, and settings
"""
import json
import sqlite3
import struct
from pathlib import Path
from uuid import uuid4
from rich.console import Console
from rich.progress import Progress, TaskID
from haiku.rag.store.engine import Store
def deserialize_sqlite_embedding(data: bytes) -> list[float]:
"""Deserialize sqlite-vec embedding from bytes."""
if not data:
return []
# sqlite-vec stores embeddings as float32 arrays
num_floats = len(data) // 4
return list(struct.unpack(f"{num_floats}f", data))
class SQLiteToLanceDBMigrator:
"""Migrates data from SQLite to LanceDB."""
def __init__(self, sqlite_path: Path, lancedb_path: Path):
self.sqlite_path = sqlite_path
self.lancedb_path = lancedb_path
self.console = Console()
def migrate(self) -> bool:
"""Perform the migration."""
try:
self.console.print(
f"[blue]Starting migration from {self.sqlite_path} to {self.lancedb_path}[/blue]"
)
# Check if SQLite database exists
if not self.sqlite_path.exists():
self.console.print(
f"[red]SQLite database not found: {self.sqlite_path}[/red]"
)
return False
# Connect to SQLite database
sqlite_conn = sqlite3.connect(self.sqlite_path)
sqlite_conn.row_factory = sqlite3.Row
# Create LanceDB store
lance_store = Store(self.lancedb_path, skip_validation=True)
with Progress() as progress:
# Migrate documents
doc_task = progress.add_task(
"[green]Migrating documents...", total=None
)
document_id_mapping = self._migrate_documents(
sqlite_conn, lance_store, progress, doc_task
)
# Migrate chunks and embeddings
chunk_task = progress.add_task(
"[yellow]Migrating chunks and embeddings...", total=None
)
self._migrate_chunks(
sqlite_conn, lance_store, progress, chunk_task, document_id_mapping
)
# Migrate settings
settings_task = progress.add_task(
"[blue]Migrating settings...", total=None
)
self._migrate_settings(
sqlite_conn, lance_store, progress, settings_task
)
sqlite_conn.close()
# Optimize the chunks table after migration
self.console.print("[blue]Optimizing LanceDB...[/blue]")
try:
lance_store.chunks_table.optimize()
self.console.print("[green]✅ Optimization completed[/green]")
except Exception as e:
self.console.print(
f"[yellow]Warning: Optimization failed: {e}[/yellow]"
)
lance_store.close()
self.console.print("[green]✅ Migration completed successfully![/green]")
self.console.print(
f"[green]✅ Migrated {len(document_id_mapping)} documents[/green]"
)
return True
except Exception as e:
self.console.print(f"[red]❌ Migration failed: {e}[/red]")
import traceback
self.console.print(f"[red]{traceback.format_exc()}[/red]")
return False
def _migrate_documents(
self,
sqlite_conn: sqlite3.Connection,
lance_store: Store,
progress: Progress,
task: TaskID,
) -> dict[int, str]:
"""Migrate documents from SQLite to LanceDB and return ID mapping."""
cursor = sqlite_conn.cursor()
cursor.execute(
"SELECT id, content, uri, metadata, created_at, updated_at FROM documents ORDER BY id"
)
documents = []
id_mapping = {} # Maps old integer ID to new UUID
for row in cursor.fetchall():
new_uuid = str(uuid4())
id_mapping[row["id"]] = new_uuid
doc_data = {
"id": new_uuid,
"content": row["content"],
"uri": row["uri"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"created_at": row["created_at"],
"updated_at": row["updated_at"],
}
documents.append(doc_data)
# Batch insert documents to LanceDB
if documents:
from haiku.rag.store.engine import DocumentRecord
doc_records = [
DocumentRecord(
id=doc["id"],
content=doc["content"],
uri=doc["uri"],
metadata=json.dumps(doc["metadata"]),
created_at=doc["created_at"],
updated_at=doc["updated_at"],
)
for doc in documents
]
lance_store.documents_table.add(doc_records)
progress.update(task, completed=len(documents), total=len(documents))
return id_mapping
def _migrate_chunks(
self,
sqlite_conn: sqlite3.Connection,
lance_store: Store,
progress: Progress,
task: TaskID,
document_id_mapping: dict[int, str],
):
"""Migrate chunks and embeddings from SQLite to LanceDB."""
cursor = sqlite_conn.cursor()
# Get chunks first
cursor.execute("""
SELECT id, document_id, content, metadata
FROM chunks
ORDER BY id
""")
chunks_data = cursor.fetchall()
# Get embeddings separately to avoid vec0 virtual table issues
embeddings_map = {}
try:
# Try to get embeddings from the vec0 tables directly
cursor.execute("""
SELECT
r.chunk_id,
v.vectors
FROM chunk_embeddings_rowids r
JOIN chunk_embeddings_vector_chunks00 v ON r.rowid = v.rowid
""")
for row in cursor.fetchall():
chunk_id = row[0]
vectors_blob = row[1]
if vectors_blob and chunk_id not in embeddings_map:
embeddings_map[chunk_id] = vectors_blob
except sqlite3.OperationalError as e:
self.console.print(
f"[yellow]Warning: Could not extract embeddings: {e}[/yellow]"
)
self.console.print(
"[yellow]Continuing migration without embeddings...[/yellow]"
)
chunks = []
for row in chunks_data:
# Generate new UUID for chunk
chunk_uuid = str(uuid4())
# Map the old document_id to new UUID
document_uuid = document_id_mapping.get(row["document_id"])
if not document_uuid:
self.console.print(
f"[yellow]Warning: Document ID {row['document_id']} not found in mapping for chunk {row['id']}[/yellow]"
)
continue
# Get embedding for this chunk
embedding = []
embedding_blob = embeddings_map.get(row["id"])
if embedding_blob:
try:
embedding = deserialize_sqlite_embedding(embedding_blob)
except Exception as e:
self.console.print(
f"[yellow]Warning: Failed to deserialize embedding for chunk {row['id']}: {e}[/yellow]"
)
# Generate a zero vector of the expected dimension
embedding = [0.0] * lance_store.embedder._vector_dim
else:
# No embedding found, generate zero vector
embedding = [0.0] * lance_store.embedder._vector_dim
chunk_data = {
"id": chunk_uuid,
"document_id": document_uuid,
"content": row["content"],
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
"vector": embedding,
}
chunks.append(chunk_data)
# Batch insert chunks to LanceDB
if chunks:
chunk_records = [
lance_store.ChunkRecord(
id=chunk["id"],
document_id=chunk["document_id"],
content=chunk["content"],
metadata=json.dumps(chunk["metadata"]),
vector=chunk["vector"],
)
for chunk in chunks
]
lance_store.chunks_table.add(chunk_records)
progress.update(task, completed=len(chunks), total=len(chunks))
def _migrate_settings(
self,
sqlite_conn: sqlite3.Connection,
lance_store: Store,
progress: Progress,
task: TaskID,
):
"""Migrate settings from SQLite to LanceDB."""
cursor = sqlite_conn.cursor()
try:
cursor.execute("SELECT id, settings FROM settings WHERE id = 1")
row = cursor.fetchone()
if row:
settings_data = json.loads(row["settings"]) if row["settings"] else {}
# Update the existing settings in LanceDB (use string ID)
lance_store.settings_table.update(
where="id = 'settings'",
values={"settings": json.dumps(settings_data)},
)
progress.update(task, completed=1, total=1)
else:
progress.update(task, completed=0, total=0)
except sqlite3.OperationalError:
# Settings table doesn't exist in old SQLite database
self.console.print(
"[yellow]No settings table found in SQLite database[/yellow]"
)
progress.update(task, completed=0, total=0)
async def migrate_sqlite_to_lancedb(
sqlite_path: Path, lancedb_path: Path | None = None
) -> bool:
"""
Migrate an existing SQLite database to LanceDB.
Args:
sqlite_path: Path to the existing SQLite database
lancedb_path: Path for the new LanceDB database (optional, will auto-generate if not provided)
Returns:
True if migration was successful, False otherwise
"""
if lancedb_path is None:
# Auto-generate LanceDB path
lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
return migrator.migrate()

View file

@ -31,10 +31,4 @@ def get_reranker() -> RerankerBase | None:
except ImportError: except ImportError:
return None return None
if Config.RERANK_PROVIDER == "ollama":
from haiku.rag.reranking.ollama import OllamaReranker
_reranker = OllamaReranker()
return _reranker
return None return None

View file

@ -1,81 +0,0 @@
from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.ollama import OllamaProvider
from haiku.rag.config import Config
from haiku.rag.reranking.base import RerankerBase
from haiku.rag.store.models.chunk import Chunk
class RerankResult(BaseModel):
"""Individual rerank result with index and relevance score."""
index: int
relevance_score: float
class RerankResponse(BaseModel):
"""Response from the reranking model containing ranked results."""
results: list[RerankResult]
class OllamaReranker(RerankerBase):
def __init__(self, model: str = Config.RERANK_MODEL):
self._model = model
# Create the reranking prompt
system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query.
Return your response as a JSON object with a "results" array. Each result should have:
- "index": the original index of the document (integer)
- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant)
Only return the top documents up to the requested limit, ordered by decreasing relevance score.
/no_think
"""
model_obj = OpenAIModel(
model_name=model,
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
)
self._agent = Agent(
model=model_obj,
output_type=RerankResponse,
system_prompt=system_prompt,
)
async def rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
documents = []
for i, chunk in enumerate(chunks):
documents.append({"index": i, "content": chunk.content})
documents_text = ""
for doc in documents:
documents_text += f"Index {doc['index']}: {doc['content']}\n\n"
user_prompt = f"""Query: {query}
Documents to rerank:
{documents_text.strip()}
Rank these documents by relevance to the query and return the top {top_n} results as JSON."""
try:
result = await self._agent.run(user_prompt)
return [
(chunks[result_item.index], result_item.relevance_score)
for result_item in result.output.results[:top_n]
]
except Exception:
# Fallback: return chunks in original order with same score
return [(chunks[i], 1.0) for i in range(min(top_n, len(chunks)))]

View file

@ -1,171 +1,203 @@
import sqlite3 import json
import struct import logging
from importlib import metadata from importlib import metadata
from pathlib import Path from pathlib import Path
from typing import Literal from uuid import uuid4
import sqlite_vec import lancedb
from packaging.version import parse from lancedb.pydantic import LanceModel, Vector
from rich.console import Console from pydantic import Field
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.embeddings import get_embedder from haiku.rag.embeddings import get_embedder
from haiku.rag.store.upgrades import upgrades
from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int logger = logging.getLogger(__name__)
class DocumentRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4()))
content: str
uri: str | None = None
metadata: str = Field(default="{}")
created_at: str = Field(default_factory=lambda: "")
updated_at: str = Field(default_factory=lambda: "")
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 ChunkRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4()))
document_id: str
content: str
metadata: str = Field(default="{}")
vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore
return ChunkRecord
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
settings: str = Field(default="{}")
class Store: class Store:
def __init__( def __init__(self, db_path: Path, skip_validation: bool = False):
self, db_path: Path | Literal[":memory:"], skip_validation: bool = False self.db_path: Path = db_path
): self.embedder = get_embedder()
self.db_path: Path | Literal[":memory:"] = db_path
# Create the ChunkRecord model with the correct vector dimension
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
# Connect to LanceDB
self.db = self._connect_to_lancedb(db_path)
# Initialize tables
self.create_or_update_db() self.create_or_update_db()
# Validate config compatibility after connection is established # Validate config compatibility after connection is established
if not skip_validation: if not skip_validation:
self._validate_configuration()
def _connect_to_lancedb(self, db_path: Path):
"""Establish connection to LanceDB (local, cloud, or object storage)."""
# Check if we have cloud configuration
if self._has_cloud_config():
return lancedb.connect(
uri=Config.LANCEDB_URI,
api_key=Config.LANCEDB_API_KEY,
region=Config.LANCEDB_REGION,
)
else:
# Local file system connection
return lancedb.connect(db_path)
def _has_cloud_config(self) -> bool:
"""Check if cloud configuration is complete."""
return bool(
Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION
)
def _validate_configuration(self) -> None:
"""Validate that the configuration is compatible with the database."""
from haiku.rag.store.repositories.settings import SettingsRepository from haiku.rag.store.repositories.settings import SettingsRepository
settings_repo = SettingsRepository(self) settings_repo = SettingsRepository(self)
settings_repo.validate_config_compatibility() settings_repo.validate_config_compatibility()
current_version = metadata.version("haiku.rag")
self.set_user_version(current_version)
def create_or_update_db(self): def create_or_update_db(self):
"""Create the database and tables with sqlite-vec support for embeddings.""" """Create the database tables."""
current_version = metadata.version("haiku.rag")
db = sqlite3.connect(self.db_path) # Get list of existing tables
db.enable_load_extension(True) existing_tables = self.db.table_names()
sqlite_vec.load(db)
# Enable WAL mode for better concurrency (skip for in-memory databases) # Create or get documents table
if self.db_path != ":memory:": if "documents" in existing_tables:
db.execute("PRAGMA journal_mode=WAL") self.documents_table = self.db.open_table("documents")
else:
self.documents_table = self.db.create_table(
"documents", schema=DocumentRecord
)
self._connection = db # Create or get chunks table
existing_tables = [ if "chunks" in existing_tables:
row[0] self.chunks_table = self.db.open_table("chunks")
for row in db.execute( else:
"SELECT name FROM sqlite_master WHERE type='table';" self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
).fetchall() # Create FTS index on the new table
] self.chunks_table.create_fts_index("content", replace=True)
# If we have a db already, perform upgrades and return # Create or get settings table
if self.db_path != ":memory:" and "documents" in existing_tables: if "settings" in existing_tables:
# Upgrade database self.settings_table = self.db.open_table("settings")
console = Console() else:
db_version = self.get_user_version() self.settings_table = self.db.create_table(
for version, steps in upgrades: "settings", schema=SettingsRecord
if parse(current_version) >= parse(version) and parse(version) > parse(
db_version
):
for step in steps:
step(db)
console.print(
f"[green][b]DB Upgrade: [/b]{step.__doc__}[/green]"
) )
return
# Create documents table
db.execute("""
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
content TEXT NOT NULL,
uri TEXT,
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 = get_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 FTS5 table for full-text search
db.execute("""
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
content,
content='chunks',
content_rowid='id'
)
""")
# Create settings table for storing current configuration
db.execute("""
CREATE TABLE IF NOT EXISTS settings (
id INTEGER PRIMARY KEY DEFAULT 1,
settings TEXT NOT NULL DEFAULT '{}'
)
""")
# Save current settings to the new database # Save current settings to the new database
settings_json = Config.model_dump_json() settings_data = Config.model_dump(mode="json")
db.execute( self.settings_table.add(
"INSERT OR IGNORE INTO settings (id, settings) VALUES (1, ?)", [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
(settings_json,),
) )
# Create indexes for better performance
db.execute( # Set current version in settings
"CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)" current_version = metadata.version("haiku.rag")
self.set_haiku_version(current_version)
# Check if we need to perform upgrades
try:
existing_settings = list(
self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
) )
db.commit() if existing_settings:
db_version = self.get_haiku_version() # noqa: F841
# TODO: Add upgrade logic here similar to SQLite version when needed
except Exception:
# Settings table might not exist yet in fresh databases
pass
def get_user_version(self) -> str: def get_haiku_version(self) -> str:
"""Returns the SQLite user version""" """Returns the user version stored in settings."""
if self._connection is None: settings_records = list(
raise ValueError("Store connection is not available") self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
)
if settings_records:
settings = (
json.loads(settings_records[0].settings)
if settings_records[0].settings
else {}
)
return settings.get("version", "0.0.0")
return "0.0.0"
cursor = self._connection.execute("PRAGMA user_version;") def set_haiku_version(self, version: str) -> None:
version = cursor.fetchone() """Updates the user version in settings."""
return int_to_semantic_version(version[0]) settings_records = list(
self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
def set_user_version(self, version: str) -> None: )
"""Updates the SQLite user version""" if settings_records:
if self._connection is None: settings = (
raise ValueError("Store connection is not available") json.loads(settings_records[0].settings)
if settings_records[0].settings
self._connection.execute( else {}
f"PRAGMA user_version = {semantic_version_to_int(version)};" )
settings["version"] = version
# Update the record
self.settings_table.update(
where="id = 'settings'", values={"settings": json.dumps(settings)}
)
else:
# Create new settings record
settings_data = Config.model_dump(mode="json")
settings_data["version"] = version
self.settings_table.add(
[SettingsRecord(id="settings", settings=json.dumps(settings_data))]
) )
def recreate_embeddings_table(self) -> None: def recreate_embeddings_table(self) -> None:
"""Recreate the embeddings table with current vector dimensions.""" """Recreate the chunks table with current vector dimensions."""
if self._connection is None: # Drop and recreate chunks table
raise ValueError("Store connection is not available") try:
self.db.drop_table("chunks")
except Exception:
pass
# Drop existing embeddings table # Update the ChunkRecord model with new vector dimension
self._connection.execute("DROP TABLE IF EXISTS chunk_embeddings") self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
# Recreate with current dimensions # Create FTS index on the new table
embedder = get_embedder() self.chunks_table.create_fts_index("content", replace=True)
self._connection.execute(f"""
CREATE VIRTUAL TABLE chunk_embeddings USING vec0(
chunk_id INTEGER PRIMARY KEY,
embedding FLOAT[{embedder._vector_dim}]
)
""")
self._connection.commit()
@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): def close(self):
"""Close the database connection if it's an in-memory database.""" """Close the database connection."""
if self._connection is not None: # LanceDB connections are automatically managed
self._connection.close() pass
self._connection = None
@property
def _connection(self):
"""Compatibility property for repositories expecting _connection."""
return self

View file

@ -6,8 +6,8 @@ class Chunk(BaseModel):
Represents a chunk with content, metadata, and optional document information. Represents a chunk with content, metadata, and optional document information.
""" """
id: int | None = None id: str | None = None
document_id: int | None = None document_id: str | None = None
content: str content: str
metadata: dict = {} metadata: dict = {}
document_uri: str | None = None document_uri: str | None = None

View file

@ -8,7 +8,7 @@ class Document(BaseModel):
Represents a document with an ID, content, and metadata. Represents a document with an ID, content, and metadata.
""" """
id: int | None = None id: str | None = None
content: str content: str
uri: str | None = None uri: str | None = None
metadata: dict = {} metadata: dict = {}

View file

@ -1,5 +1,9 @@
from haiku.rag.store.repositories.base import BaseRepository
from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.settings import SettingsRepository
__all__ = ["BaseRepository", "DocumentRepository", "ChunkRepository"] __all__ = [
"ChunkRepository",
"DocumentRepository",
"SettingsRepository",
]

View file

@ -1,40 +0,0 @@
from abc import ABC, abstractmethod
from typing import Generic, TypeVar
from haiku.rag.store.engine import Store
T = TypeVar("T")
class BaseRepository(ABC, Generic[T]):
"""Base repository interface for database operations."""
def __init__(self, store: Store):
self.store = store
@abstractmethod
async def create(self, entity: T) -> T:
"""Create a new entity in the database."""
pass
@abstractmethod
async def get_by_id(self, entity_id: int) -> T | None:
"""Get an entity by its ID."""
pass
@abstractmethod
async def update(self, entity: T) -> T:
"""Update an existing entity."""
pass
@abstractmethod
async def delete(self, entity_id: int) -> bool:
"""Delete an entity by its ID."""
pass
@abstractmethod
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[T]:
"""List all entities with optional pagination."""
pass

View file

@ -1,516 +1,381 @@
import asyncio
import json import json
import re import logging
from uuid import uuid4
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
from lancedb.rerankers import RRFReranker
from haiku.rag.chunker import chunker from haiku.rag.chunker import chunker
from haiku.rag.config import Config
from haiku.rag.embeddings import get_embedder from haiku.rag.embeddings import get_embedder
from haiku.rag.store.engine import DocumentRecord, Store
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.repositories.base import BaseRepository
logger = logging.getLogger(__name__)
class ChunkRepository(BaseRepository[Chunk]): class ChunkRepository:
"""Repository for Chunk database operations.""" """Repository for Chunk operations."""
def __init__(self, store): def __init__(self, store: Store) -> None:
super().__init__(store) self.store = store
self.embedder = get_embedder() self.embedder = get_embedder()
self._optimize_lock = asyncio.Lock()
async def create(self, entity: Chunk, commit: bool = True) -> Chunk: def _ensure_fts_index(self) -> None:
"""Ensure FTS index exists on the content column."""
try:
self.store.chunks_table.create_fts_index("content", replace=True)
except Exception as e:
# Log the error but don't fail - FTS might already exist
logger.debug(f"FTS index creation skipped: {e}")
async def _optimize(self) -> None:
"""Optimize the chunks table to refresh indexes."""
# Skip optimization for LanceDB Cloud as it handles this automatically
if Config.LANCEDB_URI and Config.LANCEDB_URI.startswith("db://"):
return
async with self._optimize_lock:
try:
self.store.chunks_table.optimize()
except (RuntimeError, OSError) as e:
# Handle "too many open files" and other resource errors gracefully
logger.debug(
f"Table optimization skipped due to resource constraints: {e}"
)
async def create(self, entity: Chunk) -> Chunk:
"""Create a chunk in the database.""" """Create a chunk in the database."""
if self.store._connection is None: assert entity.document_id, "Chunk must have a document_id to be created"
raise ValueError("Store connection is not available")
if entity.document_id is None:
raise ValueError("Chunk must have a document_id to be created")
cursor = self.store._connection.cursor() chunk_id = str(uuid4())
cursor.execute(
"""
INSERT INTO chunks (document_id, content, metadata)
VALUES (:document_id, :content, :metadata)
""",
{
"document_id": entity.document_id,
"content": entity.content,
"metadata": json.dumps(entity.metadata),
},
)
entity.id = cursor.lastrowid # Generate embedding if not provided
# Generate and store embedding - use existing one if provided
if entity.embedding is not None: if entity.embedding is not None:
# Use the provided embedding embedding = entity.embedding
serialized_embedding = self.store.serialize_embedding(entity.embedding)
else: else:
# Generate embedding from content
embedding = await self.embedder.embed(entity.content) embedding = await self.embedder.embed(entity.content)
serialized_embedding = self.store.serialize_embedding(embedding) chunk_record = self.store.ChunkRecord(
id=chunk_id,
cursor.execute( document_id=entity.document_id,
""" content=entity.content,
INSERT INTO chunk_embeddings (chunk_id, embedding) metadata=json.dumps(entity.metadata),
VALUES (:chunk_id, :embedding) vector=embedding,
""",
{"chunk_id": entity.id, "embedding": serialized_embedding},
) )
# Insert into FTS5 table for full-text search self.store.chunks_table.add([chunk_record])
cursor.execute(
""" entity.id = chunk_id
INSERT INTO chunks_fts(rowid, content)
VALUES (:rowid, :content) # Try to optimize if not currently locked (non-blocking)
""", if not self._optimize_lock.locked():
{"rowid": entity.id, "content": entity.content}, asyncio.create_task(self._optimize())
)
if commit:
self.store._connection.commit()
return entity return entity
async def get_by_id(self, entity_id: int) -> 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."""
if self.store._connection is None: results = list(
raise ValueError("Store connection is not available") self.store.chunks_table.search()
.where(f"id = '{entity_id}'")
cursor = self.store._connection.cursor() .limit(1)
cursor.execute( .to_pydantic(self.store.ChunkRecord)
"""
SELECT id, document_id, content, metadata
FROM chunks WHERE id = :id
""",
{"id": entity_id},
) )
row = cursor.fetchone() if not results:
if row is None:
return None return None
chunk_id, document_id, content, metadata_json = row chunk_record = results[0]
metadata = json.loads(metadata_json) if metadata_json else {}
return Chunk( return Chunk(
id=chunk_id, document_id=document_id, content=content, metadata=metadata id=chunk_record.id,
document_id=chunk_record.document_id,
content=chunk_record.content,
metadata=json.loads(chunk_record.metadata) if chunk_record.metadata else {},
) )
async def update(self, entity: Chunk) -> Chunk: async def update(self, entity: Chunk) -> Chunk:
"""Update an existing chunk.""" """Update an existing chunk."""
if self.store._connection is None: assert entity.id, "Chunk ID is required for update"
raise ValueError("Store connection is not available")
if entity.id is None:
raise ValueError("Chunk ID is required for update")
cursor = self.store._connection.cursor() embedding = await self.embedder.embed(entity.content)
cursor.execute(
""" self.store.chunks_table.update(
UPDATE chunks where=f"id = '{entity.id}'",
SET document_id = :document_id, content = :content, metadata = :metadata values={
WHERE id = :id
""",
{
"document_id": entity.document_id, "document_id": entity.document_id,
"content": entity.content, "content": entity.content,
"metadata": json.dumps(entity.metadata), "metadata": json.dumps(entity.metadata),
"id": entity.id, "vector": embedding,
}, },
) )
# Try to optimize if not currently locked (non-blocking)
if not self._optimize_lock.locked():
asyncio.create_task(self._optimize())
# Regenerate and update embedding
embedding = await self.embedder.embed(entity.content)
serialized_embedding = self.store.serialize_embedding(embedding)
cursor.execute(
"""
UPDATE chunk_embeddings
SET embedding = :embedding
WHERE chunk_id = :chunk_id
""",
{"embedding": serialized_embedding, "chunk_id": entity.id},
)
# Update FTS5 table
cursor.execute(
"""
UPDATE chunks_fts
SET content = :content
WHERE rowid = :rowid
""",
{"content": entity.content, "rowid": entity.id},
)
self.store._connection.commit()
return entity return entity
async def delete(self, entity_id: int, commit: bool = True) -> bool: async def delete(self, entity_id: str) -> bool:
"""Delete a chunk by its ID.""" """Delete a chunk by its ID."""
if self.store._connection is None: chunk = await self.get_by_id(entity_id)
raise ValueError("Store connection is not available") if chunk is None:
return False
cursor = self.store._connection.cursor() self.store.chunks_table.delete(f"id = '{entity_id}'")
return True
# Delete from FTS5 table first
cursor.execute(
"DELETE FROM chunks_fts WHERE rowid = :rowid", {"rowid": entity_id}
)
# Delete the embedding
cursor.execute(
"DELETE FROM chunk_embeddings WHERE chunk_id = :chunk_id",
{"chunk_id": entity_id},
)
# Delete the chunk
cursor.execute("DELETE FROM chunks WHERE id = :id", {"id": entity_id})
deleted = cursor.rowcount > 0
if commit:
self.store._connection.commit()
return deleted
async def list_all( async def list_all(
self, limit: int | None = None, offset: int | None = None self, limit: int | None = None, offset: int | None = None
) -> list[Chunk]: ) -> list[Chunk]:
"""List all chunks with optional pagination.""" """List all chunks with optional pagination."""
if self.store._connection is None: query = self.store.chunks_table.search()
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
query = "SELECT id, document_id, content, metadata FROM chunks ORDER BY document_id, id"
params = {}
if limit is not None:
query += " LIMIT :limit"
params["limit"] = limit
if offset is not None: if offset is not None:
query += " OFFSET :offset" query = query.offset(offset)
params["offset"] = offset if limit is not None:
query = query.limit(limit)
cursor.execute(query, params) results = list(query.to_pydantic(self.store.ChunkRecord))
rows = cursor.fetchall()
return [ return [
Chunk( Chunk(
id=chunk_id, id=chunk.id,
document_id=document_id, document_id=chunk.document_id,
content=content, content=chunk.content,
metadata=json.loads(metadata_json) if metadata_json else {}, metadata=json.loads(chunk.metadata) if chunk.metadata else {},
) )
for chunk_id, document_id, content, metadata_json in rows for chunk in results
] ]
async def create_chunks_for_document( async def create_chunks_for_document(
self, document_id: int, document: DoclingDocument, commit: bool = True self, document_id: str, document: DoclingDocument
) -> list[Chunk]: ) -> list[Chunk]:
"""Create chunks and embeddings for a document from DoclingDocument.""" """Create chunks and embeddings for a document from DoclingDocument."""
# Chunk the document content
chunk_texts = await chunker.chunk(document) chunk_texts = await chunker.chunk(document)
# Generate embeddings in parallel for all chunks
embeddings_tasks = []
for chunk_text in chunk_texts:
embeddings_tasks.append(self.embedder.embed(chunk_text))
# Wait for all embeddings to complete
embeddings = await asyncio.gather(*embeddings_tasks)
# Prepare all chunk records for batch insertion
chunk_records = []
created_chunks = [] created_chunks = []
# Create chunks with embeddings using the create method for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)):
for order, chunk_text in enumerate(chunk_texts): chunk_id = str(uuid4())
# Create chunk with order in metadata
chunk = Chunk( chunk_record = self.store.ChunkRecord(
document_id=document_id, content=chunk_text, metadata={"order": order} id=chunk_id,
document_id=document_id,
content=chunk_text,
metadata=json.dumps({"order": order}),
vector=embedding,
) )
chunk_records.append(chunk_record)
created_chunk = await self.create(chunk, commit=commit) chunk = Chunk(
created_chunks.append(created_chunk) id=chunk_id,
document_id=document_id,
content=chunk_text,
metadata={"order": order},
)
created_chunks.append(chunk)
# Batch insert all chunks at once
if chunk_records:
self.store.chunks_table.add(chunk_records)
# Force optimization once at the end for bulk operations
await self._optimize()
return created_chunks return created_chunks
async def delete_all(self, commit: bool = True) -> bool: async def delete_all(self) -> None:
"""Delete all chunks from the database.""" """Delete all chunks from the database."""
if self.store._connection is None: # Drop and recreate table to clear all data
raise ValueError("Store connection is not available") self.store.db.drop_table("chunks")
self.store.chunks_table = self.store.db.create_table(
"chunks", schema=self.store.ChunkRecord
)
# Create FTS index on the new table
self.store.chunks_table.create_fts_index("content", replace=True)
cursor = self.store._connection.cursor() async def delete_by_document_id(self, document_id: str) -> bool:
cursor.execute("DELETE FROM chunks_fts")
cursor.execute("DELETE FROM chunk_embeddings")
cursor.execute("DELETE FROM chunks")
deleted = cursor.rowcount > 0
if commit:
self.store._connection.commit()
return deleted
async def delete_by_document_id(
self, document_id: int, commit: bool = True
) -> bool:
"""Delete all chunks for a document.""" """Delete all chunks for a document."""
chunks = await self.get_by_document_id(document_id) chunks = await self.get_by_document_id(document_id)
deleted_any = False if not chunks:
for chunk in chunks: return False
if chunk.id is not None:
deleted = await self.delete(chunk.id, commit=False)
deleted_any = deleted_any or deleted
if commit and deleted_any and self.store._connection: self.store.chunks_table.delete(f"document_id = '{document_id}'")
self.store._connection.commit() return True
return deleted_any
async def search_chunks( async def search(
self, query: str, limit: int = 5 self, query: str, limit: int = 5, search_type: str = "hybrid"
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using vector similarity.""" """Search for relevant chunks using the specified search method.
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor() Args:
query: The search query string.
limit: Maximum number of results to return.
search_type: Type of search - "vector", "fts", or "hybrid" (default).
# Generate embedding for the query Returns:
List of (chunk, score) tuples ordered by relevance.
"""
if not query.strip():
return []
if search_type == "vector":
query_embedding = await self.embedder.embed(query) query_embedding = await self.embedder.embed(query)
serialized_query_embedding = self.store.serialize_embedding(query_embedding)
# Search for similar chunks using sqlite-vec results = self.store.chunks_table.search(
cursor.execute( query_embedding, query_type="vector", vector_column_name="vector"
""" ).limit(limit)
SELECT c.id, c.document_id, c.content, c.metadata, distance, d.uri, d.metadata as document_metadata
FROM chunk_embeddings return await self._process_search_results(results)
JOIN chunks c ON c.id = chunk_embeddings.chunk_id
JOIN documents d ON c.document_id = d.id elif search_type == "fts":
WHERE embedding MATCH :embedding AND k = :k results = self.store.chunks_table.search(query, query_type="fts").limit(
ORDER BY distance limit
""",
{"embedding": serialized_query_embedding, "k": limit},
) )
return await self._process_search_results(results)
results = cursor.fetchall() else: # hybrid (default)
return [
(
Chunk(
id=chunk_id,
document_id=document_id,
content=content,
metadata=json.loads(metadata_json) if metadata_json else {},
document_uri=document_uri,
document_meta=json.loads(document_metadata_json)
if document_metadata_json
else {},
),
1.0 / (1.0 + distance),
)
for chunk_id, document_id, content, metadata_json, distance, document_uri, document_metadata_json in results
]
async def search_chunks_fts(
self, query: str, limit: int = 5
) -> list[tuple[Chunk, float]]:
"""Search for chunks using FTS5 full-text search."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
# Clean the query for FTS5 - extract keywords for better matching
# Remove special characters and split into words
words = re.findall(r"\b\w+\b", query.lower())
# Join with OR to find chunks containing any of the keywords
fts_query = " OR ".join(words) if words else query
# Search using FTS5
cursor.execute(
"""
SELECT c.id, c.document_id, c.content, c.metadata, rank, d.uri, d.metadata as document_metadata
FROM chunks_fts
JOIN chunks c ON c.id = chunks_fts.rowid
JOIN documents d ON c.document_id = d.id
WHERE chunks_fts MATCH :query
ORDER BY rank
LIMIT :limit
""",
{"query": fts_query, "limit": limit},
)
results = cursor.fetchall()
return [
(
Chunk(
id=chunk_id,
document_id=document_id,
content=content,
metadata=json.loads(metadata_json) if metadata_json else {},
document_uri=document_uri,
document_meta=json.loads(document_metadata_json)
if document_metadata_json
else {},
),
-rank,
)
for chunk_id, document_id, content, metadata_json, rank, document_uri, document_metadata_json in results
# FTS5 rank is negative BM25 score
]
async def search_chunks_hybrid(
self, query: str, limit: int = 5, k: int = 60
) -> list[tuple[Chunk, float]]:
"""Hybrid search using Reciprocal Rank Fusion (RRF) combining vector similarity and FTS5 full-text search."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
# Generate embedding for the query
query_embedding = await self.embedder.embed(query) query_embedding = await self.embedder.embed(query)
serialized_query_embedding = self.store.serialize_embedding(query_embedding)
# Clean the query for FTS5 - extract keywords for better matching # Create RRF reranker
# Remove special characters and split into words reranker = RRFReranker()
words = re.findall(r"\b\w+\b", query.lower())
# Join with OR to find chunks containing any of the keywords
fts_query = " OR ".join(words) if words else query
# Perform hybrid search using RRF (Reciprocal Rank Fusion)
cursor.execute(
"""
WITH vector_search AS (
SELECT
c.id,
c.document_id,
c.content,
c.metadata,
ROW_NUMBER() OVER (ORDER BY ce.distance) as vector_rank
FROM chunk_embeddings ce
JOIN chunks c ON c.id = ce.chunk_id
WHERE ce.embedding MATCH :embedding AND k = :k_vector
ORDER BY ce.distance
),
fts_search AS (
SELECT
c.id,
c.document_id,
c.content,
c.metadata,
ROW_NUMBER() OVER (ORDER BY chunks_fts.rank) as fts_rank
FROM chunks_fts
JOIN chunks c ON c.id = chunks_fts.rowid
WHERE chunks_fts MATCH :fts_query
ORDER BY chunks_fts.rank
),
all_chunks AS (
SELECT id, document_id, content, metadata FROM vector_search
UNION
SELECT id, document_id, content, metadata FROM fts_search
),
rrf_scores AS (
SELECT
a.id,
a.document_id,
a.content,
a.metadata,
COALESCE(1.0 / (:k + v.vector_rank), 0) + COALESCE(1.0 / (:k + f.fts_rank), 0) as rrf_score
FROM all_chunks a
LEFT JOIN vector_search v ON a.id = v.id
LEFT JOIN fts_search f ON a.id = f.id
)
SELECT r.id, r.document_id, r.content, r.metadata, r.rrf_score, d.uri, d.metadata as document_metadata
FROM rrf_scores r
JOIN documents d ON r.document_id = d.id
ORDER BY r.rrf_score DESC
LIMIT :limit
""",
{
"embedding": serialized_query_embedding,
"k_vector": limit * 3,
"fts_query": fts_query,
"k": k,
"limit": limit,
},
)
results = cursor.fetchall() # Perform native hybrid search with RRF reranking
return [ results = (
( self.store.chunks_table.search(query_type="hybrid")
Chunk( .vector(query_embedding)
id=chunk_id, .text(query)
document_id=document_id, .rerank(reranker)
content=content, .limit(limit)
metadata=json.loads(metadata_json) if metadata_json else {},
document_uri=document_uri,
document_meta=json.loads(document_metadata_json)
if document_metadata_json
else {},
),
rrf_score,
) )
for chunk_id, document_id, content, metadata_json, rrf_score, document_uri, document_metadata_json in results return await self._process_search_results(results)
]
async def get_by_document_id(self, document_id: int) -> list[Chunk]: async def get_by_document_id(self, document_id: str) -> list[Chunk]:
"""Get all chunks for a specific document.""" """Get all chunks for a specific document."""
if self.store._connection is None: results = list(
raise ValueError("Store connection is not available") self.store.chunks_table.search()
.where(f"document_id = '{document_id}'")
cursor = self.store._connection.cursor() .to_pydantic(self.store.ChunkRecord)
cursor.execute(
"""
SELECT c.id, c.document_id, c.content, c.metadata, d.uri, d.metadata as document_metadata
FROM chunks c
JOIN documents d ON c.document_id = d.id
WHERE c.document_id = :document_id
ORDER BY JSON_EXTRACT(c.metadata, '$.order')
""",
{"document_id": document_id},
) )
rows = cursor.fetchall() # Get document info
return [ doc_results = list(
self.store.documents_table.search()
.where(f"id = '{document_id}'")
.limit(1)
.to_pydantic(DocumentRecord)
)
doc_uri = doc_results[0].uri if doc_results else None
doc_meta = doc_results[0].metadata if doc_results else "{}"
# Sort by order in metadata
chunks = [
Chunk( Chunk(
id=chunk_id, id=chunk.id,
document_id=document_id, document_id=chunk.document_id,
content=content, content=chunk.content,
metadata=json.loads(metadata_json) if metadata_json else {}, metadata=json.loads(chunk.metadata) if chunk.metadata else {},
document_uri=document_uri, document_uri=doc_uri,
document_meta=json.loads(document_metadata_json) document_meta=json.loads(doc_meta) if doc_meta else {},
if document_metadata_json
else {},
) )
for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows for chunk in results
] ]
chunks.sort(key=lambda c: c.metadata.get("order", 0))
return chunks
async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]: async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
"""Get adjacent chunks before and after the given chunk within the same document.""" """Get adjacent chunks before and after the given chunk within the same document."""
if self.store._connection is None: assert chunk.document_id, "Document id is required for adjacent chunk finding"
raise ValueError("Store connection is not available")
if chunk.document_id is None:
return []
cursor = self.store._connection.cursor()
chunk_order = chunk.metadata.get("order") chunk_order = chunk.metadata.get("order")
if chunk_order is None: if chunk_order is None:
return [] return []
# Get adjacent chunks within the same document # Get all chunks for the document
cursor.execute( all_chunks = await self.get_by_document_id(chunk.document_id)
"""
SELECT c.id, c.document_id, c.content, c.metadata, d.uri, d.metadata as document_metadata # Filter to adjacent chunks
FROM chunks c adjacent_chunks = []
JOIN documents d ON c.document_id = d.id for c in all_chunks:
WHERE c.document_id = :document_id c_order = c.metadata.get("order", 0)
AND JSON_EXTRACT(c.metadata, '$.order') BETWEEN :start_order AND :end_order if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
AND c.id != :chunk_id adjacent_chunks.append(c)
ORDER BY JSON_EXTRACT(c.metadata, '$.order')
""", return adjacent_chunks
{
"document_id": chunk.document_id, async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
"start_order": max(0, chunk_order - num_adjacent), """Process search results into chunks with document info and scores."""
"end_order": chunk_order + num_adjacent, chunks_with_scores = []
"chunk_id": chunk.id,
}, # Get both arrow and pydantic results to access scores
arrow_result = query_result.to_arrow()
pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord))
# Extract scores from arrow result based on search type
scores = []
column_names = arrow_result.column_names
if "_distance" in column_names:
# Vector search - distance (lower is better, convert to similarity)
distances = arrow_result.column("_distance").to_pylist()
scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
elif "_relevance_score" in column_names:
# Hybrid search - relevance score (higher is better)
scores = arrow_result.column("_relevance_score").to_pylist()
elif "_score" in column_names:
# FTS search - score (higher is better)
scores = arrow_result.column("_score").to_pylist()
else:
raise ValueError("Unknown search result format, cannot extract scores")
# Collect all unique document IDs for batch lookup
document_ids = list(set(chunk.document_id for chunk in pydantic_results))
# Batch fetch all documents at once
documents_map = {}
if document_ids:
# Create a WHERE clause for all document IDs
where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids)
doc_results = list(
self.store.documents_table.search()
.where(where_clause)
.to_pydantic(DocumentRecord)
)
documents_map = {doc.id: doc for doc in doc_results}
for i, chunk_record in enumerate(pydantic_results):
# Get document info from pre-fetched map
doc = documents_map.get(chunk_record.document_id)
doc_uri = doc.uri if doc else None
doc_meta = doc.metadata if doc else "{}"
chunk = Chunk(
id=chunk_record.id,
document_id=chunk_record.document_id,
content=chunk_record.content,
metadata=json.loads(chunk_record.metadata)
if chunk_record.metadata
else {},
document_uri=doc_uri,
document_meta=json.loads(doc_meta) if doc_meta else {},
) )
rows = cursor.fetchall() # Get score from arrow result
return [ score = scores[i] if i < len(scores) else 1.0
Chunk(
id=chunk_id, chunks_with_scores.append((chunk, score))
document_id=document_id,
content=content, return chunks_with_scores
metadata=json.loads(metadata_json) if metadata_json else {},
document_uri=document_uri,
document_meta=json.loads(document_metadata_json)
if document_metadata_json
else {},
)
for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows
]

View file

@ -1,27 +1,168 @@
import json import json
from datetime import datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from uuid import uuid4
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
from haiku.rag.store.engine import DocumentRecord, Store
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.base import BaseRepository
from haiku.rag.utils import text_to_docling_document
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
class DocumentRepository(BaseRepository[Document]): class DocumentRepository:
"""Repository for Document database operations.""" """Repository for Document operations."""
def __init__(self, store, chunk_repository=None): def __init__(self, store: Store) -> None:
super().__init__(store) self.store = store
# Avoid circular import by using late import if not provided self._chunk_repository = None
if chunk_repository is None:
@property
def chunk_repository(self):
"""Lazy-load ChunkRepository when needed."""
if self._chunk_repository is None:
from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.chunk import ChunkRepository
chunk_repository = ChunkRepository(store) self._chunk_repository = ChunkRepository(self.store)
self.chunk_repository = chunk_repository return self._chunk_repository
def _record_to_document(self, record: DocumentRecord) -> Document:
"""Convert a DocumentRecord to a Document model."""
return Document(
id=record.id,
content=record.content,
uri=record.uri,
metadata=json.loads(record.metadata) if record.metadata else {},
created_at=datetime.fromisoformat(record.created_at)
if record.created_at
else datetime.now(),
updated_at=datetime.fromisoformat(record.updated_at)
if record.updated_at
else datetime.now(),
)
async def create(self, entity: Document) -> Document:
"""Create a document in the database."""
# Generate new UUID
doc_id = str(uuid4())
# Create timestamp
now = datetime.now().isoformat()
# Create document record
doc_record = DocumentRecord(
id=doc_id,
content=entity.content,
uri=entity.uri,
metadata=json.dumps(entity.metadata),
created_at=now,
updated_at=now,
)
# Add to table
self.store.documents_table.add([doc_record])
entity.id = doc_id
entity.created_at = datetime.fromisoformat(now)
entity.updated_at = datetime.fromisoformat(now)
return entity
async def get_by_id(self, entity_id: str) -> Document | None:
"""Get a document by its ID."""
results = list(
self.store.documents_table.search()
.where(f"id = '{entity_id}'")
.limit(1)
.to_pydantic(DocumentRecord)
)
if not results:
return None
return self._record_to_document(results[0])
async def update(self, entity: Document) -> Document:
"""Update an existing document."""
assert entity.id, "Document ID is required for update"
# Update timestamp
now = datetime.now().isoformat()
entity.updated_at = datetime.fromisoformat(now)
# Update the record
self.store.documents_table.update(
where=f"id = '{entity.id}'",
values={
"content": entity.content,
"uri": entity.uri,
"metadata": json.dumps(entity.metadata),
"updated_at": now,
},
)
return entity
async def delete(self, entity_id: str) -> bool:
"""Delete a document by its ID."""
# Check if document exists
doc = await self.get_by_id(entity_id)
if doc is None:
return False
# Delete associated chunks first
await self.chunk_repository.delete_by_document_id(entity_id)
# Delete the document
self.store.documents_table.delete(f"id = '{entity_id}'")
return True
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[Document]:
"""List all documents with optional pagination."""
query = self.store.documents_table.search()
if offset is not None:
query = query.offset(offset)
if limit is not None:
query = query.limit(limit)
results = list(query.to_pydantic(DocumentRecord))
return [self._record_to_document(doc) for doc in results]
async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI."""
results = list(
self.store.documents_table.search()
.where(f"uri = '{uri}'")
.limit(1)
.to_pydantic(DocumentRecord)
)
if not results:
return None
return self._record_to_document(results[0])
async def delete_all(self) -> None:
"""Delete all documents from the database."""
# Delete all chunks first
await self.chunk_repository.delete_all()
# Get count before deletion
count = len(
list(
self.store.documents_table.search().limit(1).to_pydantic(DocumentRecord)
)
)
if count > 0:
# Drop and recreate table to clear all data
self.store.db.drop_table("documents")
self.store.documents_table = self.store.db.create_table(
"documents", schema=DocumentRecord
)
async def _create_with_docling( async def _create_with_docling(
self, self,
@ -30,219 +171,44 @@ class DocumentRepository(BaseRepository[Document]):
chunks: list["Chunk"] | None = None, chunks: list["Chunk"] | None = None,
) -> Document: ) -> Document:
"""Create a document with its chunks and embeddings.""" """Create a document with its chunks and embeddings."""
if self.store._connection is None: # Create the document
raise ValueError("Store connection is not available") created_doc = await self.create(entity)
cursor = self.store._connection.cursor() # Create chunks if not provided
if chunks is None:
# Start transaction assert created_doc.id is not None, (
cursor.execute("BEGIN TRANSACTION") "Document ID should not be None after creation"
try:
# Insert the document
cursor.execute(
"""
INSERT INTO documents (content, uri, metadata, created_at, updated_at)
VALUES (:content, :uri, :metadata, :created_at, :updated_at)
""",
{
"content": entity.content,
"uri": entity.uri,
"metadata": json.dumps(entity.metadata),
"created_at": entity.created_at,
"updated_at": entity.updated_at,
},
) )
document_id = cursor.lastrowid
assert document_id is not None, "Failed to create document in database"
entity.id = document_id
# Create chunks - either use provided chunks or generate from content
if chunks is not None:
# Use provided chunks, but update their document_id and set order from list position
for order, chunk in enumerate(chunks):
chunk.document_id = document_id
# Ensure order is set from list position
chunk.metadata = chunk.metadata.copy() if chunk.metadata else {}
chunk.metadata["order"] = order
await self.chunk_repository.create(chunk, commit=False)
else:
# Create chunks and embeddings using DoclingDocument
await self.chunk_repository.create_chunks_for_document( await self.chunk_repository.create_chunks_for_document(
document_id, docling_document, commit=False created_doc.id, docling_document
) )
else:
cursor.execute("COMMIT") # Use provided chunks, set order from list position
return entity assert created_doc.id is not None, (
"Document ID should not be None after creation"
except Exception:
cursor.execute("ROLLBACK")
raise
async def create(self, entity: Document) -> Document:
"""Create a document with its chunks and embeddings."""
# Convert content to DoclingDocument
docling_document = text_to_docling_document(entity.content)
return await self._create_with_docling(entity, docling_document)
async def get_by_id(self, entity_id: int) -> Document | None:
"""Get a document by its ID."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
cursor.execute(
"""
SELECT id, content, uri, metadata, created_at, updated_at
FROM documents WHERE id = :id
""",
{"id": entity_id},
) )
for order, chunk in enumerate(chunks):
chunk.document_id = created_doc.id
chunk.metadata["order"] = order
await self.chunk_repository.create(chunk)
row = cursor.fetchone() return created_doc
if row is None:
return None
document_id, content, uri, metadata_json, created_at, updated_at = row
metadata = json.loads(metadata_json) if metadata_json else {}
return Document(
id=document_id,
content=content,
uri=uri,
metadata=metadata,
created_at=created_at,
updated_at=updated_at,
)
async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
cursor.execute(
"""
SELECT id, content, uri, metadata, created_at, updated_at
FROM documents WHERE uri = :uri
""",
{"uri": uri},
)
row = cursor.fetchone()
if row is None:
return None
document_id, content, uri, metadata_json, created_at, updated_at = row
metadata = json.loads(metadata_json) if metadata_json else {}
return Document(
id=document_id,
content=content,
uri=uri,
metadata=metadata,
created_at=created_at,
updated_at=updated_at,
)
async def _update_with_docling( async def _update_with_docling(
self, entity: Document, docling_document: DoclingDocument self, entity: Document, docling_document: DoclingDocument
) -> Document: ) -> Document:
"""Update an existing document and regenerate its chunks and embeddings.""" """Update a document and regenerate its chunks."""
if self.store._connection is None: # Delete existing chunks
raise ValueError("Store connection is not available") assert entity.id is not None, "Document ID is required for update"
if entity.id is None: await self.chunk_repository.delete_by_document_id(entity.id)
raise ValueError("Document ID is required for update")
cursor = self.store._connection.cursor()
# Start transaction
cursor.execute("BEGIN TRANSACTION")
try:
# Update the document # Update the document
cursor.execute( updated_doc = await self.update(entity)
"""
UPDATE documents
SET content = :content, uri = :uri, metadata = :metadata, updated_at = :updated_at
WHERE id = :id
""",
{
"content": entity.content,
"uri": entity.uri,
"metadata": json.dumps(entity.metadata),
"updated_at": entity.updated_at,
"id": entity.id,
},
)
# Delete existing chunks and regenerate using DoclingDocument # Create new chunks
await self.chunk_repository.delete_by_document_id(entity.id, commit=False) assert updated_doc.id is not None, "Document ID should not be None after update"
await self.chunk_repository.create_chunks_for_document( await self.chunk_repository.create_chunks_for_document(
entity.id, docling_document, commit=False updated_doc.id, docling_document
) )
cursor.execute("COMMIT") return updated_doc
return entity
except Exception:
cursor.execute("ROLLBACK")
raise
async def update(self, entity: Document) -> Document:
"""Update an existing document and regenerate its chunks and embeddings."""
# Convert content to DoclingDocument
docling_document = text_to_docling_document(entity.content)
return await self._update_with_docling(entity, docling_document)
async def delete(self, entity_id: int) -> bool:
"""Delete a document and all its associated chunks and embeddings."""
# Delete chunks and embeddings first
await self.chunk_repository.delete_by_document_id(entity_id)
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
cursor.execute("DELETE FROM documents WHERE id = :id", {"id": entity_id})
deleted = cursor.rowcount > 0
self.store._connection.commit()
return deleted
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[Document]:
"""List all documents with optional pagination."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
query = "SELECT id, content, uri, metadata, created_at, updated_at FROM documents ORDER BY created_at DESC"
params = {}
if limit is not None:
query += " LIMIT :limit"
params["limit"] = limit
if offset is not None:
query += " OFFSET :offset"
params["offset"] = offset
cursor.execute(query, params)
rows = cursor.fetchall()
return [
Document(
id=document_id,
content=content,
uri=uri,
metadata=json.loads(metadata_json) if metadata_json else {},
created_at=created_at,
updated_at=updated_at,
)
for document_id, content, uri, metadata_json, created_at, updated_at in rows
]

View file

@ -1,77 +1,143 @@
import json import json
from typing import Any
from haiku.rag.store.engine import Store from haiku.rag.config import Config
from haiku.rag.store.engine import SettingsRecord, Store
class ConfigMismatchError(Exception): class ConfigMismatchError(Exception):
"""Raised when current config doesn't match stored settings.""" """Raised when stored config doesn't match current config."""
pass pass
class SettingsRepository: class SettingsRepository:
def __init__(self, store: Store): """Repository for Settings operations."""
def __init__(self, store: Store) -> None:
self.store = store self.store = store
def get(self) -> dict[str, Any]: async def create(self, entity: dict) -> dict:
"""Get all settings from the database.""" """Create settings in the database."""
if self.store._connection is None: settings_record = SettingsRecord(id="settings", settings=json.dumps(entity))
raise ValueError("Store connection is not available") self.store.settings_table.add([settings_record])
return entity
cursor = self.store._connection.execute("SELECT settings FROM settings LIMIT 1") async def get_by_id(self, entity_id: str) -> dict | None:
row = cursor.fetchone() """Get settings by ID."""
if row: results = list(
return json.loads(row[0]) self.store.settings_table.search()
return {} .where(f"id = '{entity_id}'")
.limit(1)
def save(self) -> None: .to_pydantic(SettingsRecord)
"""Sync settings from the current AppConfig to database."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
from haiku.rag.config import Config
settings_json = Config.model_dump_json()
self.store._connection.execute(
"INSERT INTO settings (id, settings) VALUES (1, ?) ON CONFLICT(id) DO UPDATE SET settings = excluded.settings",
(settings_json,),
) )
self.store._connection.commit() if not results:
return None
return json.loads(results[0].settings) if results[0].settings else {}
async def update(self, entity: dict) -> dict:
"""Update existing settings."""
self.store.settings_table.update(
where="id = 'settings'", values={"settings": json.dumps(entity)}
)
return entity
async def delete(self, entity_id: str) -> bool:
"""Delete settings by ID."""
self.store.settings_table.delete(f"id = '{entity_id}'")
return True
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[dict]:
"""List all settings."""
results = list(self.store.settings_table.search().to_pydantic(SettingsRecord))
return [
json.loads(record.settings) if record.settings else {} for record in results
]
def get_current_settings(self) -> dict:
"""Get the current settings."""
results = list(
self.store.settings_table.search()
.where("id = 'settings'")
.limit(1)
.to_pydantic(SettingsRecord)
)
if not results:
return {}
return json.loads(results[0].settings) if results[0].settings else {}
def save_current_settings(self) -> None:
"""Save the current configuration to the database."""
current_config = Config.model_dump(mode="json")
# Check if settings exist
existing = list(
self.store.settings_table.search()
.where("id = 'settings'")
.limit(1)
.to_pydantic(SettingsRecord)
)
if existing:
# Update existing settings
self.store.settings_table.update(
where="id = 'settings'", values={"settings": json.dumps(current_config)}
)
else:
# Create new settings
settings_record = SettingsRecord(
id="settings", settings=json.dumps(current_config)
)
self.store.settings_table.add([settings_record])
def validate_config_compatibility(self) -> None: def validate_config_compatibility(self) -> None:
"""Check if current config is compatible with stored settings. """Validate that the current configuration is compatible with stored settings."""
stored_settings = self.get_current_settings()
Raises ConfigMismatchError if there are incompatible differences. # If no stored settings, this is a new database - save current config and return
If no settings exist, saves current config. if not stored_settings:
""" self.save_current_settings()
db_settings = self.get()
if not db_settings:
# No settings in DB, save current config
self.save()
return return
from haiku.rag.config import Config
current_config = Config.model_dump(mode="json") current_config = Config.model_dump(mode="json")
# Critical settings that must match # Check if embedding provider or model has changed
critical_settings = [ stored_provider = stored_settings.get("EMBEDDINGS_PROVIDER")
"EMBEDDINGS_PROVIDER", current_provider = current_config.get("EMBEDDINGS_PROVIDER")
"EMBEDDINGS_MODEL",
"EMBEDDINGS_VECTOR_DIM",
"CHUNK_SIZE",
]
errors = [] stored_model = stored_settings.get("EMBEDDINGS_MODEL")
for setting in critical_settings: current_model = current_config.get("EMBEDDINGS_MODEL")
if db_settings.get(setting) != current_config.get(setting):
errors.append( stored_vector_dim = stored_settings.get("EMBEDDINGS_VECTOR_DIM")
f"{setting}: current={current_config.get(setting)}, stored={db_settings.get(setting)}" current_vector_dim = current_config.get("EMBEDDINGS_VECTOR_DIM")
# Check for incompatible changes
incompatible_changes = []
if stored_provider and stored_provider != current_provider:
incompatible_changes.append(
f"Embedding provider changed from '{stored_provider}' to '{current_provider}'"
) )
if errors: if stored_model and stored_model != current_model:
error_msg = f"Config mismatch detected: {'; '.join(errors)}. Consider rebuilding the database with the current configuration." incompatible_changes.append(
f"Embedding model changed from '{stored_model}' to '{current_model}'"
)
if stored_vector_dim and stored_vector_dim != current_vector_dim:
incompatible_changes.append(
f"Vector dimension changed from {stored_vector_dim} to {current_vector_dim}"
)
if incompatible_changes:
error_msg = (
"Database configuration is incompatible with current settings:\n"
+ "\n".join(f" - {change}" for change in incompatible_changes)
)
error_msg += "\n\nPlease rebuild the database using: haiku-rag rebuild"
raise ConfigMismatchError(error_msg) raise ConfigMismatchError(error_msg)

View file

@ -1,3 +1 @@
from haiku.rag.store.upgrades.v0_3_4 import upgrades as v0_3_4_upgrades upgrades = []
upgrades = v0_3_4_upgrades

View file

@ -1,26 +0,0 @@
from collections.abc import Callable
from sqlite3 import Connection
from haiku.rag.config import Config
def add_settings_table(db: Connection) -> None:
"""Create settings table for storing current configuration"""
db.execute("""
CREATE TABLE settings (
id INTEGER PRIMARY KEY DEFAULT 1,
settings TEXT NOT NULL DEFAULT '{}'
)
""")
settings_json = Config.model_dump_json()
db.execute(
"INSERT INTO settings (id, settings) VALUES (1, ?)",
(settings_json,),
)
db.commit()
upgrades: list[tuple[str, list[Callable[[Connection], None]]]] = [
("0.3.4", [add_settings_table])
]

View file

@ -1,4 +1,7 @@
import asyncio
import sys import sys
from collections.abc import Callable
from functools import wraps
from importlib import metadata from importlib import metadata
from io import BytesIO from io import BytesIO
from pathlib import Path from pathlib import Path
@ -10,6 +13,42 @@ from docling_core.types.io import DocumentStream
from packaging.version import Version, parse from packaging.version import Version, parse
def debounce(wait: float) -> Callable:
"""
A decorator to debounce a function, ensuring it is called only after a specified delay
and always executes after the last call.
Args:
wait (float): The debounce delay in seconds.
Returns:
Callable: The decorated function.
"""
def decorator(func: Callable) -> Callable:
last_call = None
task = None
@wraps(func)
async def debounced(*args, **kwargs):
nonlocal last_call, task
last_call = asyncio.get_event_loop().time()
if task:
task.cancel()
async def call_func():
await asyncio.sleep(wait)
if asyncio.get_event_loop().time() - last_call >= wait: # type: ignore
await func(*args, **kwargs)
task = asyncio.create_task(call_func())
return debounced
return decorator
def get_default_data_dir() -> Path: def get_default_data_dir() -> Path:
"""Get the user data directory for the current system platform. """Get the user data directory for the current system platform.
@ -32,37 +71,6 @@ def get_default_data_dir() -> Path:
return data_path return data_path
def semantic_version_to_int(version: str) -> int:
"""Convert a semantic version string to an integer.
Args:
version: Semantic version string.
Returns:
Integer representation of semantic version.
"""
major, minor, patch = version.split(".")
major = int(major) << 16
minor = int(minor) << 8
patch = int(patch)
return major + minor + patch
def int_to_semantic_version(version: int) -> str:
"""Convert an integer to a semantic version string.
Args:
version: Integer representation of semantic version.
Returns:
Semantic version string.
"""
major = version >> 16
minor = (version >> 8) & 255
patch = version & 255
return f"{major}.{minor}.{patch}"
async def is_up_to_date() -> tuple[bool, Version, Version]: async def is_up_to_date() -> tuple[bool, Version, Version]:
"""Check whether haiku.rag is current. """Check whether haiku.rag is current.

View file

@ -1,3 +1,4 @@
import tempfile
from pathlib import Path from pathlib import Path
import pytest import pytest
@ -16,3 +17,10 @@ def qa_corpus() -> Dataset:
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
corpus.save_to_disk(ds_path) corpus.save_to_disk(ds_path)
return corpus return corpus
@pytest.fixture
def temp_db_path():
"""Create a temporary database path for testing."""
with tempfile.TemporaryDirectory() as temp_dir:
yield Path(temp_dir) / "test.lancedb"

View file

@ -11,7 +11,7 @@ from haiku.rag.qa import get_qa_agent
console = Console() console = Console()
db_path = Path(__file__).parent / "data" / "benchmark.sqlite" db_path = Path(__file__).parent / "data" / "benchmark.lancedb"
async def populate_db(): async def populate_db():
@ -53,6 +53,10 @@ async def run_match_benchmark():
async with HaikuRAG(db_path) as rag: async with HaikuRAG(db_path) as rag:
for doc in corpus: for doc in corpus:
doc_id = doc["document_id"] # type: ignore doc_id = doc["document_id"] # type: ignore
expected_answer = doc["answer"] # type: ignore
if expected_answer == "The answer is not found in the document.":
progress.advance(task)
continue
matches = await rag.search( matches = await rag.search(
query=doc["question"], # type: ignore query=doc["question"], # type: ignore
limit=3, limit=3,
@ -62,6 +66,9 @@ async def run_match_benchmark():
# Check position of correct document in results # Check position of correct document in results
for position, (chunk, _) in enumerate(matches): for position, (chunk, _) in enumerate(matches):
assert chunk.document_id is not None, (
"Chunk document_id should not be None"
)
retrieved = await rag.get_document_by_id(chunk.document_id) retrieved = await rag.get_document_by_id(chunk.document_id)
if retrieved and retrieved.uri == doc_id: if retrieved and retrieved.uri == doc_id:
if position == 0: # First position if position == 0: # First position
@ -113,6 +120,8 @@ async def run_qa_benchmark(k: int | None = None):
question = doc["question"] # type: ignore question = doc["question"] # type: ignore
expected_answer = doc["answer"] # type: ignore expected_answer = doc["answer"] # type: ignore
# Really small models might fail, let's account for that in try/except
try:
generated_answer = await qa.answer(question) generated_answer = await qa.answer(question)
is_equivalent = await judge.judge_answers( is_equivalent = await judge.judge_answers(
question, generated_answer, expected_answer question, generated_answer, expected_answer
@ -124,9 +133,14 @@ async def run_qa_benchmark(k: int | None = None):
if is_equivalent: if is_equivalent:
correct_answers += 1 correct_answers += 1
except Exception as e:
console.print(f"[red]Error processing question: {question}[/red]")
console.print(f"[red]{e}[/red]")
finally:
total_questions += 1 total_questions += 1
console.print("Current score:", correct_answers, "/", total_questions) console.print(
"Current score:", correct_answers, "/", total_questions
)
progress.advance(task) progress.advance(task)
accuracy = correct_answers / total_questions if total_questions > 0 else 0 accuracy = correct_answers / total_questions if total_questions > 0 else 0

View file

@ -1,5 +1,4 @@
import asyncio import asyncio
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
import pytest import pytest
@ -9,16 +8,16 @@ from haiku.rag.store.models.document import Document
@pytest.fixture @pytest.fixture
def app(): def app(tmp_path):
return HaikuRAGApp(db_path=Path(":memory:")) return HaikuRAGApp(db_path=tmp_path / "test.lancedb")
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_list_documents(app: HaikuRAGApp, monkeypatch): async def test_list_documents(app: HaikuRAGApp, monkeypatch):
"""Test listing documents.""" """Test listing documents."""
mock_docs = [ mock_docs = [
Document(id=1, content="doc 1"), Document(id="1", content="doc 1"),
Document(id=2, content="doc 2"), Document(id="2", content="doc 2"),
] ]
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.list_documents.return_value = mock_docs mock_client.list_documents.return_value = mock_docs
@ -42,7 +41,7 @@ async def test_list_documents(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_add_document_from_text(app: HaikuRAGApp, monkeypatch): async def test_add_document_from_text(app: HaikuRAGApp, monkeypatch):
"""Test adding a document from text.""" """Test adding a document from text."""
mock_doc = Document(id=1, content="test document") mock_doc = Document(id="1", content="test document")
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.create_document.return_value = mock_doc mock_client.create_document.return_value = mock_doc
mock_client.__aenter__.return_value = mock_client mock_client.__aenter__.return_value = mock_client
@ -65,7 +64,7 @@ async def test_add_document_from_text(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch): async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch):
"""Test adding a document from a source path.""" """Test adding a document from a source path."""
mock_doc = Document(id=1, content="test document") mock_doc = Document(id="1", content="test document")
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.create_document_from_source.return_value = mock_doc mock_client.create_document_from_source.return_value = mock_doc
mock_client.__aenter__.return_value = mock_client mock_client.__aenter__.return_value = mock_client
@ -89,7 +88,7 @@ async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_get_document(app: HaikuRAGApp, monkeypatch): async def test_get_document(app: HaikuRAGApp, monkeypatch):
"""Test getting a document.""" """Test getting a document."""
mock_doc = Document(id=1, content="test document") mock_doc = Document(id="1", content="test document")
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.get_document_by_id.return_value = mock_doc mock_client.get_document_by_id.return_value = mock_doc
mock_client.__aenter__.return_value = mock_client mock_client.__aenter__.return_value = mock_client
@ -98,9 +97,9 @@ async def test_get_document(app: HaikuRAGApp, monkeypatch):
monkeypatch.setattr(app, "_rich_print_document", mock_rich_print) monkeypatch.setattr(app, "_rich_print_document", mock_rich_print)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.get_document(1) await app.get_document("1")
mock_client.get_document_by_id.assert_called_once_with(1) mock_client.get_document_by_id.assert_called_once_with("1")
mock_rich_print.assert_called_once_with(mock_doc, truncate=False) mock_rich_print.assert_called_once_with(mock_doc, truncate=False)
@ -115,9 +114,9 @@ async def test_get_document_not_found(app: HaikuRAGApp, monkeypatch):
monkeypatch.setattr(app.console, "print", mock_print) monkeypatch.setattr(app.console, "print", mock_print)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.get_document(1) await app.get_document("1")
mock_client.get_document_by_id.assert_called_once_with(1) mock_client.get_document_by_id.assert_called_once_with("1")
mock_print.assert_called_once_with("[red]Document with id 1 not found.[/red]") mock_print.assert_called_once_with("[red]Document with id 1 not found.[/red]")
@ -131,9 +130,9 @@ async def test_delete_document(app: HaikuRAGApp, monkeypatch):
monkeypatch.setattr(app.console, "print", mock_print) monkeypatch.setattr(app.console, "print", mock_print)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.delete_document(1) await app.delete_document("1")
mock_client.delete_document.assert_called_once_with(1) mock_client.delete_document.assert_called_once_with("1")
mock_print.assert_called_once_with("[b]Document 1 deleted successfully.[/b]") mock_print.assert_called_once_with("[b]Document 1 deleted successfully.[/b]")
@ -151,7 +150,7 @@ async def test_search(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.search("query") await app.search("query")
mock_client.search.assert_called_once_with("query", limit=5, k=60) mock_client.search.assert_called_once_with("query", limit=5)
assert mock_rich_print_search.call_count == len(mock_results) assert mock_rich_print_search.call_count == len(mock_results)
@ -168,7 +167,7 @@ async def test_search_no_results(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.search("query") await app.search("query")
mock_client.search.assert_called_once_with("query", limit=5, k=60) mock_client.search.assert_called_once_with("query", limit=5)
mock_print.assert_called_once_with("[red]No results found.[/red]") mock_print.assert_called_once_with("[red]No results found.[/red]")

View file

@ -10,10 +10,10 @@ from haiku.rag.utils import text_to_docling_document
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chunk_repository_operations(qa_corpus: Dataset): async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
"""Test ChunkRepository operations.""" """Test ChunkRepository operations."""
# Create an in-memory store and repositories # Create a store and repositories
store = Store(":memory:") store = Store(temp_db_path)
doc_repo = DocumentRepository(store) doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store) chunk_repo = ChunkRepository(store)
@ -21,9 +21,12 @@ async def test_chunk_repository_operations(qa_corpus: Dataset):
first_doc = qa_corpus[0] first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"] document_text = first_doc["document_extracted"]
# Create a document first # Create a document first with chunks
document = Document(content=document_text, metadata={"source": "test"}) document = Document(content=document_text, metadata={"source": "test"})
created_document = await doc_repo.create(document) from haiku.rag.utils import text_to_docling_document
docling_document = text_to_docling_document(document_text, name="test.md")
created_document = await doc_repo._create_with_docling(document, docling_document)
assert created_document.id is not None assert created_document.id is not None
# Test getting chunks by document ID # Test getting chunks by document ID
@ -32,7 +35,7 @@ async def test_chunk_repository_operations(qa_corpus: Dataset):
assert all(chunk.document_id == created_document.id for chunk in chunks) assert all(chunk.document_id == created_document.id for chunk in chunks)
# Test chunk search # Test chunk search
results = await chunk_repo.search_chunks("election", limit=2) results = await chunk_repo.search("election", limit=2, search_type="vector")
assert len(results) <= 2 assert len(results) <= 2
assert all(hasattr(chunk, "content") for chunk, _ in results) assert all(hasattr(chunk, "content") for chunk, _ in results)
@ -48,11 +51,12 @@ async def test_chunk_repository_operations(qa_corpus: Dataset):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_create_chunks_for_document(qa_corpus: Dataset): async def test_create_chunks_for_document(qa_corpus: Dataset, temp_db_path):
"""Test creating chunks for a document.""" """Test creating chunks for a document."""
# Create an in-memory store and repositories # Create a store and repositories
store = Store(":memory:") store = Store(temp_db_path)
chunk_repo = ChunkRepository(store) chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# Get the first document from the corpus # Get the first document from the corpus
first_doc = qa_corpus[0] first_doc = qa_corpus[0]
@ -60,21 +64,8 @@ async def test_create_chunks_for_document(qa_corpus: Dataset):
# Create a document first (without chunks) # Create a document first (without chunks)
document = Document(content=document_text, metadata={"source": "test"}) document = Document(content=document_text, metadata={"source": "test"})
created_document = await doc_repo.create(document)
# Insert document manually to test chunk creation independently document_id = created_document.id
document_id = None
if store._connection is not None:
cursor = store._connection.cursor()
cursor.execute(
"""
INSERT INTO documents (content, metadata, created_at, updated_at)
VALUES (?, ?, ?, ?)
""",
(document.content, "{}", document.created_at, document.updated_at),
)
document_id = cursor.lastrowid
document.id = document_id
store._connection.commit()
assert document_id is not None, "Document ID should not be None" assert document_id is not None, "Document ID should not be None"
@ -101,25 +92,17 @@ async def test_create_chunks_for_document(qa_corpus: Dataset):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chunk_repository_crud(): async def test_chunk_repository_crud(temp_db_path):
"""Test basic CRUD operations in ChunkRepository.""" """Test basic CRUD operations in ChunkRepository."""
# Create an in-memory store # Create a store
store = Store(":memory:") store = Store(temp_db_path)
chunk_repo = ChunkRepository(store) chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# First create a document to reference # First create a document to reference
document_id = None document = Document(content="Test document content", metadata={})
if store._connection is not None: created_document = await doc_repo.create(document)
cursor = store._connection.cursor() document_id = created_document.id
cursor.execute(
"""
INSERT INTO documents (content, metadata, created_at, updated_at)
VALUES (?, ?, datetime('now'), datetime('now'))
""",
("Test document content", "{}"),
)
document_id = cursor.lastrowid
store._connection.commit()
assert document_id is not None, "Document ID should not be None" assert document_id is not None, "Document ID should not be None"
@ -162,9 +145,9 @@ async def test_chunk_repository_crud():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_adjacent_chunks(): async def test_adjacent_chunks(temp_db_path):
"""Test the get_adjacent_chunks repository method.""" """Test the get_adjacent_chunks repository method."""
store = Store(":memory:") store = Store(temp_db_path)
doc_repo = DocumentRepository(store) doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store) chunk_repo = ChunkRepository(store)

View file

@ -54,7 +54,7 @@ def test_get_document():
result = runner.invoke(cli, ["get", "1"]) result = runner.invoke(cli, ["get", "1"])
assert result.exit_code == 0 assert result.exit_code == 0
mock_app_instance.get_document.assert_called_once_with(doc_id=1) mock_app_instance.get_document.assert_called_once_with(doc_id="1")
def test_delete_document(): def test_delete_document():
@ -66,7 +66,7 @@ def test_delete_document():
result = runner.invoke(cli, ["delete", "1"]) result = runner.invoke(cli, ["delete", "1"])
assert result.exit_code == 0 assert result.exit_code == 0
mock_app_instance.delete_document.assert_called_once_with(doc_id=1) mock_app_instance.delete_document.assert_called_once_with(doc_id="1")
def test_search(): def test_search():
@ -78,7 +78,7 @@ def test_search():
result = runner.invoke(cli, ["search", "query"]) result = runner.invoke(cli, ["search", "query"])
assert result.exit_code == 0 assert result.exit_code == 0
mock_app_instance.search.assert_called_once_with(query="query", limit=5, k=60) mock_app_instance.search.assert_called_once_with(query="query", limit=5)
def test_serve(): def test_serve():

View file

@ -11,9 +11,9 @@ from haiku.rag.store.models.chunk import Chunk
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_document_crud(qa_corpus: Dataset): async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
"""Test HaikuRAG CRUD operations for documents.""" """Test HaikuRAG CRUD operations for documents."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Get test data # Get test data
first_doc = qa_corpus[0] first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"] document_text = first_doc["document_extracted"]
@ -77,9 +77,9 @@ async def test_client_document_crud(qa_corpus: Dataset):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_source(): async def test_client_create_document_from_source(temp_db_path):
"""Test creating a document from a file source.""" """Test creating a document from a file source."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir: with tempfile.TemporaryDirectory() as temp_dir:
test_content = "This is test content from a file." test_content = "This is test content from a file."
temp_path = Path(temp_dir) / "test.txt" temp_path = Path(temp_dir) / "test.txt"
@ -106,9 +106,9 @@ async def test_client_create_document_from_source():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_source_unsupported(): async def test_client_create_document_from_source_unsupported(temp_db_path):
"""Test creating a document from an unsupported file type.""" """Test creating a document from an unsupported file type."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create a temporary file with unsupported extension # Create a temporary file with unsupported extension
with tempfile.NamedTemporaryFile( with tempfile.NamedTemporaryFile(
mode="w", suffix=".unsupported", delete=False mode="w", suffix=".unsupported", delete=False
@ -122,9 +122,9 @@ async def test_client_create_document_from_source_unsupported():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_source_nonexistent(): async def test_client_create_document_from_source_nonexistent(temp_db_path):
"""Test creating a document from a non-existent file.""" """Test creating a document from a non-existent file."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
non_existent_path = Path("/non/existent/file.txt") non_existent_path = Path("/non/existent/file.txt")
# Should raise ValueError when file doesn't exist # Should raise ValueError when file doesn't exist
@ -133,9 +133,9 @@ async def test_client_create_document_from_source_nonexistent():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url(): async def test_client_create_document_from_url(temp_db_path):
"""Test creating a document from a URL.""" """Test creating a document from a URL."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Mock the HTTP response # Mock the HTTP response
mock_response = AsyncMock() mock_response = AsyncMock()
mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>" mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
@ -158,9 +158,11 @@ async def test_client_create_document_from_url():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url_with_different_content_types(): async def test_client_create_document_from_url_with_different_content_types(
temp_db_path,
):
"""Test creating documents from URLs with different content types.""" """Test creating documents from URLs with different content types."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Test JSON content # Test JSON content
mock_json_response = AsyncMock() mock_json_response = AsyncMock()
mock_json_response.content = ( mock_json_response.content = (
@ -201,9 +203,9 @@ async def test_client_create_document_from_url_with_different_content_types():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url_unsupported_content(): async def test_client_create_document_from_url_unsupported_content(temp_db_path):
"""Test creating a document from URL with unsupported content type.""" """Test creating a document from URL with unsupported content type."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Mock response with unsupported content type # Mock response with unsupported content type
mock_response = AsyncMock() mock_response = AsyncMock()
mock_response.content = b"binary content" mock_response.content = b"binary content"
@ -218,9 +220,9 @@ async def test_client_create_document_from_url_unsupported_content():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url_http_error(): async def test_client_create_document_from_url_http_error(temp_db_path):
"""Test handling HTTP errors when creating document from URL.""" """Test handling HTTP errors when creating document from URL."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
with patch("httpx.AsyncClient.get") as mock_get: with patch("httpx.AsyncClient.get") as mock_get:
mock_get.side_effect = httpx.HTTPStatusError( mock_get.side_effect = httpx.HTTPStatusError(
"404 Not Found", "404 Not Found",
@ -235,9 +237,9 @@ async def test_client_create_document_from_url_http_error():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_get_extension_from_content_type_or_url(): async def test_get_extension_from_content_type_or_url(temp_db_path):
"""Test the helper method for determining file extensions.""" """Test the helper method for determining file extensions."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Test content type mappings # Test content type mappings
assert ( assert (
client._get_extension_from_content_type_or_url("", "text/html") == ".html" client._get_extension_from_content_type_or_url("", "text/html") == ".html"
@ -280,11 +282,11 @@ async def test_get_extension_from_content_type_or_url():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_metadata_content_type_and_md5(): async def test_client_metadata_content_type_and_md5(temp_db_path):
"""Test that contentType and md5 metadata are correctly set.""" """Test that contentType and md5 metadata are correctly set."""
import hashlib import hashlib
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create a temporary file with known content # Create a temporary file with known content
test_content = "Test content for MD5 calculation." test_content = "Test content for MD5 calculation."
expected_md5 = hashlib.md5(test_content.encode()).hexdigest() expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
@ -313,9 +315,9 @@ async def test_client_metadata_content_type_and_md5():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_update_no_op_behavior(): async def test_client_create_update_no_op_behavior(temp_db_path):
"""Test create/update/no-op behavior based on MD5 changes.""" """Test create/update/no-op behavior based on MD5 changes."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create a temporary file # Create a temporary file
test_content = "Original content for testing." test_content = "Original content for testing."
with tempfile.TemporaryDirectory() as temp_dir: with tempfile.TemporaryDirectory() as temp_dir:
@ -349,9 +351,9 @@ async def test_client_create_update_no_op_behavior():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_url_create_update_no_op_behavior(): async def test_client_url_create_update_no_op_behavior(temp_db_path):
"""Test create/update/no-op behavior for URLs based on MD5 changes.""" """Test create/update/no-op behavior for URLs based on MD5 changes."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
url = "https://example.com/test.txt" url = "https://example.com/test.txt"
original_content = b"Original URL content" original_content = b"Original URL content"
updated_content = b"Updated URL content" updated_content = b"Updated URL content"
@ -385,9 +387,9 @@ async def test_client_url_create_update_no_op_behavior():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_search(): async def test_client_search(temp_db_path):
"""Test HaikuRAG search functionality.""" """Test HaikuRAG search functionality."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Add multiple documents to search from # Add multiple documents to search from
doc1_text = "Python is a high-level programming language known for its simplicity and readability." doc1_text = "Python is a high-level programming language known for its simplicity and readability."
doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming." doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
@ -428,11 +430,11 @@ async def test_client_search():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_async_context_manager(): async def test_client_async_context_manager(temp_db_path):
"""Test HaikuRAG as async context manager.""" """Test HaikuRAG as async context manager."""
# Test that context manager works and auto-closes # Test that context manager works and auto-closes
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create a document to ensure the client works # Create a document to ensure the client works
doc = await client.create_document( doc = await client.create_document(
content="Test content for context manager", content="Test content for context manager",
@ -453,9 +455,9 @@ async def test_client_async_context_manager():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_with_custom_chunks(): async def test_client_create_document_with_custom_chunks(temp_db_path):
"""Test creating a document with pre-created chunks.""" """Test creating a document with pre-created chunks."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create some custom chunks with and without embeddings # Create some custom chunks with and without embeddings
chunks = [ chunks = [
Chunk(content="This is the first chunk", metadata={"custom": "metadata1"}), Chunk(content="This is the first chunk", metadata={"custom": "metadata1"}),
@ -492,9 +494,9 @@ async def test_client_create_document_with_custom_chunks():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_ask_without_cite(): async def test_client_ask_without_cite(temp_db_path):
"""Test asking questions without citations.""" """Test asking questions without citations."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Mock the QA agent # Mock the QA agent
mock_qa_agent = AsyncMock() mock_qa_agent = AsyncMock()
mock_qa_agent.answer.return_value = "Test answer" mock_qa_agent.answer.return_value = "Test answer"
@ -507,9 +509,9 @@ async def test_client_ask_without_cite():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_ask_with_cite(): async def test_client_ask_with_cite(temp_db_path):
"""Test asking questions with citations.""" """Test asking questions with citations."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Mock the QA agent # Mock the QA agent
mock_qa_agent = AsyncMock() mock_qa_agent = AsyncMock()
mock_qa_agent.answer.return_value = "Test answer with citations [1]" mock_qa_agent.answer.return_value = "Test answer with citations [1]"
@ -522,11 +524,11 @@ async def test_client_ask_with_cite():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context(): async def test_client_expand_context(temp_db_path):
"""Test expanding search results with adjacent chunks.""" """Test expanding search results with adjacent chunks."""
# Mock Config to have CONTEXT_CHUNK_RADIUS = 2 # Mock Config to have CONTEXT_CHUNK_RADIUS = 2
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2): with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2):
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create chunks manually # Create chunks manually
manual_chunks = [ manual_chunks = [
Chunk(content="Chunk 0 content", metadata={"order": 0}), Chunk(content="Chunk 0 content", metadata={"order": 0}),
@ -571,10 +573,10 @@ async def test_client_expand_context():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context_radius_zero(): async def test_client_expand_context_radius_zero(temp_db_path):
"""Test expand_context with radius 0 returns original results.""" """Test expand_context with radius 0 returns original results."""
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 0): with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 0):
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create a simple document # Create a simple document
doc = await client.create_document(content="Simple test content") doc = await client.create_document(content="Simple test content")
assert doc.id is not None assert doc.id is not None
@ -588,10 +590,10 @@ async def test_client_expand_context_radius_zero():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context_multiple_chunks(): async def test_client_expand_context_multiple_chunks(temp_db_path):
"""Test expand_context with multiple search results.""" """Test expand_context with multiple search results."""
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1): with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1):
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create first document with manual chunks # Create first document with manual chunks
doc1_chunks = [ doc1_chunks = [
Chunk(content="Doc1 Part A", metadata={"order": 0}), Chunk(content="Doc1 Part A", metadata={"order": 0}),
@ -642,9 +644,9 @@ async def test_client_expand_context_multiple_chunks():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context_merges_overlapping_chunks(): async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
"""Test that overlapping expanded chunks are merged into one.""" """Test that overlapping expanded chunks are merged into one."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create document with 5 chunks # Create document with 5 chunks
manual_chunks = [ manual_chunks = [
Chunk(content="Chunk 0", metadata={"order": 0}), Chunk(content="Chunk 0", metadata={"order": 0}),
@ -689,9 +691,9 @@ async def test_client_expand_context_merges_overlapping_chunks():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context_keeps_separate_non_overlapping(): async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path):
"""Test that non-overlapping expanded chunks remain separate.""" """Test that non-overlapping expanded chunks remain separate."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
# Create document with chunks far apart # Create document with chunks far apart
manual_chunks = [ manual_chunks = [
Chunk(content="Chunk 0", metadata={"order": 0}), Chunk(content="Chunk 0", metadata={"order": 0}),
@ -716,7 +718,7 @@ async def test_client_expand_context_keeps_separate_non_overlapping():
) # Content: "Chunk 0" ) # Content: "Chunk 0"
chunk5 = next( chunk5 = next(
c for c in chunks if c.metadata.get("order") == 5 c for c in chunks if c.metadata.get("order") == 5
) # Content: "Chunk 7" ) # Content: "Chunk 7" but now at order 5
# chunk0 expanded: [0,1] with radius=1 (orders 0,1) # chunk0 expanded: [0,1] with radius=1 (orders 0,1)
# chunk5 expanded: [4,5] with radius=1 (orders 4,5) # chunk5 expanded: [4,5] with radius=1 (orders 4,5)
@ -743,7 +745,7 @@ async def test_client_expand_context_keeps_separate_non_overlapping():
assert score1 == 0.8 assert score1 == 0.8
# Second chunk (order=5) expanded should contain orders [4,5] # Second chunk (order=5) expanded should contain orders [4,5]
# Content should be "Chunk 6" + "Chunk 7" (orders 4 and 5) # Content should be "Chunk 6" + "Chunk 7" (but they are now at orders 4 and 5)
assert "Chunk 6" in chunk5_expanded.content # Order 4 content assert "Chunk 6" in chunk5_expanded.content # Order 4 content
assert "Chunk 7" in chunk5_expanded.content # Order 5 content assert "Chunk 7" in chunk5_expanded.content # Order 5 content
assert "Chunk 0" not in chunk5_expanded.content assert "Chunk 0" not in chunk5_expanded.content

View file

@ -7,10 +7,10 @@ from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_create_document_with_chunks(qa_corpus: Dataset): async def test_create_document_with_chunks(qa_corpus: Dataset, temp_db_path):
"""Test creating a document with chunks from the qa_corpus using repository.""" """Test creating a document with chunks from the qa_corpus using repository."""
# Create an in-memory store and repository # Create a store and repository
store = Store(":memory:") store = Store(temp_db_path)
doc_repo = DocumentRepository(store) doc_repo = DocumentRepository(store)
# Get the first document from the corpus # Get the first document from the corpus
@ -23,58 +23,39 @@ async def test_create_document_with_chunks(qa_corpus: Dataset):
metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")}, metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")},
) )
# Convert text to DoclingDocument for chunk creation
from haiku.rag.utils import text_to_docling_document
docling_document = text_to_docling_document(document_text, name="test.md")
# Create the document with chunks in the database # Create the document with chunks in the database
created_document = await doc_repo.create(document) created_document = await doc_repo._create_with_docling(document, docling_document)
# Verify the document was created # Verify the document was created
assert created_document.id is not None assert created_document.id is not None
assert created_document.content == document_text assert created_document.content == document_text
# Check that chunks were created in the database # Check that chunks were created using repository
if store._connection is not None: from haiku.rag.store.repositories.chunk import ChunkRepository
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 chunk_repo = ChunkRepository(store)
chunks = await chunk_repo.get_by_document_id(created_document.id)
# Check that embeddings were created assert len(chunks) > 0
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 # Verify chunk metadata contains order information
cursor.execute( for i, chunk in enumerate(chunks):
"SELECT metadata FROM chunks WHERE document_id = ? ORDER BY id", assert "order" in chunk.metadata
(created_document.id,), assert chunk.metadata["order"] == i
)
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() store.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_document_repository_crud(qa_corpus: Dataset): async def test_document_repository_crud(qa_corpus: Dataset, temp_db_path):
"""Test CRUD operations in DocumentRepository.""" """Test CRUD operations in DocumentRepository."""
# Create an in-memory store and repository # Create a store and repository
store = Store(":memory:") store = Store(temp_db_path)
doc_repo = DocumentRepository(store) doc_repo = DocumentRepository(store)
# Get the first document from the corpus # Get the first document from the corpus

View file

@ -0,0 +1,86 @@
from unittest.mock import patch
import pytest
from haiku.rag.config import Config
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_lancedb_cloud_skips_optimization(temp_db_path):
"""Test that optimization is skipped when using LanceDB Cloud (db:// URI)."""
# Create a store
store = Store(temp_db_path)
chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# Create a document
document = Document(content="Test document content", metadata={})
created_document = await doc_repo.create(document)
document_id = created_document.id
# Mock LANCEDB_URI to simulate LanceDB Cloud usage
with patch.object(Config, "LANCEDB_URI", "db://test-database"):
# Mock the optimize method to track if it's called
with patch.object(store.chunks_table, "optimize") as mock_optimize:
# Create a chunk - this should trigger optimization logic
chunk = Chunk(
document_id=document_id,
content="Test chunk content",
metadata={"test": "value"},
)
created_chunk = await chunk_repo.create(chunk)
assert created_chunk.id is not None
# Wait a moment to ensure any async optimization would complete
import asyncio
await asyncio.sleep(0.1)
# The optimize method should NOT have been called for LanceDB Cloud
mock_optimize.assert_not_called()
store.close()
@pytest.mark.asyncio
async def test_local_storage_calls_optimization(temp_db_path):
"""Test that optimization is called for local storage."""
# Create a store
store = Store(temp_db_path)
chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# Create a document
document = Document(content="Test document content", metadata={})
created_document = await doc_repo.create(document)
document_id = created_document.id
# Ensure LANCEDB_URI is empty (local storage)
with patch.object(Config, "LANCEDB_URI", ""):
# Mock the optimize method to track if it's called
with patch.object(store.chunks_table, "optimize") as mock_optimize:
# Create a chunk - this should trigger optimization logic
chunk = Chunk(
document_id=document_id,
content="Test chunk content",
metadata={"test": "value"},
)
created_chunk = await chunk_repo.create(chunk)
assert created_chunk.id is not None
# Wait a moment to ensure async optimization completes
import asyncio
await asyncio.sleep(0.1)
# The optimize method SHOULD have been called for local storage
mock_optimize.assert_called()
store.close()

View file

@ -18,7 +18,7 @@ async def test_file_watcher_upsert_document():
temp_path.write_text("Test content for file watcher") temp_path.write_text("Test content for file watcher")
mock_client = AsyncMock(spec=HaikuRAG) mock_client = AsyncMock(spec=HaikuRAG)
mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri()) mock_doc = Document(id="1", content="Test content", uri=temp_path.as_uri())
mock_client.create_document_from_source.return_value = mock_doc mock_client.create_document_from_source.return_value = mock_doc
mock_client.get_document_by_uri.return_value = None # No existing document mock_client.get_document_by_uri.return_value = None # No existing document
@ -27,7 +27,7 @@ async def test_file_watcher_upsert_document():
result = await watcher._upsert_document(temp_path) result = await watcher._upsert_document(temp_path)
assert result is not None assert result is not None
assert result.id == 1 assert result.id == "1"
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
mock_client.create_document_from_source.assert_called_once_with(str(temp_path)) mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
@ -41,8 +41,10 @@ async def test_file_watcher_upsert_existing_document():
temp_path.write_text("Test content for file watcher") temp_path.write_text("Test content for file watcher")
mock_client = AsyncMock(spec=HaikuRAG) mock_client = AsyncMock(spec=HaikuRAG)
existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri()) existing_doc = Document(id="1", content="Old content", uri=temp_path.as_uri())
updated_doc = Document(id=1, content="Updated content", uri=temp_path.as_uri()) updated_doc = Document(
id="1", content="Updated content", uri=temp_path.as_uri()
)
mock_client.get_document_by_uri.return_value = existing_doc mock_client.get_document_by_uri.return_value = existing_doc
mock_client.create_document_from_source.return_value = updated_doc mock_client.create_document_from_source.return_value = updated_doc
@ -63,7 +65,7 @@ async def test_file_watcher_delete_document():
temp_path = Path("/tmp/test_file.txt") temp_path = Path("/tmp/test_file.txt")
mock_client = AsyncMock(spec=HaikuRAG) mock_client = AsyncMock(spec=HaikuRAG)
existing_doc = Document(id=1, content="Content to delete", uri=temp_path.as_uri()) existing_doc = Document(id="1", content="Content to delete", uri=temp_path.as_uri())
mock_client.get_document_by_uri.return_value = existing_doc mock_client.get_document_by_uri.return_value = existing_doc
mock_client.delete_document.return_value = True mock_client.delete_document.return_value = True
@ -72,7 +74,7 @@ async def test_file_watcher_delete_document():
await watcher._delete_document(temp_path) await watcher._delete_document(temp_path)
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
mock_client.delete_document.assert_called_once_with(1) mock_client.delete_document.assert_called_once_with("1")
@pytest.mark.asyncio @pytest.mark.asyncio

View file

@ -12,9 +12,9 @@ ANTHROPIC_AVAILABLE = bool(Config.ANTHROPIC_API_KEY)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_qa_ollama(qa_corpus: Dataset): async def test_qa_ollama(qa_corpus: Dataset, temp_db_path):
"""Test Ollama QA with LLM judge.""" """Test Ollama QA with LLM judge."""
client = HaikuRAG(":memory:") client = HaikuRAG(temp_db_path)
qa = QuestionAnswerAgent(client, "ollama", "qwen3") qa = QuestionAnswerAgent(client, "ollama", "qwen3")
llm_judge = LLMJudge() llm_judge = LLMJudge()
@ -36,9 +36,9 @@ async def test_qa_ollama(qa_corpus: Dataset):
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.skipif(not OPENAI_AVAILABLE, reason="OpenAI not available") @pytest.mark.skipif(not OPENAI_AVAILABLE, reason="OpenAI not available")
async def test_qa_openai(qa_corpus: Dataset): async def test_qa_openai(qa_corpus: Dataset, temp_db_path):
"""Test OpenAI QA with LLM judge.""" """Test OpenAI QA with LLM judge."""
client = HaikuRAG(":memory:") client = HaikuRAG(temp_db_path)
qa = QuestionAnswerAgent(client, "openai", "gpt-4o-mini") qa = QuestionAnswerAgent(client, "openai", "gpt-4o-mini")
llm_judge = LLMJudge() llm_judge = LLMJudge()
@ -60,9 +60,9 @@ async def test_qa_openai(qa_corpus: Dataset):
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.skipif(not ANTHROPIC_AVAILABLE, reason="Anthropic not available") @pytest.mark.skipif(not ANTHROPIC_AVAILABLE, reason="Anthropic not available")
async def test_qa_anthropic(qa_corpus: Dataset): async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path):
"""Test Anthropic QA with LLM judge.""" """Test Anthropic QA with LLM judge."""
client = HaikuRAG(":memory:") client = HaikuRAG(temp_db_path)
qa = QuestionAnswerAgent(client, "anthropic", "claude-3-5-haiku-20241022") qa = QuestionAnswerAgent(client, "anthropic", "claude-3-5-haiku-20241022")
llm_judge = LLMJudge() llm_judge = LLMJudge()

View file

@ -6,9 +6,9 @@ from haiku.rag.store.models.document import Document
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_rebuild_database(qa_corpus: Dataset): async def test_rebuild_database(qa_corpus: Dataset, temp_db_path):
"""Test rebuild functionality with existing documents.""" """Test rebuild functionality with existing documents."""
async with HaikuRAG(":memory:") as client: async with HaikuRAG(temp_db_path) as client:
created_docs: list[Document] = [] created_docs: list[Document] = []
for content in qa_corpus["document_extracted"][:3]: for content in qa_corpus["document_extracted"][:3]:
doc = await client.create_document( doc = await client.create_document(

View file

@ -7,7 +7,7 @@ from haiku.rag.store.models.chunk import Chunk
COHERE_AVAILABLE = bool(Config.COHERE_API_KEY) COHERE_AVAILABLE = bool(Config.COHERE_API_KEY)
chunks = [ chunks = [
Chunk(content=content, document_id=i) Chunk(content=content, document_id=str(i))
for i, content in enumerate( for i, content in enumerate(
[ [
"To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.", "To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.",
@ -24,7 +24,7 @@ chunks = [
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_reranker_base(): async def test_reranker_base():
reranker = RerankerBase() reranker = RerankerBase()
assert reranker._model == "qwen3" assert reranker._model == ""
with pytest.raises(NotImplementedError): with pytest.raises(NotImplementedError):
await reranker.rerank("query", []) await reranker.rerank("query", [])
@ -35,12 +35,16 @@ async def test_mxbai_reranker():
try: try:
from haiku.rag.reranking.mxbai import MxBAIReranker from haiku.rag.reranking.mxbai import MxBAIReranker
Config.RERANK_MODEL = "mixedbread-ai/mxbai-rerank-base-v2"
reranker = MxBAIReranker() reranker = MxBAIReranker()
# reranker._model = "mixedbread-ai/mxbai-rerank-base-v2"
reranked = await reranker.rerank( reranked = await reranker.rerank(
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
) )
assert [chunk.document_id for chunk, score in reranked] == [0, 2] assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
assert all(isinstance(score, float) for chunk, score in reranked) assert all(isinstance(score, float) for chunk, score in reranked)
Config.RERANK_MODEL = ""
except ImportError: except ImportError:
pytest.skip("MxBAI package not installed") pytest.skip("MxBAI package not installed")
@ -57,21 +61,8 @@ async def test_cohere_reranker():
reranked = await reranker.rerank( reranked = await reranker.rerank(
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
) )
assert [chunk.document_id for chunk, score in reranked] == [0, 2] assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
assert all(isinstance(score, float) for chunk, score in reranked) assert all(isinstance(score, float) for chunk, score in reranked)
except ImportError: except ImportError:
pytest.skip("Cohere package not installed") pytest.skip("Cohere package not installed")
@pytest.mark.asyncio
async def test_ollama_reranker():
from haiku.rag.reranking.ollama import OllamaReranker
reranker = OllamaReranker()
reranked = await reranker.rerank(
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
)
assert [chunk.document_id for chunk, score in reranked] == [0, 2]
assert all(isinstance(score, float) for chunk, score in reranked)

View file

@ -8,54 +8,60 @@ from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_search_qa_corpus(qa_corpus: Dataset): async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
"""Test that documents can be found by searching with their associated questions.""" """Test that documents can be found by searching with their associated questions."""
# Create an in-memory store and repositories # Create a store and repositories
store = Store(":memory:") store = Store(temp_db_path)
doc_repo = DocumentRepository(store) doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store) chunk_repo = ChunkRepository(store)
num_documents = 20
# Load first 10 documents with embeddings (reduced for faster testing) # Load unique documents (limited to 10)
seen_documents = set()
documents = [] documents = []
for i in range(num_documents):
doc_data = qa_corpus[i] for doc_data in qa_corpus:
document_text = doc_data["document_extracted"] if len(seen_documents) >= 10:
break
document_text = doc_data["document_extracted"] # type: ignore
document_id = doc_data.get("document_id", "") # type: ignore
if document_id in seen_documents:
continue
seen_documents.add(document_id)
# Create a Document instance # Create a Document instance
document = Document( document = Document(content=document_text)
content=document_text,
metadata={
"source": "qa_corpus",
"topic": doc_data.get("document_topic", ""),
"document_id": doc_data.get("document_id", ""),
"question": doc_data["question"],
},
)
# Create the document with chunks and embeddings # Create the document with chunks and embeddings
created_document = await doc_repo.create(document) from haiku.rag.utils import text_to_docling_document
docling_document = text_to_docling_document(document_text, name="test.md")
created_document = await doc_repo._create_with_docling(
document, docling_document
)
documents.append((created_document, doc_data)) documents.append((created_document, doc_data))
for i in range(5): # Test with first few documents # Test with first few unique documents
target_document, doc_data = documents[i]
for target_document, doc_data in documents:
question = doc_data["question"] question = doc_data["question"]
# Test vector search # Test vector search
vector_results = await chunk_repo.search_chunks(question, limit=5) vector_results = await chunk_repo.search(
question, limit=5, search_type="vector"
)
target_document_ids = {chunk.document_id for chunk, _ in vector_results} target_document_ids = {chunk.document_id for chunk, _ in vector_results}
assert target_document.id in target_document_ids assert target_document.id in target_document_ids
# Test FTS search # Test FTS search
fts_results = await chunk_repo.search_chunks_fts(question, limit=5) fts_results = await chunk_repo.search(question, limit=5, search_type="fts")
target_document_ids = {chunk.document_id for chunk, _ in fts_results} target_document_ids = {chunk.document_id for chunk, _ in fts_results}
assert target_document.id in target_document_ids assert target_document.id in target_document_ids
for i in range(num_documents): # Test with first few documents
target_document, doc_data = documents[i]
question = doc_data["question"]
# Test hybrid search # Test hybrid search
hybrid_results = await chunk_repo.search_chunks_hybrid(question, limit=5) hybrid_results = await chunk_repo.search(
question, limit=5, search_type="hybrid"
)
target_document_ids = {chunk.document_id for chunk, _ in hybrid_results} target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}
assert target_document.id in target_document_ids assert target_document.id in target_document_ids
@ -63,9 +69,9 @@ async def test_search_qa_corpus(qa_corpus: Dataset):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chunks_include_document_info(): async def test_chunks_include_document_info(temp_db_path):
"""Test that search results include document URI and metadata.""" """Test that search results include document URI and metadata."""
store = Store(":memory:") store = Store(temp_db_path)
doc_repo = DocumentRepository(store) doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store) chunk_repo = ChunkRepository(store)
@ -76,13 +82,21 @@ async def test_chunks_include_document_info():
metadata={"title": "Test Document", "author": "Test Author"}, metadata={"title": "Test Document", "author": "Test Author"},
) )
created_document = await doc_repo.create(document) # Create the document with chunks
from haiku.rag.utils import text_to_docling_document
docling_document = text_to_docling_document(document.content, name="test.md")
created_document = await doc_repo._create_with_docling(document, docling_document)
# Search for chunks # Search for chunks
results = await chunk_repo.search_chunks_hybrid("test document", limit=1) results = await chunk_repo.search("test document", limit=1, search_type="hybrid")
assert len(results) > 0 assert len(results) > 0
chunk, _ = results[0] chunk, score = results[0]
# Test that score is valid
assert isinstance(score, int | float), f"Score should be numeric, got {type(score)}"
assert score >= 0, f"Score should be non-negative, got {score}"
# Verify the chunk includes document information # Verify the chunk includes document information
assert chunk.document_uri == "https://example.com/test.html" assert chunk.document_uri == "https://example.com/test.html"
@ -90,3 +104,73 @@ async def test_chunks_include_document_info():
assert chunk.document_id == created_document.id assert chunk.document_id == created_document.id
store.close() store.close()
@pytest.mark.asyncio
async def test_search_score_types(temp_db_path):
"""Test that different search types return appropriate score ranges."""
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
# Create multiple documents with different content
documents_content = [
"Machine learning algorithms are powerful tools for data analysis and pattern recognition.",
"Deep learning neural networks can process complex datasets and identify hidden patterns.",
"Natural language processing enables computers to understand and generate human text.",
"Computer vision systems can interpret and analyze visual information from images.",
]
for content in documents_content:
document = Document(content=content)
from haiku.rag.utils import text_to_docling_document
docling_document = text_to_docling_document(content, name="test.md")
await doc_repo._create_with_docling(document, docling_document)
query = "machine learning"
# Test vector search scores (should be converted from distances)
vector_results = await chunk_repo.search(query, limit=3, search_type="vector")
assert len(vector_results) > 0
vector_scores = [score for _, score in vector_results]
# Test FTS search scores (should be native LanceDB FTS scores)
fts_results = await chunk_repo.search(query, limit=3, search_type="fts")
assert len(fts_results) > 0
fts_scores = [score for _, score in fts_results]
# Test hybrid search scores (should be native LanceDB relevance scores)
hybrid_results = await chunk_repo.search(query, limit=3, search_type="hybrid")
assert len(hybrid_results) > 0
hybrid_scores = [score for _, score in hybrid_results]
# All scores should be numeric and non-negative
for scores, search_type in [
(vector_scores, "vector"),
(fts_scores, "fts"),
(hybrid_scores, "hybrid"),
]:
for score in scores:
assert isinstance(score, int | float), (
f"{search_type} score should be numeric"
)
assert score >= 0, f"{search_type} score should be non-negative"
# Vector scores should typically be small (0-1 range due to distance conversion)
assert all(0 <= score <= 1 for score in vector_scores), (
"Vector scores should be in 0-1 range"
)
# Scores should be sorted in descending order (most relevant first)
for scores, search_type in [
(vector_scores, "vector"),
(fts_scores, "fts"),
(hybrid_scores, "hybrid"),
]:
for i in range(len(scores) - 1):
assert scores[i] >= scores[i + 1], (
f"{search_type} results should be sorted by score descending"
)
store.close()

View file

@ -1,55 +1,59 @@
import tempfile
from pathlib import Path
import pytest import pytest
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import ( from haiku.rag.store.repositories.settings import (
ConfigMismatchError, ConfigMismatchError,
SettingsRepository,
) )
def test_settings_table_populated_on_store_init(): def test_settings_table_populated_on_store_init(temp_db_path):
"""Test that settings table is populated with current config when store is initialized.""" """Test that settings table is populated with current config when store is initialized."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
store = Store(":memory:") store = Store(temp_db_path)
settings_repo = SettingsRepository(store) settings_repo = SettingsRepository(store)
db_settings = settings_repo.get() db_settings = settings_repo.get_current_settings()
config_dict = Config.model_dump(mode="json") config_dict = Config.model_dump(mode="json")
assert db_settings == config_dict # Remove version from db_settings since it's added automatically
db_settings_without_version = {
k: v for k, v in db_settings.items() if k != "version"
}
assert db_settings_without_version == config_dict
store.close() store.close()
def test_settings_save_and_retrieve(): def test_settings_save_and_retrieve(temp_db_path):
"""Test saving and retrieving settings after config change.""" """Test saving and retrieving settings after config change."""
store = Store(":memory:") from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
store = Store(temp_db_path)
settings_repo = SettingsRepository(store) settings_repo = SettingsRepository(store)
original_chunk_size = Config.CHUNK_SIZE original_chunk_size = Config.CHUNK_SIZE
Config.CHUNK_SIZE = 2 * original_chunk_size Config.CHUNK_SIZE = 2 * original_chunk_size
settings_repo.save() settings_repo.save_current_settings()
retrieved_settings = settings_repo.get() retrieved_settings = settings_repo.get_current_settings()
assert retrieved_settings["CHUNK_SIZE"] == 2 * original_chunk_size assert retrieved_settings["CHUNK_SIZE"] == 2 * original_chunk_size
Config.CHUNK_SIZE = original_chunk_size Config.CHUNK_SIZE = original_chunk_size
store.close() store.close()
async def test_config_validation_on_db_load(): @pytest.mark.skip(reason="Config validation not fully implemented for LanceDB")
async def test_config_validation_on_db_load(temp_db_path):
"""Test that config validation fails when loading db with mismatched settings.""" """Test that config validation fails when loading db with mismatched settings."""
# Create a temporary database file from haiku.rag.store.engine import Store
with tempfile.NamedTemporaryFile(suffix=".sqlite") as tmp: from haiku.rag.store.repositories.settings import SettingsRepository
db_path = Path(tmp.name)
# Create store and save settings # Create store and save settings
store1 = Store(db_path) store1 = Store(temp_db_path)
store1.close() store1.close()
# Change config # Change config
@ -59,20 +63,20 @@ async def test_config_validation_on_db_load():
try: try:
# Loading the database should raise ConfigMismatchError # Loading the database should raise ConfigMismatchError
with pytest.raises(ConfigMismatchError) as exc_info: with pytest.raises(ConfigMismatchError) as exc_info:
Store(db_path) Store(temp_db_path)
assert "CHUNK_SIZE" in str(exc_info.value) assert "CHUNK_SIZE" in str(exc_info.value)
assert "Consider rebuilding" in str(exc_info.value) assert "Consider rebuilding" in str(exc_info.value)
# Rebuild # Rebuild
async with HaikuRAG(db_path=db_path, skip_validation=True) as client: async with HaikuRAG(db_path=temp_db_path, skip_validation=True) as client:
async for _ in client.rebuild_database(): async for _ in client.rebuild_database():
pass # Process all documents pass # Process all documents
# Verify we can now load the database without exception (settings were updated) # Verify we can now load the database without exception (settings were updated)
store2 = Store(db_path) store2 = Store(temp_db_path)
settings_repo2 = SettingsRepository(store2) settings_repo2 = SettingsRepository(store2)
db_settings = settings_repo2.get() db_settings = settings_repo2.get_current_settings()
assert db_settings["CHUNK_SIZE"] == 999 assert db_settings["CHUNK_SIZE"] == 999
store2.close() store2.close()

View file

@ -1,22 +1,4 @@
from haiku.rag.utils import ( from haiku.rag.utils import text_to_docling_document
int_to_semantic_version,
semantic_version_to_int,
text_to_docling_document,
)
def test_sqlite_user_version():
version = "0.1.5"
assert semantic_version_to_int(version) == 261
assert int_to_semantic_version(261) == version
version = "0.0.0"
assert semantic_version_to_int(version) == 0
assert int_to_semantic_version(0) == version
version = "255.255.255"
assert semantic_version_to_int(version) == 16777215
assert int_to_semantic_version(16777215) == version
def test_text_to_docling_document(): def test_text_to_docling_document():

68
uv.lock
View file

@ -567,6 +567,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/20/34/a08b0ee99715eaba118cbe19a71f7b5e2425c2718ef96007c325944a1152/datasets-3.6.0-py3-none-any.whl", hash = "sha256:25000c4a2c0873a710df127d08a202a06eab7bf42441a6bc278b499c2f72cd1b", size = 491546, upload-time = "2025-05-07T15:14:59.742Z" }, { url = "https://files.pythonhosted.org/packages/20/34/a08b0ee99715eaba118cbe19a71f7b5e2425c2718ef96007c325944a1152/datasets-3.6.0-py3-none-any.whl", hash = "sha256:25000c4a2c0873a710df127d08a202a06eab7bf42441a6bc278b499c2f72cd1b", size = 491546, upload-time = "2025-05-07T15:14:59.742Z" },
] ]
[[package]]
name = "deprecation"
version = "2.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "packaging" },
]
sdist = { url = "https://files.pythonhosted.org/packages/5a/d3/8ae2869247df154b64c1884d7346d412fed0c49df84db635aab2d1c40e62/deprecation-2.1.0.tar.gz", hash = "sha256:72b3bde64e5d778694b0cf68178aed03d15e15477116add3fb773e581f9518ff", size = 173788, upload-time = "2020-04-20T14:23:38.738Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/02/c3/253a89ee03fc9b9682f1541728eb66db7db22148cd94f89ab22528cd1e1b/deprecation-2.1.0-py2.py3-none-any.whl", hash = "sha256:a10811591210e1fb0e768a8c25517cabeabcba6f0bf96564f8ff45189f90b14a", size = 11178, upload-time = "2020-04-20T14:23:36.581Z" },
]
[[package]] [[package]]
name = "dill" name = "dill"
version = "0.3.8" version = "0.3.8"
@ -1015,12 +1027,12 @@ dependencies = [
{ name = "docling" }, { name = "docling" },
{ name = "fastmcp" }, { name = "fastmcp" },
{ name = "httpx" }, { name = "httpx" },
{ name = "lancedb" },
{ name = "ollama" }, { name = "ollama" },
{ name = "pydantic" }, { name = "pydantic" },
{ name = "pydantic-ai" }, { name = "pydantic-ai" },
{ name = "python-dotenv" }, { name = "python-dotenv" },
{ name = "rich" }, { name = "rich" },
{ name = "sqlite-vec" },
{ name = "tiktoken" }, { name = "tiktoken" },
{ name = "typer" }, { name = "typer" },
{ name = "watchfiles" }, { name = "watchfiles" },
@ -1052,13 +1064,13 @@ requires-dist = [
{ name = "docling", specifier = ">=2.15.0" }, { name = "docling", specifier = ">=2.15.0" },
{ name = "fastmcp", specifier = ">=2.8.1" }, { name = "fastmcp", specifier = ">=2.8.1" },
{ name = "httpx", specifier = ">=0.28.1" }, { name = "httpx", specifier = ">=0.28.1" },
{ name = "lancedb", specifier = ">=0.17.0" },
{ name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" },
{ name = "ollama", specifier = ">=0.5.3" }, { name = "ollama", specifier = ">=0.5.3" },
{ name = "pydantic", specifier = ">=2.11.7" }, { name = "pydantic", specifier = ">=2.11.7" },
{ name = "pydantic-ai", specifier = ">=0.7.2" }, { name = "pydantic-ai", specifier = ">=0.7.2" },
{ name = "python-dotenv", specifier = ">=1.1.0" }, { name = "python-dotenv", specifier = ">=1.1.0" },
{ name = "rich", specifier = ">=14.0.0" }, { name = "rich", specifier = ">=14.0.0" },
{ name = "sqlite-vec", specifier = ">=0.1.6" },
{ name = "tiktoken", specifier = ">=0.9.0" }, { name = "tiktoken", specifier = ">=0.9.0" },
{ name = "typer", specifier = ">=0.16.0" }, { name = "typer", specifier = ">=0.16.0" },
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" },
@ -1072,7 +1084,7 @@ dev = [
{ name = "mkdocs", specifier = ">=1.6.1" }, { name = "mkdocs", specifier = ">=1.6.1" },
{ name = "mkdocs-material", specifier = ">=9.6.14" }, { name = "mkdocs-material", specifier = ">=9.6.14" },
{ name = "pre-commit", specifier = ">=4.2.0" }, { name = "pre-commit", specifier = ">=4.2.0" },
{ name = "pyright", specifier = ">=1.1.403" }, { name = "pyright", specifier = ">=1.1.404" },
{ name = "pytest", specifier = ">=8.4.0" }, { name = "pytest", specifier = ">=8.4.0" },
{ name = "pytest-asyncio", specifier = ">=1.0.0" }, { name = "pytest-asyncio", specifier = ">=1.0.0" },
{ name = "pytest-cov", specifier = ">=6.2.1" }, { name = "pytest-cov", specifier = ">=6.2.1" },
@ -1336,6 +1348,29 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/01/0e/b27cdbaccf30b890c40ed1da9fd4a3593a5cf94dae54fb34f8a4b74fcd3f/jsonschema_specifications-2025.4.1-py3-none-any.whl", hash = "sha256:4653bffbd6584f7de83a67e0d620ef16900b390ddc7939d56684d6c81e33f1af", size = 18437, upload-time = "2025-04-23T12:34:05.422Z" }, { url = "https://files.pythonhosted.org/packages/01/0e/b27cdbaccf30b890c40ed1da9fd4a3593a5cf94dae54fb34f8a4b74fcd3f/jsonschema_specifications-2025.4.1-py3-none-any.whl", hash = "sha256:4653bffbd6584f7de83a67e0d620ef16900b390ddc7939d56684d6c81e33f1af", size = 18437, upload-time = "2025-04-23T12:34:05.422Z" },
] ]
[[package]]
name = "lancedb"
version = "0.24.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "deprecation" },
{ name = "numpy" },
{ name = "overrides" },
{ name = "packaging" },
{ name = "pyarrow" },
{ name = "pydantic" },
{ name = "tqdm" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/73/41/6e51c15c0ea9dc19c6eb037a0c284ccae49fc2e8425b934a6a1ccdc0e6d6/lancedb-0.24.3-cp39-abi3-macosx_10_15_x86_64.whl", hash = "sha256:65b5aa6201d6f1d921694dce614a2686c1272e25926cc7809ad8eb89a8eb85fa", size = 33401550, upload-time = "2025-08-15T19:02:35.274Z" },
{ url = "https://files.pythonhosted.org/packages/1b/78/d020464db8f189923caba69eaa9281c8b0a6c5c6cef3841cca3a17feb00e/lancedb-0.24.3-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:f2b86928b3175afde9507ca2e39f545df380ca52c415511ffd52fab8b09937ec", size = 30810069, upload-time = "2025-08-15T19:15:55.771Z" },
{ url = "https://files.pythonhosted.org/packages/b7/30/3317ebf9c6397591b9beccd091a87e23f7d8186ddaa7cd0c1aa318e09323/lancedb-0.24.3-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e7756d020c2e6d22130d8f59f32b816fbfda5f65ecc80c5b1bbeb01bf55dc834", size = 31708034, upload-time = "2025-08-15T18:23:36.783Z" },
{ url = "https://files.pythonhosted.org/packages/7b/61/b91a7bfe0138cec86f2e9930160112f38da923333da16f2bfaac85649d45/lancedb-0.24.3-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0c933d2680e7bb28d99cb32b2395e808593a845a2b5ab59030e7d840d97e91f2", size = 34916331, upload-time = "2025-08-15T18:26:11.104Z" },
{ url = "https://files.pythonhosted.org/packages/e0/c6/b475d3addf841f803f7c54a64427ef0e973ce340295f59adbd0358097a69/lancedb-0.24.3-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:60ecfaabd1d33f435b498624f988e610cafacd5dfea26368e5582647730495e0", size = 31718200, upload-time = "2025-08-15T18:22:17.472Z" },
{ url = "https://files.pythonhosted.org/packages/9a/b9/3e0e25b7c6dcd4f6b0e977cb886965070ca05d7994819834484cbe8c8d00/lancedb-0.24.3-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:78632bd98e6c317609ce844eb737d1dd306786291ed2d98009e37377be13e1b4", size = 34963956, upload-time = "2025-08-15T18:25:47.58Z" },
{ url = "https://files.pythonhosted.org/packages/c4/3b/063e78eaf61ee8c1428e1063d9c29c5fab8b2d7f5d324a13716419a29a57/lancedb-0.24.3-cp39-abi3-win_amd64.whl", hash = "sha256:d5eb70b8a3b66d728c183f9b77ddfed13894c12880f647bf5cd85e477ad5368b", size = 36945181, upload-time = "2025-08-15T18:44:06.182Z" },
]
[[package]] [[package]]
name = "latex2mathml" name = "latex2mathml"
version = "3.78.0" version = "3.78.0"
@ -2097,6 +2132,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/bb/ee/6b08dde0a022c463b88f55ae81149584b125a42183407dc1045c486cc870/opentelemetry_api-1.36.0-py3-none-any.whl", hash = "sha256:02f20bcacf666e1333b6b1f04e647dc1d5111f86b8e510238fcc56d7762cda8c", size = 65564, upload-time = "2025-07-29T15:11:47.998Z" }, { url = "https://files.pythonhosted.org/packages/bb/ee/6b08dde0a022c463b88f55ae81149584b125a42183407dc1045c486cc870/opentelemetry_api-1.36.0-py3-none-any.whl", hash = "sha256:02f20bcacf666e1333b6b1f04e647dc1d5111f86b8e510238fcc56d7762cda8c", size = 65564, upload-time = "2025-07-29T15:11:47.998Z" },
] ]
[[package]]
name = "overrides"
version = "7.7.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/36/86/b585f53236dec60aba864e050778b25045f857e17f6e5ea0ae95fe80edd2/overrides-7.7.0.tar.gz", hash = "sha256:55158fa3d93b98cc75299b1e67078ad9003ca27945c76162c1c0766d6f91820a", size = 22812, upload-time = "2024-01-27T21:01:33.423Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/2c/ab/fc8290c6a4c722e5514d80f62b2dc4c4df1a68a41d1364e625c35990fcf3/overrides-7.7.0-py3-none-any.whl", hash = "sha256:c7ed9d062f78b8e4c1a7b70bd8796b35ead4d9f510227ef9c5dc7626c60d7e49", size = 17832, upload-time = "2024-01-27T21:01:31.393Z" },
]
[[package]] [[package]]
name = "packaging" name = "packaging"
version = "25.0" version = "25.0"
@ -2729,15 +2773,15 @@ wheels = [
[[package]] [[package]]
name = "pyright" name = "pyright"
version = "1.1.403" version = "1.1.404"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
dependencies = [ dependencies = [
{ name = "nodeenv" }, { name = "nodeenv" },
{ name = "typing-extensions" }, { name = "typing-extensions" },
] ]
sdist = { url = "https://files.pythonhosted.org/packages/fe/f6/35f885264ff08c960b23d1542038d8da86971c5d8c955cfab195a4f672d7/pyright-1.1.403.tar.gz", hash = "sha256:3ab69b9f41c67fb5bbb4d7a36243256f0d549ed3608678d381d5f51863921104", size = 3913526, upload-time = "2025-07-09T07:15:52.882Z" } sdist = { url = "https://files.pythonhosted.org/packages/e2/6e/026be64c43af681d5632722acd100b06d3d39f383ec382ff50a71a6d5bce/pyright-1.1.404.tar.gz", hash = "sha256:455e881a558ca6be9ecca0b30ce08aa78343ecc031d37a198ffa9a7a1abeb63e", size = 4065679, upload-time = "2025-08-20T18:46:14.029Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/49/b6/b04e5c2f41a5ccad74a1a4759da41adb20b4bc9d59a5e08d29ba60084d07/pyright-1.1.403-py3-none-any.whl", hash = "sha256:c0eeca5aa76cbef3fcc271259bbd785753c7ad7bcac99a9162b4c4c7daed23b3", size = 5684504, upload-time = "2025-07-09T07:15:50.958Z" }, { url = "https://files.pythonhosted.org/packages/84/30/89aa7f7d7a875bbb9a577d4b1dc5a3e404e3d2ae2657354808e905e358e0/pyright-1.1.404-py3-none-any.whl", hash = "sha256:c7b7ff1fdb7219c643079e4c3e7d4125f0dafcc19d253b47e898d130ea426419", size = 5902951, upload-time = "2025-08-20T18:46:12.096Z" },
] ]
[[package]] [[package]]
@ -3455,18 +3499,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e7/9c/0e6afc12c269578be5c0c1c9f4b49a8d32770a080260c333ac04cc1c832d/soupsieve-2.7-py3-none-any.whl", hash = "sha256:6e60cc5c1ffaf1cebcc12e8188320b72071e922c2e897f737cadce79ad5d30c4", size = 36677, upload-time = "2025-04-20T18:50:07.196Z" }, { url = "https://files.pythonhosted.org/packages/e7/9c/0e6afc12c269578be5c0c1c9f4b49a8d32770a080260c333ac04cc1c832d/soupsieve-2.7-py3-none-any.whl", hash = "sha256:6e60cc5c1ffaf1cebcc12e8188320b72071e922c2e897f737cadce79ad5d30c4", size = 36677, upload-time = "2025-04-20T18:50:07.196Z" },
] ]
[[package]]
name = "sqlite-vec"
version = "0.1.6"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/88/ed/aabc328f29ee6814033d008ec43e44f2c595447d9cccd5f2aabe60df2933/sqlite_vec-0.1.6-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:77491bcaa6d496f2acb5cc0d0ff0b8964434f141523c121e313f9a7d8088dee3", size = 164075, upload-time = "2024-11-20T16:40:29.847Z" },
{ url = "https://files.pythonhosted.org/packages/a7/57/05604e509a129b22e303758bfa062c19afb020557d5e19b008c64016704e/sqlite_vec-0.1.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fdca35f7ee3243668a055255d4dee4dea7eed5a06da8cad409f89facf4595361", size = 165242, upload-time = "2024-11-20T16:40:31.206Z" },
{ url = "https://files.pythonhosted.org/packages/f2/48/dbb2cc4e5bad88c89c7bb296e2d0a8df58aab9edc75853728c361eefc24f/sqlite_vec-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b0519d9cd96164cd2e08e8eed225197f9cd2f0be82cb04567692a0a4be02da3", size = 103704, upload-time = "2024-11-20T16:40:33.729Z" },
{ url = "https://files.pythonhosted.org/packages/80/76/97f33b1a2446f6ae55e59b33869bed4eafaf59b7f4c662c8d9491b6a714a/sqlite_vec-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:823b0493add80d7fe82ab0fe25df7c0703f4752941aee1c7b2b02cec9656cb24", size = 151556, upload-time = "2024-11-20T16:40:35.387Z" },
{ url = "https://files.pythonhosted.org/packages/6a/98/e8bc58b178266eae2fcf4c9c7a8303a8d41164d781b32d71097924a6bebe/sqlite_vec-0.1.6-py3-none-win_amd64.whl", hash = "sha256:c65bcfd90fa2f41f9000052bcb8bb75d38240b2dae49225389eca6c3136d3f0c", size = 281540, upload-time = "2024-11-20T16:40:37.296Z" },
]
[[package]] [[package]]
name = "sse-starlette" name = "sse-starlette"
version = "2.3.6" version = "2.3.6"