From a1079fd1f52340d8c70f1b793a9b95b001e1a32d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 18 Jun 2025 09:42:35 +0200 Subject: [PATCH] VoyageAI embeddings --- .gitignore | 3 ++ README.md | 78 ++++++++++------------------ pyproject.toml | 3 ++ src/haiku/rag/embeddings/__init__.py | 12 +++++ src/haiku/rag/embeddings/voyageai.py | 14 +++++ uv.lock | 73 ++++++++++++++++++++++++++ 6 files changed, 132 insertions(+), 51 deletions(-) create mode 100644 src/haiku/rag/embeddings/voyageai.py diff --git a/.gitignore b/.gitignore index a51b694a..c38836d3 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,6 @@ wheels/ tests/data/ .pytest_cache/ .ruff_cache/ + +# environment variables +.env diff --git a/README.md b/README.md index 34dfd829..a3ce1e53 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 042ec815..971629e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/haiku/rag/embeddings/__init__.py b/src/haiku/rag/embeddings/__init__.py index bcaee4d6..f54a1978 100644 --- a/src/haiku/rag/embeddings/__init__.py +++ b/src/haiku/rag/embeddings/__init__.py @@ -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}") diff --git a/src/haiku/rag/embeddings/voyageai.py b/src/haiku/rag/embeddings/voyageai.py new file mode 100644 index 00000000..188ec4a6 --- /dev/null +++ b/src/haiku/rag/embeddings/voyageai.py @@ -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] diff --git a/uv.lock b/uv.lock index a2f87e12..8dcb2c20 100644 --- a/uv.lock +++ b/uv.lock @@ -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"