VoyageAI embeddings

This commit is contained in:
Yiorgis Gozadinos 2025-06-18 09:42:35 +02:00
parent 51723308ed
commit a1079fd1f5
No known key found for this signature in database
6 changed files with 132 additions and 51 deletions

3
.gitignore vendored
View file

@ -14,3 +14,6 @@ wheels/
tests/data/ tests/data/
.pytest_cache/ .pytest_cache/
.ruff_cache/ .ruff_cache/
# environment variables
.env

View file

@ -3,14 +3,12 @@
A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient document storage, chunking, and hybrid search capabilities. A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient document storage, chunking, and hybrid search capabilities.
## Features ## Features
- **Local SQLite**: No need to run additional servers
- **Document Management**: Store and manage documents with automatic content parsing - **Support for various embedding providers**: You can use Ollama, VoyageAI, OpenAI or add your own
- **Smart Updates**: Intelligent file/URL monitoring with MD5-based change detection
- **Hybrid Search**: Full-text search (FTS5) combined with vector embeddings
- **Multi-format Support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, and more
- **Web Content**: Direct URL ingestion with automatic content type detection
- **Vector Embeddings**: Uses sqlite-vec for efficient similarity search - **Vector Embeddings**: Uses sqlite-vec for efficient similarity search
- **Automatic Chunking**: Intelligent document segmentation for better retrieval - **Hybrid Search**: Full-text search (FTS5) combined with vector embeddings using Reciprocal Rank Fusion
- **Multi-format Support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more
- **Web Content**: Direct URL ingestion with automatic content type detection
## Installation ## Installation
@ -18,14 +16,28 @@ A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient d
uv pip install haiku.rag uv pip install haiku.rag
``` ```
or for development, checkout the repository and then, By default Ollama (with the `mxbai-embed-large` model) is used for the embeddings.
For other providers use:
- **VoyageAI**: `uv pip install haiku.rag --extra voyageai`
## Configuration
If you want to use an alternative embeddings provider (Ollama being the default) you will need to set the provider details through environment variables:
By default:
```bash ```bash
# Install dependencies EMBEDDING_PROVIDER="ollama"
uv sync EMBEDDING_MODEL="mxbai-embed-large" # or any other model
EMBEDDING_VECTOR_DIM=1024
```
# Activate virtual environment For VoyageAI:
source .venv/bin/activate ```bash
EMBEDDING_PROVIDER="voyageai"
EMBEDDING_MODEL="voyage-3.5" # or any other model
EMBEDDING_VECTOR_DIM=1024
``` ```
## Quick Start ## Quick Start
@ -90,14 +102,11 @@ finally:
```python ```python
async with HaikuRAG("database.db") as client: async with HaikuRAG("database.db") as client:
# Basic search
results = await client.search("your query here")
# Search with custom parameters
results = await client.search( results = await client.search(
query="machine learning", query="machine learning",
limit=10, # Maximum results to return limit=5, # Maximum results to return, defaults to 5
k=60 # RRF parameter for reciprocal rank fusion k=60 # RRF parameter for reciprocal rank fusion, defaults to 60
) )
# Process results # Process results
@ -107,29 +116,10 @@ async with HaikuRAG("database.db") as client:
print(f"From document: {chunk.document_id}") print(f"From document: {chunk.document_id}")
``` ```
## Smart Document Updates
The system automatically tracks file changes using MD5 hashes:
```python
async with HaikuRAG("database.db") as client:
# First call - creates new document
doc1 = await client.create_document_from_source("document.txt")
# Second call - no changes, returns existing document (no processing)
doc2 = await client.create_document_from_source("document.txt")
assert doc1.id == doc2.id
# After file modification - automatically updates existing document
# File content changed...
doc3 = await client.create_document_from_source("document.txt")
assert doc1.id == doc3.id # Same document
assert doc3.content != doc1.content # Updated content
```
## Supported File Formats ## Supported File Formats
The system supports 40+ file formats through MarkItDown: `haiku.rag` supports 40+ file formats through MarkItDown:
- **Documents**: PDF, DOCX, PPTX, XLSX - **Documents**: PDF, DOCX, PPTX, XLSX
- **Web**: HTML, XML - **Web**: HTML, XML
@ -137,20 +127,6 @@ The system supports 40+ file formats through MarkItDown:
- **Code**: PY, JS, TS, C, CPP, JAVA, GO, RS, and more - **Code**: PY, JS, TS, C, CPP, JAVA, GO, RS, and more
- **Media**: MP3, WAV (transcription) - **Media**: MP3, WAV (transcription)
## Document Metadata
Documents automatically include metadata:
```python
doc = await client.create_document_from_source("example.pdf")
print(doc.metadata)
# {
# "contentType": "application/pdf",
# "md5": "abc123...",
# "custom_field": "value" # Your custom metadata
# }
```
## Contributing ## Contributing
1. Fork the repository 1. Fork the repository

View file

@ -16,6 +16,9 @@ dependencies = [
"watchfiles>=1.1.0", "watchfiles>=1.1.0",
] ]
[project.optional-dependencies]
voyageai = ["voyageai>=0.3.2"]
[build-system] [build-system]
requires = ["hatchling"] requires = ["hatchling"]
build-backend = "hatchling.build" build-backend = "hatchling.build"

View file

@ -1,6 +1,7 @@
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.embeddings.base import EmbedderBase from haiku.rag.embeddings.base import EmbedderBase
from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder
from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
def get_embedder() -> EmbedderBase: def get_embedder() -> EmbedderBase:
@ -10,4 +11,15 @@ def get_embedder() -> EmbedderBase:
if Config.EMBEDDING_PROVIDER == "ollama": if Config.EMBEDDING_PROVIDER == "ollama":
return OllamaEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM) return OllamaEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM)
if Config.EMBEDDING_PROVIDER == "voyageai":
try:
import voyageai
except ImportError:
raise ImportError(
"VoyageAI embedder requires the 'voyageai' package. "
"Please install haiku.rag with the 'voyageai' extra:"
"uv pip install haiku.rag --extra voyageai"
)
return VoyageAIEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM)
raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDING_PROVIDER}") raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDING_PROVIDER}")

View file

@ -0,0 +1,14 @@
from voyageai.client import Client
from haiku.rag.config import Config
from haiku.rag.embeddings.base import EmbedderBase
class Embedder(EmbedderBase):
_model: str = Config.EMBEDDING_MODEL
_vector_dim: int = 1024
async def embed(self, text: str) -> list[float]:
client = Client()
res = client.embed([text], model=self._model, output_dtype="float")
return res.embeddings[0] # type: ignore[return-value]

73
uv.lock
View file

@ -45,6 +45,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/9d/47/b11d0089875a23bff0abd3edb5516bcd454db3fefab8604f5e4b07bd6210/aiohttp-3.12.13-cp313-cp313-win_amd64.whl", hash = "sha256:5a178390ca90419bfd41419a809688c368e63c86bd725e1186dd97f6b89c2706", size = 446735, upload-time = "2025-06-14T15:15:02.858Z" }, { url = "https://files.pythonhosted.org/packages/9d/47/b11d0089875a23bff0abd3edb5516bcd454db3fefab8604f5e4b07bd6210/aiohttp-3.12.13-cp313-cp313-win_amd64.whl", hash = "sha256:5a178390ca90419bfd41419a809688c368e63c86bd725e1186dd97f6b89c2706", size = 446735, upload-time = "2025-06-14T15:15:02.858Z" },
] ]
[[package]]
name = "aiolimiter"
version = "1.2.1"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/f1/23/b52debf471f7a1e42e362d959a3982bdcb4fe13a5d46e63d28868807a79c/aiolimiter-1.2.1.tar.gz", hash = "sha256:e02a37ea1a855d9e832252a105420ad4d15011505512a1a1d814647451b5cca9", size = 7185, upload-time = "2024-12-08T15:31:51.496Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f3/ba/df6e8e1045aebc4778d19b8a3a9bc1808adb1619ba94ca354d9ba17d86c3/aiolimiter-1.2.1-py3-none-any.whl", hash = "sha256:d3f249e9059a20badcb56b61601a83556133655c11d1eb3dd3e04ff069e5f3c7", size = 6711, upload-time = "2024-12-08T15:31:49.874Z" },
]
[[package]] [[package]]
name = "aiosignal" name = "aiosignal"
version = "1.3.2" version = "1.3.2"
@ -470,6 +479,14 @@ dependencies = [
{ name = "watchfiles" }, { name = "watchfiles" },
] ]
[package.optional-dependencies]
all = [
{ name = "voyageai" },
]
voyageai = [
{ name = "voyageai" },
]
[package.dev-dependencies] [package.dev-dependencies]
dev = [ dev = [
{ name = "datasets" }, { name = "datasets" },
@ -490,8 +507,11 @@ requires-dist = [
{ name = "python-dotenv", specifier = ">=1.1.0" }, { name = "python-dotenv", specifier = ">=1.1.0" },
{ name = "sqlite-vec", specifier = ">=0.1.6" }, { name = "sqlite-vec", specifier = ">=0.1.6" },
{ name = "tiktoken", specifier = ">=0.9.0" }, { name = "tiktoken", specifier = ">=0.9.0" },
{ name = "voyageai", marker = "extra == 'all'", specifier = ">=0.3.2" },
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" },
{ name = "watchfiles", specifier = ">=1.1.0" }, { name = "watchfiles", specifier = ">=1.1.0" },
] ]
provides-extras = ["voyageai", "all"]
[package.metadata.requires-dev] [package.metadata.requires-dev]
dev = [ dev = [
@ -1403,6 +1423,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
] ]
[[package]]
name = "tenacity"
version = "9.1.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/0a/d4/2b0cd0fe285e14b36db076e78c93766ff1d529d70408bd1d2a5a84f1d929/tenacity-9.1.2.tar.gz", hash = "sha256:1169d376c297e7de388d18b4481760d478b0e99a777cad3a9c86e556f4b697cb", size = 48036, upload-time = "2025-04-02T08:25:09.966Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/30/643397144bfbfec6f6ef821f36f33e57d35946c44a2352d3c9f0ae847619/tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138", size = 28248, upload-time = "2025-04-02T08:25:07.678Z" },
]
[[package]] [[package]]
name = "tiktoken" name = "tiktoken"
version = "0.9.0" version = "0.9.0"
@ -1421,6 +1450,31 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/de/a8/8f499c179ec900783ffe133e9aab10044481679bb9aad78436d239eee716/tiktoken-0.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:5ea0edb6f83dc56d794723286215918c1cde03712cbbafa0348b33448faf5b95", size = 894669, upload-time = "2025-02-14T06:02:47.341Z" }, { url = "https://files.pythonhosted.org/packages/de/a8/8f499c179ec900783ffe133e9aab10044481679bb9aad78436d239eee716/tiktoken-0.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:5ea0edb6f83dc56d794723286215918c1cde03712cbbafa0348b33448faf5b95", size = 894669, upload-time = "2025-02-14T06:02:47.341Z" },
] ]
[[package]]
name = "tokenizers"
version = "0.21.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "huggingface-hub" },
]
sdist = { url = "https://files.pythonhosted.org/packages/92/76/5ac0c97f1117b91b7eb7323dcd61af80d72f790b4df71249a7850c195f30/tokenizers-0.21.1.tar.gz", hash = "sha256:a1bb04dc5b448985f86ecd4b05407f5a8d97cb2c0532199b2a302a604a0165ab", size = 343256, upload-time = "2025-03-13T10:51:18.189Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a5/1f/328aee25f9115bf04262e8b4e5a2050b7b7cf44b59c74e982db7270c7f30/tokenizers-0.21.1-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:e78e413e9e668ad790a29456e677d9d3aa50a9ad311a40905d6861ba7692cf41", size = 2780767, upload-time = "2025-03-13T10:51:09.459Z" },
{ url = "https://files.pythonhosted.org/packages/ae/1a/4526797f3719b0287853f12c5ad563a9be09d446c44ac784cdd7c50f76ab/tokenizers-0.21.1-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:cd51cd0a91ecc801633829fcd1fda9cf8682ed3477c6243b9a095539de4aecf3", size = 2650555, upload-time = "2025-03-13T10:51:07.692Z" },
{ url = "https://files.pythonhosted.org/packages/4d/7a/a209b29f971a9fdc1da86f917fe4524564924db50d13f0724feed37b2a4d/tokenizers-0.21.1-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:28da6b72d4fb14ee200a1bd386ff74ade8992d7f725f2bde2c495a9a98cf4d9f", size = 2937541, upload-time = "2025-03-13T10:50:56.679Z" },
{ url = "https://files.pythonhosted.org/packages/3c/1e/b788b50ffc6191e0b1fc2b0d49df8cff16fe415302e5ceb89f619d12c5bc/tokenizers-0.21.1-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:34d8cfde551c9916cb92014e040806122295a6800914bab5865deb85623931cf", size = 2819058, upload-time = "2025-03-13T10:50:59.525Z" },
{ url = "https://files.pythonhosted.org/packages/36/aa/3626dfa09a0ecc5b57a8c58eeaeb7dd7ca9a37ad9dd681edab5acd55764c/tokenizers-0.21.1-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:aaa852d23e125b73d283c98f007e06d4595732104b65402f46e8ef24b588d9f8", size = 3133278, upload-time = "2025-03-13T10:51:04.678Z" },
{ url = "https://files.pythonhosted.org/packages/a4/4d/8fbc203838b3d26269f944a89459d94c858f5b3f9a9b6ee9728cdcf69161/tokenizers-0.21.1-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a21a15d5c8e603331b8a59548bbe113564136dc0f5ad8306dd5033459a226da0", size = 3144253, upload-time = "2025-03-13T10:51:01.261Z" },
{ url = "https://files.pythonhosted.org/packages/d8/1b/2bd062adeb7c7511b847b32e356024980c0ffcf35f28947792c2d8ad2288/tokenizers-0.21.1-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2fdbd4c067c60a0ac7eca14b6bd18a5bebace54eb757c706b47ea93204f7a37c", size = 3398225, upload-time = "2025-03-13T10:51:03.243Z" },
{ url = "https://files.pythonhosted.org/packages/8a/63/38be071b0c8e06840bc6046991636bcb30c27f6bb1e670f4f4bc87cf49cc/tokenizers-0.21.1-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2dd9a0061e403546f7377df940e866c3e678d7d4e9643d0461ea442b4f89e61a", size = 3038874, upload-time = "2025-03-13T10:51:06.235Z" },
{ url = "https://files.pythonhosted.org/packages/ec/83/afa94193c09246417c23a3c75a8a0a96bf44ab5630a3015538d0c316dd4b/tokenizers-0.21.1-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:db9484aeb2e200c43b915a1a0150ea885e35f357a5a8fabf7373af333dcc8dbf", size = 9014448, upload-time = "2025-03-13T10:51:10.927Z" },
{ url = "https://files.pythonhosted.org/packages/ae/b3/0e1a37d4f84c0f014d43701c11eb8072704f6efe8d8fc2dcdb79c47d76de/tokenizers-0.21.1-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:ed248ab5279e601a30a4d67bdb897ecbe955a50f1e7bb62bd99f07dd11c2f5b6", size = 8937877, upload-time = "2025-03-13T10:51:12.688Z" },
{ url = "https://files.pythonhosted.org/packages/ac/33/ff08f50e6d615eb180a4a328c65907feb6ded0b8f990ec923969759dc379/tokenizers-0.21.1-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:9ac78b12e541d4ce67b4dfd970e44c060a2147b9b2a21f509566d556a509c67d", size = 9186645, upload-time = "2025-03-13T10:51:14.723Z" },
{ url = "https://files.pythonhosted.org/packages/5f/aa/8ae85f69a9f6012c6f8011c6f4aa1c96154c816e9eea2e1b758601157833/tokenizers-0.21.1-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e5a69c1a4496b81a5ee5d2c1f3f7fbdf95e90a0196101b0ee89ed9956b8a168f", size = 9384380, upload-time = "2025-03-13T10:51:16.526Z" },
{ url = "https://files.pythonhosted.org/packages/e8/5b/a5d98c89f747455e8b7a9504910c865d5e51da55e825a7ae641fb5ff0a58/tokenizers-0.21.1-cp39-abi3-win32.whl", hash = "sha256:1039a3a5734944e09de1d48761ade94e00d0fa760c0e0551151d4dd851ba63e3", size = 2239506, upload-time = "2025-03-13T10:51:20.643Z" },
{ url = "https://files.pythonhosted.org/packages/e6/b6/072a8e053ae600dcc2ac0da81a23548e3b523301a442a6ca900e92ac35be/tokenizers-0.21.1-cp39-abi3-win_amd64.whl", hash = "sha256:0f0dcbcc9f6e13e675a66d7a5f2f225a736745ce484c1a4e07476a89ccdad382", size = 2435481, upload-time = "2025-03-13T10:51:19.243Z" },
]
[[package]] [[package]]
name = "tqdm" name = "tqdm"
version = "4.67.1" version = "4.67.1"
@ -1486,6 +1540,25 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/f3/40/b1c265d4b2b62b58576588510fc4d1fe60a86319c8de99fd8e9fec617d2c/virtualenv-20.31.2-py3-none-any.whl", hash = "sha256:36efd0d9650ee985f0cad72065001e66d49a6f24eb44d98980f630686243cf11", size = 6057982, upload-time = "2025-05-08T17:58:21.15Z" }, { url = "https://files.pythonhosted.org/packages/f3/40/b1c265d4b2b62b58576588510fc4d1fe60a86319c8de99fd8e9fec617d2c/virtualenv-20.31.2-py3-none-any.whl", hash = "sha256:36efd0d9650ee985f0cad72065001e66d49a6f24eb44d98980f630686243cf11", size = 6057982, upload-time = "2025-05-08T17:58:21.15Z" },
] ]
[[package]]
name = "voyageai"
version = "0.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohttp" },
{ name = "aiolimiter" },
{ name = "numpy" },
{ name = "pillow" },
{ name = "pydantic" },
{ name = "requests" },
{ name = "tenacity" },
{ name = "tokenizers" },
]
sdist = { url = "https://files.pythonhosted.org/packages/c7/b7/c9f633149f1bdea95b43d38aa50404df2bdf769f0ccc0b402ca922d454e3/voyageai-0.3.2.tar.gz", hash = "sha256:bd1b52d26179d91853cbd2a0e52dc95cb0d526760c6c830959e01eb5ff9eaa12", size = 18979, upload-time = "2024-12-03T00:33:53.471Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/86/e1/0b2defa3a83aabe67db05d5f494d617dd4764b2a043d83ddc26be5e6e0db/voyageai-0.3.2-py3-none-any.whl", hash = "sha256:1398d6c6bfb1dd3b484f400713e538f00ce8a335250442b0902c21116d9705a8", size = 25518, upload-time = "2024-12-03T00:33:51.927Z" },
]
[[package]] [[package]]
name = "watchfiles" name = "watchfiles"
version = "1.1.0" version = "1.1.0"