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/
.pytest_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.
## Features
- **Document Management**: Store and manage documents with automatic content parsing
- **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
- **Local SQLite**: No need to run additional servers
- **Support for various embedding providers**: You can use Ollama, VoyageAI, OpenAI or add your own
- **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
@ -18,14 +16,28 @@ A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient d
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
# Install dependencies
uv sync
EMBEDDING_PROVIDER="ollama"
EMBEDDING_MODEL="mxbai-embed-large" # or any other model
EMBEDDING_VECTOR_DIM=1024
```
# Activate virtual environment
source .venv/bin/activate
For VoyageAI:
```bash
EMBEDDING_PROVIDER="voyageai"
EMBEDDING_MODEL="voyage-3.5" # or any other model
EMBEDDING_VECTOR_DIM=1024
```
## Quick Start
@ -90,14 +102,11 @@ finally:
```python
async with HaikuRAG("database.db") as client:
# Basic search
results = await client.search("your query here")
# Search with custom parameters
results = await client.search(
query="machine learning",
limit=10, # Maximum results to return
k=60 # RRF parameter for reciprocal rank fusion
limit=5, # Maximum results to return, defaults to 5
k=60 # RRF parameter for reciprocal rank fusion, defaults to 60
)
# Process results
@ -107,29 +116,10 @@ async with HaikuRAG("database.db") as client:
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
The system supports 40+ file formats through MarkItDown:
`haiku.rag` supports 40+ file formats through MarkItDown:
- **Documents**: PDF, DOCX, PPTX, XLSX
- **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
- **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
1. Fork the repository

View file

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

View file

@ -1,6 +1,7 @@
from haiku.rag.config import Config
from haiku.rag.embeddings.base import EmbedderBase
from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder
from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
def get_embedder() -> EmbedderBase:
@ -10,4 +11,15 @@ def get_embedder() -> EmbedderBase:
if Config.EMBEDDING_PROVIDER == "ollama":
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}")

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" },
]
[[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]]
name = "aiosignal"
version = "1.3.2"
@ -470,6 +479,14 @@ dependencies = [
{ name = "watchfiles" },
]
[package.optional-dependencies]
all = [
{ name = "voyageai" },
]
voyageai = [
{ name = "voyageai" },
]
[package.dev-dependencies]
dev = [
{ name = "datasets" },
@ -490,8 +507,11 @@ requires-dist = [
{ name = "python-dotenv", specifier = ">=1.1.0" },
{ name = "sqlite-vec", specifier = ">=0.1.6" },
{ 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" },
]
provides-extras = ["voyageai", "all"]
[package.metadata.requires-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" },
]
[[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]]
name = "tiktoken"
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" },
]
[[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]]
name = "tqdm"
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" },
]
[[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]]
name = "watchfiles"
version = "1.1.0"