From ecdf250e2181c6c97b9c7a5db45102215e72be2a Mon Sep 17 00:00:00 2001 From: doubleinfinity <6169958+doubleinfinity@users.noreply.github.com> Date: Thu, 30 Oct 2025 14:34:13 +1100 Subject: [PATCH 1/6] feat: add ZeusDB vector store integration --- .../vector_stores/ZeusDBIndexDemo.ipynb | 759 +++++++++++ .../community/integrations/vector_stores.md | 44 + .../module_guides/storing/vector_stores.md | 2 + .../.gitignore | 153 +++ .../llama-index-vector-stores-zeusdb/BUILD | 1 + .../CHANGELOG.md | 100 ++ .../llama-index-vector-stores-zeusdb/LICENSE | 21 + .../llama-index-vector-stores-zeusdb/Makefile | 17 + .../README.md | 174 +++ .../examples/async_examples.py | 456 +++++++ .../examples/delete_records_examples.py | 119 ++ .../examples/metadata_filter_examples.py | 184 +++ .../examples/mmr_examples.py | 291 ++++ .../examples/persistence_examples.py | 400 ++++++ .../examples/quantization_examples.py | 369 ++++++ .../examples/quickstart.py | 38 + .../vector_stores/zeusdb/__init__.py | 5 + .../llama_index/vector_stores/zeusdb/base.py | 1177 +++++++++++++++++ .../llama_index/vector_stores/zeusdb/py.typed | 0 .../pyproject.toml | 127 ++ .../tests/__init__.py | 0 .../tests/test_vector_stores_zeusdb.py | 408 ++++++ 22 files changed, 4845 insertions(+) create mode 100644 docs/examples/vector_stores/ZeusDBIndexDemo.ipynb create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/.gitignore create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/BUILD create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/LICENSE create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/Makefile create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/delete_records_examples.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/metadata_filter_examples.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/mmr_examples.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/persistence_examples.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quantization_examples.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quickstart.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/__init__.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/py.typed create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/__init__.py create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py diff --git a/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb b/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb new file mode 100644 index 0000000000..8270d2aa2d --- /dev/null +++ b/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb @@ -0,0 +1,759 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1c177cbd-3774-4317-b9ae-1afedcf00b17", + "metadata": {}, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "id": "0c760fca-af3a-4116-b75d-27b87656a9e9", + "metadata": {}, + "source": [ + "# ZeusDB Vector Store" + ] + }, + { + "cell_type": "markdown", + "id": "b4dcf44a-6aa5-4e8b-8dba-d096f0d74209", + "metadata": {}, + "source": [ + "This document explains how to use ZeusDB as a vector store in LlamaIndex. " + ] + }, + { + "cell_type": "markdown", + "id": "e512769f-6cb2-458a-a673-9cfe5d0f5d9f", + "metadata": {}, + "source": [ + "[ZeusDB](https://www.zeusdb.com) is a high-performance vector database written in Rust, offering features like product quantization, persistent storage, and enterprise-grade logging. " + ] + }, + { + "cell_type": "markdown", + "id": "14c4e2d7-80b6-4b7e-9ba3-89e60d2ea528", + "metadata": {}, + "source": [ + "Follow these instructions and examples below to enhance your LlamaIndex apps with ZeusDB's production capabilities." + ] + }, + { + "cell_type": "markdown", + "id": "59ed4632-0ae2-45cd-bbb7-87034f47396f", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "31ce9d72-00da-4738-b18b-b752a14f1586", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "markdown", + "id": "abb5ebe6-6b72-468a-b6e1-54f342518a06", + "metadata": {}, + "source": [ + "Install the ZeusDB LlamaIndex integration package from PyPi:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d694a171-023f-4873-a406-e26bce46df20", + "metadata": {}, + "outputs": [], + "source": [ + "pip install llama-index-vector-stores-zeusdb" + ] + }, + { + "cell_type": "markdown", + "id": "5863b8c9-d723-4093-bd19-7db3ffd94fde", + "metadata": {}, + "source": [ + "*Setup in Jupyter Notebooks*" + ] + }, + { + "cell_type": "markdown", + "id": "52f320c4-f9ea-45a9-9267-da84b0409c58", + "metadata": {}, + "source": [ + "> πŸ’‘ Tip: If you’re working inside Jupyter or Google Colab, use the %pip magic command so the package is installed into the active kernel:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1e101a7a-eb30-459f-9252-9de93a54ee31", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install llama-index-vector-stores-zeusdb" + ] + }, + { + "cell_type": "markdown", + "id": "82d894de-fb54-4a51-aa3c-6278fe364ec0", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "38bbd268-108c-49c2-a253-8b4644480ea9", + "metadata": {}, + "source": [ + "## Getting Started" + ] + }, + { + "cell_type": "markdown", + "id": "d717cfc5-9d0a-4eec-9768-9677d82893dd", + "metadata": {}, + "source": [ + "This example uses OpenAIEmbeddings, which requires an OpenAI API key – [Get your OpenAI API key here](https://platform.openai.com/api-keys)" + ] + }, + { + "cell_type": "markdown", + "id": "a83b4dd9-0f80-4251-919c-e6d2fe258abd", + "metadata": {}, + "source": [ + "Install the LlamaIndex Core and OpenAI integration packages from PyPi:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1cbd590d-58ef-49da-a5ff-b9673ff078fd", + "metadata": {}, + "outputs": [], + "source": [ + "pip install llama-index-core\n", + "pip install llama-index-llms-openai\n", + "pip install llama-index-embeddings-openai\n", + "\n", + "# Use these commands if inside Jupyter Notebooks\n", + "# %pip install llama-index-core\n", + "# %pip install llama-index-llms-openai\n", + "# %pip install llama-index-embeddings-openai" + ] + }, + { + "cell_type": "markdown", + "id": "3a9f61f4-742c-40b4-93ed-bf61cb1e6d85", + "metadata": {}, + "source": [ + "#### Please choose an option below for your OpenAI key integration" + ] + }, + { + "cell_type": "markdown", + "id": "fbc0d74c-dc82-4b48-82f6-179618ecfbdb", + "metadata": {}, + "source": [ + "*Option 1: πŸ”‘ Enter your API key each time* " + ] + }, + { + "cell_type": "markdown", + "id": "44bb6169-e34c-4c25-b607-c6e29ec60b60", + "metadata": {}, + "source": [ + "Use getpass in Jupyter to securely input your key for the current session:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e2dd1f59-7308-4ba5-96f3-84f99c263c5f", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import getpass\n", + "\n", + "os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')" + ] + }, + { + "cell_type": "markdown", + "id": "5422513a-4f60-4de7-ade2-3d1baff8f8e5", + "metadata": {}, + "source": [ + "*Option 2: πŸ—‚οΈ Use a .env file*" + ] + }, + { + "cell_type": "markdown", + "id": "787eb08c-703c-4fc2-b7f1-26d66bc32829", + "metadata": {}, + "source": [ + "Keep your key in a local .env file and load it automatically with python-dotenv" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b696af7-77ef-49c5-8198-8b077f550243", + "metadata": {}, + "outputs": [], + "source": [ + "from dotenv import load_dotenv\n", + "\n", + "load_dotenv() # reads .env and sets OPENAI_API_KEY" + ] + }, + { + "cell_type": "markdown", + "id": "4b97b7a5-7d1b-4690-ba5b-b506976638b0", + "metadata": {}, + "source": [ + "πŸŽ‰πŸŽ‰ That's it! You are good to go." + ] + }, + { + "cell_type": "markdown", + "id": "a206d1c9-5b75-4386-b69f-03a6f21afb36", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "cd0201e1-3863-4ceb-a338-0a535bf2d770", + "metadata": {}, + "source": [ + "## Initialization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14584c6d-d5ee-4fc6-8426-7cbb48fe6d5d", + "metadata": {}, + "outputs": [], + "source": [ + "# Import required Packages and Classes\n", + "from llama_index.core import VectorStoreIndex, Document, StorageContext\n", + "from llama_index.vector_stores.zeusdb import ZeusDBVectorStore\n", + "from llama_index.embeddings.openai import OpenAIEmbedding\n", + "from llama_index.llms.openai import OpenAI\n", + "from llama_index.core import Settings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3798bb8-0f85-4fb0-b4d5-c0b55695e338", + "metadata": {}, + "outputs": [], + "source": [ + "# Set up embedding model and LLM\n", + "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")\n", + "Settings.llm = OpenAI(model=\"gpt-5\")" + ] + }, + { + "cell_type": "markdown", + "id": "6097ef9a-4ab0-43f5-b7e1-61df1821018d", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "5ca8ee11-5adb-43e8-8e20-ced5b91d9f8c", + "metadata": {}, + "source": [ + "## Quickstart" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42960e47-48ba-41ca-a1fe-03ba695d1714", + "metadata": {}, + "outputs": [], + "source": [ + "# Create ZeusDB vector store\n", + "vector_store = ZeusDBVectorStore(\n", + " dim=1536, # OpenAI embedding dimension\n", + " distance=\"cosine\",\n", + " index_type=\"hnsw\"\n", + ")\n", + "\n", + "# Create storage context\n", + "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n", + "\n", + "# Create documents\n", + "documents = [\n", + " Document(text=\"ZeusDB is a high-performance vector database.\"),\n", + " Document(text=\"LlamaIndex provides RAG capabilities.\"),\n", + " Document(text=\"Vector search enables semantic similarity.\")\n", + "]\n", + "\n", + "# Create index and store documents\n", + "index = VectorStoreIndex.from_documents(\n", + " documents,\n", + " storage_context=storage_context\n", + ")\n", + "\n", + "# Query the index\n", + "query_engine = index.as_query_engine()\n", + "response = query_engine.query(\"What is ZeusDB?\")\n", + "print(response)" + ] + }, + { + "cell_type": "markdown", + "id": "34a53714-e531-4cfe-95c5-a957bb33a1c2", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "373dd643-b8c0-49ff-9d27-41e99a6ca3a0", + "metadata": {}, + "source": [ + "## Direct Query Interface Example" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f7509be-1953-442a-9f02-3114f610777a", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.vector_stores.types import VectorStoreQuery\n", + "\n", + "# Create query\n", + "embed_model = Settings.embed_model\n", + "query_embedding = embed_model.get_text_embedding(\"machine learning\")\n", + "\n", + "query_obj = VectorStoreQuery(\n", + " query_embedding=query_embedding,\n", + " similarity_top_k=2\n", + ")\n", + "\n", + "# Execute query\n", + "results = vector_store.query(query_obj)\n", + "\n", + "# Results contain IDs and similarities\n", + "print(f\"Found {len(results.ids or [])} results:\")\n", + "for node_id, similarity in zip(results.ids or [], results.similarities or []):\n", + " print(f\" ID: {node_id}, Similarity: {similarity:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "053344dd-871e-4e1c-b851-31548e46fc80", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "4b6e17e7-c3f8-476d-9d4d-8207d1e1dd82", + "metadata": {}, + "source": [ + "## MMR Search" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8ce77669-8ee6-4d65-9307-92d7ca74c008", + "metadata": {}, + "outputs": [], + "source": [ + "# MMR search via direct query\n", + "mmr_results = vector_store.query(\n", + " query_obj,\n", + " mmr=True,\n", + " fetch_k=10,\n", + " mmr_lambda=0.7 # 0.0=max diversity, 1.0=pure relevance\n", + ")\n", + "\n", + "print(f\"MMR Results: {len(mmr_results.ids or [])} items (with diversity)\")" + ] + }, + { + "cell_type": "markdown", + "id": "9340fa93-2018-4423-859a-d9a6a5096f21", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "f71a5be3-b7e9-4cf8-bd93-30cfce02ed2a", + "metadata": {}, + "source": [ + "## Search with Metadata Filtering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "696bad5f-1c32-4137-8349-65b89f8bd9c3", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.vector_stores.types import (\n", + " MetadataFilters,\n", + " FilterOperator,\n", + " FilterCondition\n", + ")\n", + "\n", + "# Create a fresh vector store for this example\n", + "vector_store = ZeusDBVectorStore(dim=1536, distance=\"cosine\")\n", + "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n", + "\n", + "# Create documents with metadata\n", + "documents_with_meta = [\n", + " Document(\n", + " text=\"Python is great for data science\",\n", + " metadata={\"category\": \"tech\", \"year\": 2024}\n", + " ),\n", + " Document(\n", + " text=\"JavaScript is for web development\",\n", + " metadata={\"category\": \"tech\", \"year\": 2023}\n", + " ),\n", + " Document(\n", + " text=\"Climate change impacts ecosystems\",\n", + " metadata={\"category\": \"environment\", \"year\": 2024}\n", + " ),\n", + "]\n", + "\n", + "# Build index with metadata\n", + "index = VectorStoreIndex.from_documents(\n", + " documents_with_meta,\n", + " storage_context=storage_context\n", + ")\n", + "\n", + "# Create metadata filter\n", + "filters = MetadataFilters.from_dicts([\n", + " {\"key\": \"category\", \"value\": \"tech\", \"operator\": FilterOperator.EQ},\n", + " {\"key\": \"year\", \"value\": 2024, \"operator\": FilterOperator.GTE}\n", + "], condition=FilterCondition.AND)\n", + "\n", + "# Use the retriever with filters (recommended approach)\n", + "retriever = index.as_retriever(similarity_top_k=5, filters=filters)\n", + "filtered_results = retriever.retrieve(\"programming\")\n", + "\n", + "# Process results\n", + "for r in filtered_results:\n", + " print(f\"- {r.node.get_content(metadata_mode='none')}\")\n", + " print(f\" Metadata: {r.node.metadata}\\n\")" + ] + }, + { + "cell_type": "markdown", + "id": "f288b835-6150-4a07-bb9d-891c073ec377", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "556fb9ef-ec59-496e-a054-7424ea4b36e7", + "metadata": {}, + "source": [ + "## Save and Load indexes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50747f8a-e010-412b-9c0e-e118503488fd", + "metadata": {}, + "outputs": [], + "source": [ + "# Save index\n", + "save_path = \"my_index.zdb\"\n", + "vector_store.save_index(save_path)\n", + "print(f\"βœ… Index saved to {save_path}\")\n", + "print(f\" Vector count: {vector_store.get_vector_count()}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13ea3708-b096-4597-a394-dddb15e144ad", + "metadata": {}, + "outputs": [], + "source": [ + "# Load index\n", + "loaded_store = ZeusDBVectorStore.load_index(save_path)\n", + "print(f\"βœ… Index loaded from {save_path}\")\n", + "print(f\" Vector count: {loaded_store.get_vector_count()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "92122449-26ee-481c-9ead-70c82316fcf5", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "4665eb68-7d5e-40f0-8ab3-25b62ff97b18", + "metadata": {}, + "source": [ + "## Quantization Example" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "80b7f87f-c399-4efb-b109-6993ee0aadf7", + "metadata": {}, + "outputs": [], + "source": [ + "# Create quantized vector store for memory efficiency\n", + "quantization_config = {\n", + " 'type': 'pq',\n", + " 'subvectors': 8,\n", + " 'bits': 8,\n", + " 'training_size': 1000,\n", + " 'storage_mode': 'quantized_only'\n", + "}\n", + "\n", + "vector_store = ZeusDBVectorStore(\n", + " dim=1536,\n", + " distance=\"cosine\",\n", + " index_type=\"hnsw\",\n", + " quantization_config=quantization_config\n", + ")\n", + "\n", + "# Check quantization status\n", + "print(f\"Is quantized: {vector_store.is_quantized()}\")\n", + "print(f\"Can use quantization: {vector_store.can_use_quantization()}\")\n", + "print(f\"Training progress: {vector_store.get_training_progress():.1f}%\")\n", + "print(f\"Storage mode: {vector_store.get_storage_mode()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "3f8b0a1c-5bf1-458f-af50-5accb0950033", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "fdb4f859-2b2d-481a-ae34-3fc551d6824b", + "metadata": {}, + "source": [ + "## Delete Operations Example" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "63d467cd-7ea7-4afe-8f61-826c9f2b1299", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core import VectorStoreIndex, Document, StorageContext\n", + "from llama_index.vector_stores.zeusdb import ZeusDBVectorStore\n", + "\n", + "# Create a fresh vector store for this example\n", + "delete_vs = ZeusDBVectorStore(dim=1536, distance=\"cosine\")\n", + "delete_sc = StorageContext.from_defaults(vector_store=delete_vs)\n", + "\n", + "# Create documents\n", + "delete_docs = [\n", + " Document(text=f\"Document {i}\", metadata={\"doc_id\": i})\n", + " for i in range(5)\n", + "]\n", + "\n", + "# Build index\n", + "delete_index = VectorStoreIndex.from_documents(\n", + " delete_docs,\n", + " storage_context=delete_sc\n", + ")\n", + "\n", + "print(f\"Before delete: {delete_vs.get_vector_count()} vectors\")\n", + "\n", + "# Get node IDs to delete\n", + "retriever = delete_index.as_retriever(similarity_top_k=10)\n", + "results = retriever.retrieve(\"document\")\n", + "\n", + "if results:\n", + " # Extract node IDs from results\n", + " node_ids_to_delete = [result.node.node_id for result in results[:2]]\n", + " print(f\"Deleting node IDs: {node_ids_to_delete[0][:8]}...\")\n", + " \n", + " # Delete by node IDs\n", + " delete_vs.delete_nodes(node_ids=node_ids_to_delete)\n", + " print(f\"After delete: {delete_vs.get_vector_count()} vectors\")\n", + " print(\"βœ… delete_nodes(node_ids=[...]) works!\")\n", + "\n", + "# Demonstrate unsupported delete by ref_doc_id\n", + "try:\n", + " delete_vs.delete(ref_doc_id=\"doc_1\")\n", + " print(\"❌ Should have raised NotImplementedError\")\n", + "except NotImplementedError as e:\n", + " print(\"❌ delete(ref_doc_id='...') raises NotImplementedError\")\n", + " print(f\" (This is expected - not supported by backend)\")" + ] + }, + { + "cell_type": "markdown", + "id": "6a455833-7fd5-4fe5-9d1c-d7c20fab2441", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "818e9a2b-590a-46a8-8f97-bf10f35235ee", + "metadata": {}, + "source": [ + "## Async Operations Example" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b51f3dde-3247-4480-9e05-3465614ed82d", + "metadata": {}, + "outputs": [], + "source": [ + "import asyncio\n", + "from llama_index.core.schema import TextNode\n", + "\n", + "# In Jupyter, use nest_asyncio to handle event loops\n", + "try:\n", + " import nest_asyncio\n", + " nest_asyncio.apply()\n", + "except ImportError:\n", + " pass\n", + "\n", + "async def async_operations():\n", + " # Create nodes\n", + " nodes = [\n", + " TextNode(text=f\"Document {i}\", metadata={\"doc_id\": i})\n", + " for i in range(10)\n", + " ]\n", + " \n", + " # Generate embeddings (required before adding)\n", + " embed_model = Settings.embed_model\n", + " for node in nodes:\n", + " node.embedding = embed_model.get_text_embedding(node.text)\n", + " \n", + " # Add nodes asynchronously\n", + " node_ids = await vector_store.async_add(nodes)\n", + " print(f\"Added {len(node_ids)} nodes\")\n", + " \n", + " # Query asynchronously\n", + " query_embedding = embed_model.get_text_embedding(\"document\")\n", + " query_obj = VectorStoreQuery(\n", + " query_embedding=query_embedding,\n", + " similarity_top_k=3\n", + " )\n", + " \n", + " results = await vector_store.aquery(query_obj)\n", + " print(f\"Found {len(results.ids or [])} results\")\n", + " \n", + " # Delete asynchronously\n", + " await vector_store.adelete_nodes(node_ids=node_ids[:2])\n", + " print(f\"Deleted 2 nodes, {vector_store.get_vector_count()} remaining\")\n", + "\n", + "# Run async function\n", + "await async_operations() # In Jupyter\n", + "# asyncio.run(async_operations()) # In regular Python scripts" + ] + }, + { + "cell_type": "markdown", + "id": "12926879-24b7-46cb-9c0c-f5e46e36a0c3", + "metadata": {}, + "source": [ + "---" + ] + }, + { + "cell_type": "markdown", + "id": "747d9f49-dac7-4d7d-9519-c6a5f68cb85a", + "metadata": {}, + "source": [ + "## Performance Monitoring" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08ed7b40-1f87-4307-909c-80ddb1e082f7", + "metadata": {}, + "outputs": [], + "source": [ + "# Get index statistics\n", + "stats = vector_store.get_zeusdb_stats()\n", + "print(f\"Key stats: vectors={stats.get('total_vectors')}, space={stats.get('space')}\")\n", + "\n", + "# Get vector count\n", + "count = vector_store.get_vector_count()\n", + "print(f\"Vector count: {count}\")\n", + "\n", + "# Get detailed index info\n", + "info = vector_store.info()\n", + "print(f\"Index info: {info}\")\n", + "\n", + "# Check quantization status\n", + "if vector_store.is_quantized():\n", + " progress = vector_store.get_training_progress()\n", + " quant_info = vector_store.get_quantization_info()\n", + " print(f\"Quantization: {progress:.1f}% complete\")\n", + " print(f\"Compression: {quant_info['compression_ratio']:.1f}x\")\n", + "else:\n", + " print(\"Index is not quantized\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/src/content/docs/framework/community/integrations/vector_stores.md b/docs/src/content/docs/framework/community/integrations/vector_stores.md index 2cba3b1166..05602f3fe4 100644 --- a/docs/src/content/docs/framework/community/integrations/vector_stores.md +++ b/docs/src/content/docs/framework/community/integrations/vector_stores.md @@ -56,6 +56,7 @@ as the storage backend for `VectorStoreIndex`. - Weaviate (`WeaviateVectorStore`). [Installation](https://weaviate.io/developers/weaviate/installation). [Python Client](https://weaviate.io/developers/weaviate/client-libraries/python). - WordLift (`WordliftVectorStore`). [Quickstart](https://docs.wordlift.io/llm-connectors/wordlift-vector-store/). [Python Client](https://pypi.org/project/wordlift-client/). - Zep (`ZepVectorStore`). [Installation](https://docs.getzep.com/deployment/quickstart/). [Python Client](https://docs.getzep.com/sdk/). +- ZeusDB (`ZeusDBVectorStore`). [Installation/Quickstart](https://docs.zeusdb.com/en/latest/vector_database/getting_started.html) - Zilliz (`MilvusVectorStore`). [Quickstart](https://zilliz.com/doc/quick_start) A detailed API reference is [found here](/python/framework-api-reference/storage/vector_store). @@ -958,6 +959,48 @@ retriever = index.as_retriever(filters=filters) result = retriever.retrieve("What is inception about?") ``` +**ZeusDB** + +```python +from llama_index.core import VectorStoreIndex, Document, StorageContext +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.core import Settings + +# Set up embedding model and LLM +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-5") + +# Create ZeusDB vector store +vector_store = ZeusDBVectorStore( + dim=1536, # OpenAI embedding dimension + distance="cosine", + index_type="hnsw" +) + +# Create storage context +storage_context = StorageContext.from_defaults(vector_store=vector_store) + +# Create documents +documents = [ + Document(text="ZeusDB is a high-performance vector database."), + Document(text="LlamaIndex provides RAG capabilities."), + Document(text="Vector search enables semantic similarity.") +] + +# Create index and store documents +index = VectorStoreIndex.from_documents( + documents, + storage_context=storage_context +) + +# Query the index +query_engine = index.as_query_engine() +response = query_engine.query("What is ZeusDB?") +print(response) +``` + **Zilliz** - Zilliz Cloud (hosted version of Milvus) uses the Milvus Index with some extra arguments. @@ -1200,3 +1243,4 @@ documents = reader.load_data( - [Weaviate Hybrid Search](/python/examples/vector_stores/weaviateindexdemo-hybrid) - [WordLift](/python/examples/vector_stores/wordliftdemo) - [Zep](/python/examples/vector_stores/zepindexdemo) +- [ZeusDB](/python/examples/vector_stores/zeusdbindexdemo) diff --git a/docs/src/content/docs/framework/module_guides/storing/vector_stores.md b/docs/src/content/docs/framework/module_guides/storing/vector_stores.md index 5c2a3fc1b8..3ae70e74b8 100644 --- a/docs/src/content/docs/framework/module_guides/storing/vector_stores.md +++ b/docs/src/content/docs/framework/module_guides/storing/vector_stores.md @@ -68,6 +68,7 @@ We are actively adding more integrations and improving feature coverage for each | Vertex AI Vector Search | cloud | βœ“ | | βœ“ | βœ“ | | | Weaviate | self-hosted / cloud | βœ“ | βœ“ | βœ“ | βœ“ | | | WordLift | cloud | βœ“ | βœ“ | βœ“ | βœ“ | βœ“ | +| ZeusDB | self-hosted / cloud | βœ“ | βœ“ | βœ“ | βœ“ | βœ“ | For more details, see [Vector Store Integrations](/python/framework/community/integrations/vector_stores). @@ -130,3 +131,4 @@ For more details, see [Vector Store Integrations](/python/framework/community/in - [Weaviate Hybrid Search](/python/examples/vector_stores/weaviateindexdemo-hybrid) - [WordLift](/python/examples/vector_stores/wordliftdemo) - [Zep](/python/examples/vector_stores/zepindexdemo) +- [ZeusDB](/python/examples/vector_stores/zeusdbindexdemo) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/.gitignore b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/.gitignore new file mode 100644 index 0000000000..990c18de22 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/.gitignore @@ -0,0 +1,153 @@ +llama_index/_static +.DS_Store +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +bin/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +etc/ +include/ +lib/ +lib64/ +parts/ +sdist/ +share/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +.ruff_cache + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints +notebooks/ + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +pyvenv.cfg + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# Jetbrains +.idea +modules/ +*.swp + +# VsCode +.vscode + +# pipenv +Pipfile +Pipfile.lock + +# pyright +pyrightconfig.json diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/BUILD new file mode 100644 index 0000000000..db46e8d6c9 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md new file mode 100644 index 0000000000..000063c69b --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md @@ -0,0 +1,100 @@ + +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +--- + +## [0.1.3] - 2025-10-28 + +### Added + +- Added `.venv` to the `mypy` exclude list to improve compatibility with local development environments. + +### Changed + +- Bumped `llama-index-core` to `>=0.14.6` to match the latest LlamaIndex release (Oct 26, 2025). +- Updated `llama-index-llms-openai` development dependency to `>=0.6.6`. +- Updated `pyproject.toml` build includes for proper source distribution packaging (relative paths under `[tool.hatch.build.targets.sdist]`). + +--- + +## [0.1.2] - 2025-10-16 + +### Changed + +- **Async API naming**: Renamed `aadd()` to `async_add()` to align with LlamaIndex standard async method naming conventions. The adapter now uses the official LlamaIndex async API (`async_add`, `aquery`, `adelete`, `adelete_nodes`, `aclear`) for consistency across the ecosystem. +- Updated async examples to use standard LlamaIndex async method names instead of custom aliases. +- Updated documentation strings to reference `async_add()` instead of `aadd()`. + +--- + +## [0.1.1] - 2025-10-14 + +### Added + +- Delete examples - Added `examples/delete_examples.py` demonstrating supported deletion operations and workarounds. +- Test coverage - Added tests for ID-based deletion and proper error handling for unsupported operations. +- Addeed Product **Quantization (PQ) support** - Full support for memory-efficient vector compression with automatic training +- **Persistence Support**: Complete save/load functionality for ZeusDB indexes + - Save indexes to disk with `save_index(path)` + - Load indexes from disk with `load_index(path)` + - Preserves vectors, metadata, HNSW graph structure, and quantization configuration + - Directory-based format (.zdb) with JSON metadata and binary data files + - Cross-platform compatibility for sharing indexes between systems + - Added comprehensive persistence examples (`examples/persistence_examples.py`) +- **Async Support**: Full asynchronous operation support for non-blocking workflows + - `aadd()` - Add nodes asynchronously + - `aquery()` - Query asynchronously + - `adelete_nodes()` - Delete nodes by IDs asynchronously + - Thread-offloaded async wrappers using `asyncio.to_thread()` +- **MMR (Maximal Marginal Relevance) Search**: Diversity-focused retrieval for comprehensive results + - Balance relevance and diversity with `mmr_lambda` parameter (0.0-1.0) + - Control candidate pool size with `fetch_k` parameter + - Prevents redundant/similar results in search responses + - Perfect for RAG applications, research, and multi-perspective retrieval +- Added comprehensive async examples (`examples/async_examples.py`) +- Added MMR examples (`examples/mmr_examples.py`) + +### Changed + +- Filter format alignment β€” Updated _filters_to_zeusdb() to produce a flat dictionary with implicit AND (e.g., { "key": value, "other": { "op": value } }) instead of a nested {"and": [...]} structure, matching the Rust implementation. +- Test infrastructure β€” Updated _match_filter() in the test fake from the nested format to the flat format to reflect production behavior. + +### Fixed + +- Filter translation for metadata queries - _filters_to_zeusdb() now emits the flat format expected by the ZeusDB backend. Single filters and AND combinations are handled correctly. The previous nested format could cause filtered queries to return zero results. +- Deletion behavior - Correctly implemented ID-based deletion using `remove_point()`. Delete operations now properly remove vectors from the index and update vector counts. + +--- + +## [0.1.0] - 2025-09-19 + +### Added + +- Initial release of ZeusDB vector database integration for LlamaIndex. +- Support for connecting LlamaIndex’s RAG framework with ZeusDB for high-performance retrieval. +- Trusted Publishing workflow for PyPI releases via GitHub Actions. +- Build validation workflow to check distributions without publishing. +- Documentation updates including project description and setup instructions. + +--- + +## [Unreleased] + +### Added + + +### Changed + + +### Fixed + + +### Removed + + +--- diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/LICENSE b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/LICENSE new file mode 100644 index 0000000000..0839457ac5 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 ZeusDB + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/Makefile b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/Makefile new file mode 100644 index 0000000000..b9eab05aa3 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/Makefile @@ -0,0 +1,17 @@ +GIT_ROOT ?= $(shell git rev-parse --show-toplevel) + +help: ## Show all Makefile targets. + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}' + +format: ## Run code autoformatters (black). + pre-commit install + git ls-files | xargs pre-commit run black --files + +lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy + pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files + +test: ## Run tests via pytest. + pytest tests + +watch-docs: ## Build and watch documentation. + sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/ diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md new file mode 100644 index 0000000000..8eb7977b4c --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md @@ -0,0 +1,174 @@ +# LlamaIndex ZeusDB Integration + +ZeusDB vector database integration for LlamaIndex. Connect LlamaIndex's RAG framework with high-performance, enterprise-grade vector database. + +## Features + +- **Production Ready**: Built for enterprise-scale RAG applications +- **Persistence**: Complete save/load functionality with cross-platform compatibility +- **Advanced Filtering**: Comprehensive metadata filtering with complex operators +- **MMR Support**: Maximal Marginal Relevance for diverse, non-redundant results +- **Quantization**: Product Quantization (PQ) for memory-efficient vector storage +- **Async Support**: Async wrappers for non-blocking operations (`aadd`, `aquery`, `adelete_nodes`) + +## Installation + +```bash +pip install llama-index-vector-stores-zeusdb +``` + +## Quick Start + +```python +from llama_index.core import VectorStoreIndex, Document, StorageContext +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.core import Settings + +# Set up embedding model and LLM +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-5") + +# Create ZeusDB vector store +vector_store = ZeusDBVectorStore( + dim=1536, # OpenAI embedding dimension + distance="cosine", + index_type="hnsw" +) + +# Create storage context +storage_context = StorageContext.from_defaults(vector_store=vector_store) + +# Create documents +documents = [ + Document(text="ZeusDB is a high-performance vector database."), + Document(text="LlamaIndex provides RAG capabilities."), + Document(text="Vector search enables semantic similarity.") +] + +# Create index and store documents +index = VectorStoreIndex.from_documents( + documents, + storage_context=storage_context +) + +# Query the index +query_engine = index.as_query_engine() +response = query_engine.query("What is ZeusDB?") +print(response) +``` + +## Advanced Features + +### Persistence + +Save and load indexes with complete state preservation: + +```python +# Save index to disk +vector_store.save_index("my_index.zdb") + +# Load index from disk +loaded_store = ZeusDBVectorStore.load_index("my_index.zdb") +``` + +### MMR Search + +Balance relevance and diversity for comprehensive results: + +```python +from llama_index.core.vector_stores.types import VectorStoreQuery + +# Query with MMR for diverse results +query_embedding = embed_model.get_text_embedding("your query") +results = vector_store.query( + VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=5), + mmr=True, + fetch_k=20, + mmr_lambda=0.7 # 0.0=max diversity, 1.0=pure relevance +) + +# Note: MMR automatically enables return_vector=True for diversity calculation +# Results contain ids and similarities (nodes=None) +``` + +### Quantization + +Reduce memory usage with Product Quantization: + +```python +vector_store = ZeusDBVectorStore( + dim=1536, + distance="cosine", + quantization_config={ + 'type': 'pq', + 'subvectors': 8, + 'bits': 8, + 'training_size': 1000, + 'storage_mode': 'quantized_only' + } +) +``` + +### Async Operations + +Non-blocking operations for web servers and concurrent workflows: + +```python +import asyncio + +# Async add +node_ids = await vector_store.aadd(nodes) + +# Async query +results = await vector_store.aquery(query_obj) + +# Async delete +await vector_store.adelete_nodes(node_ids=["id1", "id2"]) +``` + +### Metadata Filtering + +Filter results by metadata: + +```python +from llama_index.core.vector_stores.types import ( + MetadataFilters, + FilterOperator, + FilterCondition +) + +# Create metadata filter +filters = MetadataFilters.from_dicts([ + {"key": "category", "value": "tech", "operator": FilterOperator.EQ}, + {"key": "year", "value": 2024, "operator": FilterOperator.GTE} +], condition=FilterCondition.AND) + +# Query with filters +results = vector_store.query( + VectorStoreQuery( + query_embedding=query_embedding, + similarity_top_k=5, + filters=filters + ) +) +``` + +**Supported operators**: EQ, NE, GT, GTE, LT, LTE, IN, NIN, ANY, ALL, CONTAINS, TEXT_MATCH, TEXT_MATCH_INSENSITIVE + +## Configuration + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `dim` | Vector dimension | Required | +| `distance` | Distance metric (`cosine`, `l2`, `l1`) | `cosine` | +| `index_type` | Index type (`hnsw`) | `hnsw` | +| `m` | HNSW connectivity parameter | 16 | +| `ef_construction` | HNSW build-time search depth | 200 | +| `expected_size` | Expected number of vectors | 10000 | +| `quantization_config` | PQ quantization settings | None | + +## License + +This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py new file mode 100644 index 0000000000..4cd8f0e2dd --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py @@ -0,0 +1,456 @@ +# examples/async_examples.py +""" +ZeusDB Async Examples for LlamaIndex + +Demonstrates asynchronous operations for non-blocking, concurrent vector operations. + +When to use async: +- Web servers (FastAPI/Starlette) handling multiple requests +- Agents/pipelines doing parallel searches +- Concurrent document processing +- Notebooks where you want non-blocking operations + +For simple scripts, sync methods are fine. +""" + +import asyncio +import time + +from dotenv import load_dotenv + +from llama_index.core import Settings +from llama_index.core.schema import TextNode +from llama_index.core.vector_stores.types import VectorStoreQuery +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() + +# Configure OpenAI +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") + +print("=" * 70) +print("ZeusDB Async Examples") +print("=" * 70) +print() + + +# ============================================================================= +# Example 1: Basic Async Operations +# ============================================================================= +async def example_basic_async(): + """Demonstrate basic async add, query, and delete operations.""" + print("=" * 70) + print("Example 1: Basic Async Operations") + print("=" * 70) + print() + + # Create vector store + vector_store = ZeusDBVectorStore(dim=1536, distance="cosine", index_type="hnsw") + + # Create sample nodes + nodes = [ + TextNode( + text=f"Document {i}: Sample content about technology.", + metadata={"doc_id": i, "category": "tech"}, + ) + for i in range(10) + ] + + # Generate embeddings + print("Generating embeddings...") + embed_model = Settings.embed_model + for node in nodes: + node.embedding = embed_model.get_text_embedding(node.text) + + # Async add + print(f"Adding {len(nodes)} nodes asynchronously...") + start = time.perf_counter() + # node_ids = await vector_store.aadd(nodes) + node_ids = await vector_store.async_add(nodes) + add_time = (time.perf_counter() - start) * 1000 + print(f" βœ… Added {len(node_ids)} nodes in {add_time:.2f}ms") + print() + + # Async query + print("Querying asynchronously...") + query_text = "technology" + query_embedding = embed_model.get_text_embedding(query_text) + query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=3) + + start = time.perf_counter() + results = await vector_store.aquery(query_obj) + query_time = (time.perf_counter() - start) * 1000 + + result_ids = results.ids or [] + result_sims = results.similarities or [] + + print(f" βœ… Query completed in {query_time:.2f}ms") + print(f" Found {len(result_ids)} results:") + for i, (node_id, similarity) in enumerate(zip(result_ids, result_sims), 1): + print(f" {i}. Score: {similarity:.4f}, ID: {node_id}") + print() + + # Async delete + print("Deleting nodes asynchronously...") + start = time.perf_counter() + await vector_store.adelete_nodes(node_ids=node_ids[:3]) + delete_time = (time.perf_counter() - start) * 1000 + print(f" βœ… Deleted 3 nodes in {delete_time:.2f}ms") + print(f" Remaining nodes: {vector_store.get_vector_count()}") + print() + + print(f"βœ… Example 1 completed. Final count: {vector_store.get_vector_count()}") + print() + + +# ============================================================================= +# Example 2: Concurrent Queries +# ============================================================================= +async def example_concurrent_queries(): + """Demonstrate running multiple queries concurrently.""" + print("=" * 70) + print("Example 2: Concurrent Queries") + print("=" * 70) + print() + + # Setup + vector_store = ZeusDBVectorStore(dim=1536, distance="cosine") + embed_model = Settings.embed_model + + # Create diverse documents + documents = [ + "Python is a programming language for data science and web development.", + "Machine learning models require training data and compute resources.", + "Climate change is affecting global weather patterns and ecosystems.", + "Quantum computing uses qubits for parallel computation.", + "Electric vehicles are becoming more popular due to environmental concerns.", + "Artificial intelligence can analyze large datasets efficiently.", + "Renewable energy sources include solar, wind, and hydroelectric power.", + "Neural networks are inspired by biological brain structures.", + ] + + nodes = [] + print("Preparing documents...") + for i, text in enumerate(documents): + node = TextNode(text=text, metadata={"doc_id": i, "source": "examples"}) + node.embedding = embed_model.get_text_embedding(text) + nodes.append(node) + + # await vector_store.aadd(nodes) + await vector_store.async_add(nodes) + print(f" βœ… Added {len(nodes)} documents") + print() + + # Define multiple queries + queries = [ + "programming languages", + "machine learning", + "climate and environment", + "quantum computing", + ] + + print(f"Running {len(queries)} queries concurrently...") + print() + + # Sequential version (for comparison) + print("πŸ“Š Sequential execution:") + start = time.perf_counter() + seq_results = [] + for query_text in queries: + query_embedding = embed_model.get_text_embedding(query_text) + query_obj = VectorStoreQuery( + query_embedding=query_embedding, similarity_top_k=2 + ) + result = await vector_store.aquery(query_obj) + seq_results.append(result) + seq_time = (time.perf_counter() - start) * 1000 + print(f" ⏱️ Total time: {seq_time:.2f}ms") + print() + + # Concurrent version + print("⚑ Concurrent execution:") + + async def query_single(query_text: str): + query_embedding = embed_model.get_text_embedding(query_text) + query_obj = VectorStoreQuery( + query_embedding=query_embedding, similarity_top_k=2 + ) + return await vector_store.aquery(query_obj) + + start = time.perf_counter() + concurrent_results = await asyncio.gather(*[query_single(q) for q in queries]) + concurrent_time = (time.perf_counter() - start) * 1000 + + print(f" ⏱️ Total time: {concurrent_time:.2f}ms") + speedup = seq_time / concurrent_time if concurrent_time > 0 else 0 + print(f" πŸš€ Speedup: {speedup:.2f}x faster") + print() + + # Display results + print("Results:") + for query_text, result in zip(queries, concurrent_results): + result_ids = result.ids or [] + result_sims = result.similarities or [] + print(f" Query: '{query_text}'") + for i, (node_id, sim) in enumerate(zip(result_ids, result_sims), 1): + print(f" {i}. Score: {sim:.4f}") + print() + + print(f"βœ… Example 2 completed. Total documents: {vector_store.get_vector_count()}") + print() + + +# ============================================================================= +# Example 3: Concurrent Document Processing +# ============================================================================= +async def example_concurrent_processing(): + """Demonstrate processing multiple document batches concurrently.""" + print("=" * 70) + print("Example 3: Concurrent Document Processing") + print("=" * 70) + print() + + vector_store = ZeusDBVectorStore(dim=1536, distance="cosine") + embed_model = Settings.embed_model + + # Simulate multiple batches of documents + batches = [ + [f"Tech document {i}" for i in range(5)], + [f"Science document {i}" for i in range(5)], + [f"Business document {i}" for i in range(5)], + ] + + print(f"Processing {len(batches)} batches concurrently...") + print() + + async def process_batch(batch: list[str], batch_id: int): + """Process a single batch of documents.""" + start = time.perf_counter() + + nodes = [] + for i, text in enumerate(batch): + node = TextNode( + text=text, + metadata={ + "batch_id": batch_id, + "doc_id": i, + }, + ) + node.embedding = embed_model.get_text_embedding(text) + nodes.append(node) + + # node_ids = await vector_store.aadd(nodes) + node_ids = await vector_store.async_add(nodes) + elapsed = (time.perf_counter() - start) * 1000 + + print(f" βœ… Batch {batch_id}: Added {len(node_ids)} nodes in {elapsed:.2f}ms") + return node_ids + + # Process all batches concurrently + start = time.perf_counter() + all_ids = await asyncio.gather( + *[process_batch(batch, i) for i, batch in enumerate(batches)] + ) + total_time = (time.perf_counter() - start) * 1000 + + total_docs = sum(len(ids) for ids in all_ids) + print() + print("πŸ“Š Summary:") + print(f" Total documents: {total_docs}") + print(f" Total time: {total_time:.2f}ms") + print(f" Avg per document: {total_time / total_docs:.2f}ms") + print() + + print(f"βœ… Example 3 completed. Total documents: {vector_store.get_vector_count()}") + print() + + +# ============================================================================= +# Example 4: Async with Error Handling +# ============================================================================= +async def example_error_handling(): + """Demonstrate proper error handling in async operations.""" + print("=" * 70) + print("Example 4: Async Error Handling") + print("=" * 70) + print() + + vector_store = ZeusDBVectorStore(dim=1536, distance="cosine") + embed_model = Settings.embed_model + + # Add some documents + nodes = [] + for i in range(5): + node = TextNode(text=f"Document {i}", metadata={"doc_id": i}) + node.embedding = embed_model.get_text_embedding(node.text) + nodes.append(node) + + # node_ids = await vector_store.aadd(nodes) + node_ids = await vector_store.async_add(nodes) + print(f"Added {len(node_ids)} documents") + print() + + # Demonstrate handling of unsupported operation + print("Testing unsupported delete by ref_doc_id...") + try: + await vector_store.adelete(ref_doc_id="some_doc") + print(" ❌ Should have raised NotImplementedError") + except NotImplementedError as e: + print(f" βœ… Correctly caught error: {str(e)[:60]}...") + print() + + # Test unsupported clear operation + print("Testing unsupported clear operation...") + try: + await vector_store.aclear() + print(" ❌ Should have raised NotImplementedError") + except NotImplementedError as e: + print(f" βœ… Correctly caught error: {str(e)[:50]}...") + print() + + # Demonstrate successful delete by node IDs + print("Testing supported delete by node IDs...") + try: + await vector_store.adelete_nodes(node_ids=node_ids[:2]) + print(" βœ… Successfully deleted 2 nodes") + print(f" Remaining: {vector_store.get_vector_count()}") + except Exception as e: + print(f" ❌ Unexpected error: {e}") + print() + + # Demonstrate concurrent operations with error handling + print("Testing concurrent operations with error handling...") + + async def safe_query(query_text: str, query_id: int): + """Query with error handling.""" + try: + query_embedding = embed_model.get_text_embedding(query_text) + query_obj = VectorStoreQuery( + query_embedding=query_embedding, similarity_top_k=2 + ) + result = await vector_store.aquery(query_obj) + return {"id": query_id, "success": True, "result": result} + except Exception as e: + return {"id": query_id, "success": False, "error": str(e)} + + queries = ["technology", "science", "business"] + results = await asyncio.gather( + *[safe_query(q, i) for i, q in enumerate(queries)], return_exceptions=True + ) + + success_count = sum(1 for r in results if isinstance(r, dict) and r.get("success")) + print(f" βœ… {success_count}/{len(queries)} queries succeeded") + print() + + print(f"βœ… Example 4 completed. Final count: {vector_store.get_vector_count()}") + print() + + +# ============================================================================= +# Example 5: Async with Timeouts +# ============================================================================= +async def example_timeouts(): + """Demonstrate using timeouts with async operations.""" + print("=" * 70) + print("Example 5: Async Operations with Timeouts") + print("=" * 70) + print() + + vector_store = ZeusDBVectorStore(dim=1536, distance="cosine") + embed_model = Settings.embed_model + + # Add documents + nodes = [] + for i in range(10): + node = TextNode( + text=f"Sample document {i} about various topics.", metadata={"doc_id": i} + ) + node.embedding = embed_model.get_text_embedding(node.text) + nodes.append(node) + + print("Adding documents with timeout...") + try: + node_ids = await asyncio.wait_for( + # vector_store.aadd(nodes), + vector_store.async_add(nodes), + timeout=10.0, # 10 second timeout + ) + print(f" βœ… Added {len(node_ids)} nodes within timeout") + except asyncio.TimeoutError: + print(" ❌ Operation timed out") + print() + + # Query with timeout + print("Querying with timeout...") + query_text = "sample topics" + query_embedding = embed_model.get_text_embedding(query_text) + query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=3) + + try: + result = await asyncio.wait_for( + vector_store.aquery(query_obj), + timeout=5.0, # 5 second timeout + ) + result_ids = result.ids or [] + print(f" βœ… Query returned {len(result_ids)} results within timeout") + except asyncio.TimeoutError: + print(" ❌ Query timed out") + print() + + print(f"βœ… Example 5 completed. Final count: {vector_store.get_vector_count()}") + print() + + +# ============================================================================= +# Main Execution +# ============================================================================= +async def main(): + """Run all async examples.""" + try: + await example_basic_async() + await example_concurrent_queries() + await example_concurrent_processing() + await example_error_handling() + await example_timeouts() + + print("=" * 70) + print("Summary: Async Benefits") + print("=" * 70) + print() + print("βœ… Non-blocking operations - don't freeze your application") + print("βœ… Concurrent execution - run multiple operations simultaneously") + print("βœ… Better resource utilization - efficient I/O handling") + print("βœ… Improved throughput - process more requests per second") + print("βœ… Timeout support - prevent operations from hanging") + print() + print("πŸ’‘ When to use async:") + print(" β€’ Web servers handling multiple requests (FastAPI, Starlette)") + print(" β€’ Agents/pipelines with parallel searches") + print(" β€’ Concurrent document processing") + print(" β€’ Notebooks with non-blocking operations") + print() + print("πŸ’‘ When sync is fine:") + print(" β€’ Simple scripts with sequential operations") + print(" β€’ Single-threaded batch processing") + print(" β€’ Quick prototypes and experiments") + print() + print("πŸ“š Available async methods:") + # print(" β€’ aadd(nodes) - Add nodes asynchronously") + print(" β€’ async_add(nodes) - Add nodes asynchronously") + print(" β€’ aquery(query) - Query asynchronously") + print(" β€’ adelete_nodes(node_ids) - Delete by IDs asynchronously") + print() + print("⚠️ Note: aclear() not yet implemented in ZeusDB backend") + print() + + except Exception as e: + print(f"❌ Error in examples: {e}") + import traceback + + traceback.print_exc() + + +if __name__ == "__main__": + # For running as a script + asyncio.run(main()) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/delete_records_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/delete_records_examples.py new file mode 100644 index 0000000000..ca09b1fbbe --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/delete_records_examples.py @@ -0,0 +1,119 @@ +# examples/delete_examples.py +from dotenv import load_dotenv + +from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex +from llama_index.core.vector_stores.types import FilterOperator, MetadataFilters +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() + +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-4") + +vector_store = ZeusDBVectorStore(dim=1536, distance="cosine", index_type="hnsw") +storage_context = StorageContext.from_defaults(vector_store=vector_store) + +documents = [ + Document( + text="Document A1 from project Alpha", + metadata={"project": "alpha", "version": 1}, + doc_id="doc_alpha_1", + ), + Document( + text="Document A2 from project Alpha", + metadata={"project": "alpha", "version": 2}, + doc_id="doc_alpha_2", + ), + Document( + text="Document B1 from project Beta", + metadata={"project": "beta", "version": 1}, + doc_id="doc_beta_1", + ), +] + +index = VectorStoreIndex.from_documents(documents, storage_context=storage_context) + +print("=== ZeusDB Delete Operation Tests ===\n") +print(" Note: ZeusDB HNSW only supports deletion by node ID\n") + +# Show initial state +retriever = index.as_retriever(similarity_top_k=10) +results = retriever.retrieve("document") +initial_count = vector_store.get_vector_count() + +print(f"Initial state: {len(results)} documents, {initial_count} vectors") +for node in results: + print(f" - {node.node.ref_doc_id} (node_id: {node.node.node_id[:8]}...)") + +# Test 1: delete() by ref_doc_id - Should fail +print("\n" + "=" * 60) +print("Test 1: Delete by ref_doc_id (NOT SUPPORTED)") +print("=" * 60) +try: + index.delete_ref_doc(ref_doc_id="doc_alpha_1", delete_from_docstore=False) + print("❌ FAIL: Should have raised NotImplementedError") +except NotImplementedError as e: + print("βœ… PASS: Correctly raised NotImplementedError") + print(f" Message: {str(e)[:80]}...") + +# Test 2: delete_nodes() by ID - Should work +print("\n" + "=" * 60) +print("Test 2: Delete by Node ID (SUPPORTED)") +print("=" * 60) + +results = retriever.retrieve("document") +before_count = vector_store.get_vector_count() + +if results: + node_to_delete = results[0] + node_id = node_to_delete.node.node_id + ref_doc_id = node_to_delete.node.ref_doc_id + + print(f"Before: {before_count} vectors") + print(f"Deleting: {ref_doc_id} (node: {node_id[:8]}...)") + + vector_store.delete_nodes(node_ids=[node_id]) + + after_count = vector_store.get_vector_count() + new_results = retriever.retrieve("document") + + print(f"After: {after_count} vectors, {len(new_results)} retrievable documents") + + if after_count == before_count - 1: + print(f"βœ… PASS: Vector count decreased ({before_count} β†’ {after_count})") + else: + print(f"❌ FAIL: Expected {before_count - 1}, got {after_count}") + + if len(new_results) == len(results) - 1: + print( + f"βœ… PASS: Retrievable docs decreased ({len(results)} β†’ {len(new_results)})" + ) + else: + print(f"❌ FAIL: Expected {len(results) - 1} docs, got {len(new_results)}") + +# Test 3: delete_nodes() with filters - Should fail +print("\n" + "=" * 60) +print("Test 3: Delete with Filters (NOT SUPPORTED)") +print("=" * 60) + +filters = MetadataFilters.from_dicts( + [{"key": "project", "value": "beta", "operator": FilterOperator.EQ}] +) + +try: + vector_store.delete_nodes(filters=filters) + print("❌ FAIL: Should have raised NotImplementedError") +except NotImplementedError as e: + print("βœ… PASS: Correctly raised NotImplementedError") + print(f" Message: {str(e)[:80]}...") + +# Summary +print("\n" + "=" * 60) +print("Summary") +print("=" * 60) +print("βœ… delete_nodes(node_ids=[...]) works correctly") +print("βœ… delete(ref_doc_id='...') correctly raises NotImplementedError") +print("βœ… delete_nodes(filters=...) correctly raises NotImplementedError") +print("\nπŸ“ Recommendation: Track node IDs at application level for deletion") diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/metadata_filter_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/metadata_filter_examples.py new file mode 100644 index 0000000000..bc84fbdaee --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/metadata_filter_examples.py @@ -0,0 +1,184 @@ +# metadata_filter_examples.py +from dotenv import load_dotenv + +from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex +from llama_index.core.vector_stores.types import ( + FilterCondition, + FilterOperator, + MetadataFilters, +) +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() + +# Configure OpenAI +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-4") + +# Create ZeusDB vector store +vector_store = ZeusDBVectorStore(dim=1536, distance="cosine", index_type="hnsw") + +storage_context = StorageContext.from_defaults(vector_store=vector_store) + +# Create documents with rich metadata +documents = [ + Document( + text="ZeusDB is a high-performance vector database optimized for " + "semantic search.", + metadata={ + "source": "doc", + "topic": "zeusdb", + "year": 2024, + "tags": ["db", "vector", "search"], + }, + ), + Document( + text="LlamaIndex provides RAG capabilities including retrievers, query " + "engines, and rerankers.", + metadata={ + "source": "doc", + "topic": "llamaindex", + "year": 2023, + "tags": ["rag", "framework"], + }, + ), + Document( + text="HNSW is a graph-based ANN index enabling fast approximate nearest " + "neighbor search.", + metadata={ + "source": "blog", + "topic": "ann", + "year": 2022, + "tags": ["hnsw", "ann"], + }, + ), + Document( + text="ZeusDB supports cosine distance and integrates with LlamaIndex as a " + "vector store.", + metadata={ + "source": "doc", + "topic": "zeusdb", + "year": 2025, + "tags": ["integration", "vector"], + }, + ), + Document( + text="BM25 and keyword methods focus on exact term matching rather than " + "semantic similarity.", + metadata={ + "source": "blog", + "topic": "ir", + "year": 2021, + "tags": ["bm25", "keyword", "ir"], + }, + ), + Document( + text="Vector search enables semantic similarity. It's commonly paired with " + "metadata filters.", + metadata={ + "source": "note", + "topic": "search", + "year": 2024, + "tags": ["vector", "filters"], + }, + ), +] + +# Build index +index = VectorStoreIndex.from_documents(documents, storage_context=storage_context) + +print("=== Metadata Filter Examples ===\n") + +print("\n=== Example 1: No Filter (baseline) ===") +# First, test without filters to see if basic retrieval works +print("--- Filter: None ---\n") +retriever_baseline = index.as_retriever(similarity_top_k=5) +results_baseline = retriever_baseline.retrieve("ZeusDB database") + +print(f"Found {len(results_baseline)} results without filters:") +for i, node in enumerate(results_baseline, 1): + print(f"{i}. Score: {node.score:.4f}") + print(f" Text: {node.text[:100]}...") + print(f" Metadata: {node.metadata}\n") + + +print("\n=== Example 2: Topic and Year Filter ===") +# Test (a): topic == 'zeusdb' AND year >= 2024 +print("--- Filter: topic=='zeusdb' AND year>=2024 ---\n") +filters_a = MetadataFilters.from_dicts( + [ + {"key": "topic", "value": "zeusdb", "operator": FilterOperator.EQ}, + {"key": "year", "value": 2024, "operator": FilterOperator.GTE}, + ], + condition=FilterCondition.AND, +) +retriever_a = index.as_retriever(similarity_top_k=5, filters=filters_a) +results_a = retriever_a.retrieve("integration with LlamaIndex") + +for i, node in enumerate(results_a, 1): + print(f"{i}. Score: {node.score:.4f}") + print(f" Text: {node.text[:100]}...") + print(f" Metadata: {node.metadata}\n") + + +print("\n=== Example 3: IN Operator - Multiple Values ===") +print("--- Filter: source IN ['blog', 'note'] ---\n") +filters_in = MetadataFilters.from_dicts( + [ + {"key": "source", "value": ["blog", "note"], "operator": FilterOperator.IN}, + ], + condition=FilterCondition.AND, +) +retriever_in = index.as_retriever(similarity_top_k=5, filters=filters_in) +results_in = retriever_in.retrieve("information retrieval methods") + +print(f"Found {len(results_in)} documents from blogs or notes:") +for i, node in enumerate(results_in, 1): + print(f"{i}. Score: {node.score:.4f}") + print(f" Source: {node.metadata.get('source')}") + print(f" Topic: {node.metadata.get('topic')}") + print(f" Text: {node.text[:80]}...\n") + + +print("\n=== Example 4: Array Contains Value ===") +print("--- Filter: tags CONTAINS 'vector' ---\n") +filters_contains = MetadataFilters.from_dicts( + [ + {"key": "tags", "value": "vector", "operator": FilterOperator.CONTAINS}, + ], + condition=FilterCondition.AND, +) +retriever_contains = index.as_retriever(similarity_top_k=5, filters=filters_contains) +results_contains = retriever_contains.retrieve("semantic search") + +print(f"Found {len(results_contains)} documents tagged with 'vector':") +for i, node in enumerate(results_contains, 1): + print(f"{i}. Score: {node.score:.4f}") + print(f" Tags: {node.metadata.get('tags')}") + print(f" Topic: {node.metadata.get('topic')}") + print(f" Text: {node.text[:80]}...\n") + + +print("\n=== Example 5: Exclusion Filter ===") +print("--- Filter: topic != 'zeusdb' AND year >= 2022 ---\n") +filters_ne = MetadataFilters.from_dicts( + [ + {"key": "topic", "value": "zeusdb", "operator": FilterOperator.NE}, + {"key": "year", "value": 2022, "operator": FilterOperator.GTE}, + ], + condition=FilterCondition.AND, +) +retriever_ne = index.as_retriever(similarity_top_k=5, filters=filters_ne) +results_ne = retriever_ne.retrieve("advanced search techniques") + +print(f"Found {len(results_ne)} non-ZeusDB documents from 2022 onwards:") +for i, node in enumerate(results_ne, 1): + print(f"{i}. Score: {node.score:.4f}") + print(f" Topic: {node.metadata.get('topic')} (not 'zeusdb')") + print(f" Year: {node.metadata.get('year')}") + print(f" Text: {node.text[:80]}...\n") + + +print("βœ“ Metadata filtering tests complete") diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/mmr_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/mmr_examples.py new file mode 100644 index 0000000000..067f06f8a1 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/mmr_examples.py @@ -0,0 +1,291 @@ +# examples/mmr_examples.py +""" +ZeusDB MMR (Maximal Marginal Relevance) Examples for LlamaIndex + +Demonstrates using MMR to balance relevance and diversity in search results. + +MMR is useful when you want to avoid redundant results and ensure diverse +perspectives in your retrieved documents. + +Common use cases: +- RAG applications (avoiding repetitive context) +- Document summarization (covering different aspects) +- Research/exploration (finding varied perspectives) +- Question answering (providing comprehensive coverage) +""" + +from dotenv import load_dotenv + +from llama_index.core import Settings +from llama_index.core.schema import TextNode +from llama_index.core.vector_stores.types import VectorStoreQuery +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() + +# Configure OpenAI +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") + +print("=" * 70) +print("ZeusDB MMR (Maximal Marginal Relevance) Examples") +print("=" * 70) +print() + +# ============================================================================= +# Example 1: The Problem - Redundant Results Without MMR +# ============================================================================= +print("=" * 70) +print("Example 1: The Problem - Redundant Results") +print("=" * 70) +print() + +# Create vector store +vector_store = ZeusDBVectorStore(dim=1536, distance="cosine") +embed_model = Settings.embed_model + +# Create documents with some redundancy +documents = [ + # Cluster 1: Python programming (very similar) + "Python is a high-level programming language known for its simplicity.", + "Python is an easy-to-learn programming language with clear syntax.", + "Python programming language is popular for its readability.", + # Cluster 2: Python data science (similar to cluster 1) + "Python is widely used in data science and machine learning.", + "Data scientists prefer Python for analytics and ML tasks.", + # Cluster 3: Other languages (different topic) + "JavaScript is essential for web development and frontend applications.", + "Java is a robust language used in enterprise applications.", + "Rust provides memory safety without garbage collection.", +] + +print(f"Adding {len(documents)} documents...") +nodes = [] +for i, text in enumerate(documents): + node = TextNode(text=text, metadata={"doc_id": i, "source": "examples"}) + node.embedding = embed_model.get_text_embedding(text) + nodes.append(node) + +vector_store.add(nodes) +print(f" βœ… Added {len(nodes)} documents") +print() + +# Query without MMR (standard similarity search) +query_text = "Python programming" +print(f"Query: '{query_text}'") +print() + +query_embedding = embed_model.get_text_embedding(query_text) +query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=5) + +print("πŸ“‹ Standard Similarity Search (No MMR):") +results = vector_store.query(query_obj) + +result_ids = results.ids or [] +result_sims = results.similarities or [] + +# Create ID to node mapping +id_to_node = {node.node_id: node for node in nodes} + +for i, (node_id, similarity) in enumerate(zip(result_ids, result_sims), 1): + # Get the actual text using the mapping + node = id_to_node.get(node_id) + text = node.text if node else "Unknown" + print(f" {i}. Score: {similarity:.4f}") + print(f" {text[:70]}...") + print() + +print("⚠️ Notice: Top results are very similar (redundant information)") +print() + +# ============================================================================= +# Example 2: The Solution - Diverse Results With MMR +# ============================================================================= +print("=" * 70) +print("Example 2: The Solution - MMR for Diversity") +print("=" * 70) +print() + +print(f"Query: '{query_text}'") +print() + +# Query WITH MMR +print("🎯 MMR Search (Balanced relevance + diversity):") +mmr_results = vector_store.query( + query_obj, + mmr=True, + fetch_k=8, # Fetch more candidates + mmr_lambda=0.5, # Balance: 0.5 relevance, 0.5 diversity +) + +mmr_ids = mmr_results.ids or [] +mmr_sims = mmr_results.similarities or [] + +for i, (node_id, similarity) in enumerate(zip(mmr_ids, mmr_sims), 1): + # Get the actual text using the mapping + node = id_to_node.get(node_id) + text = node.text if node else "Unknown" + print(f" {i}. Score: {similarity:.4f}") + print(f" {text[:70]}...") + print() + +print("βœ… Notice: Results are more diverse (different aspects/topics)") +print() + +# ============================================================================= +# Example 3: Tuning Lambda - Relevance vs Diversity Tradeoff +# ============================================================================= +print("=" * 70) +print("Example 3: Tuning Lambda (Relevance vs Diversity)") +print("=" * 70) +print() + +print("Lambda controls the relevance-diversity tradeoff:") +print(" β€’ lambda=1.0: Pure relevance (like standard search)") +print(" β€’ lambda=0.5: Balanced (default)") +print(" β€’ lambda=0.0: Maximum diversity") +print() + +lambda_values = [1.0, 0.7, 0.5, 0.3, 0.0] + +for lambda_val in lambda_values: + print(f"πŸ“Š MMR with lambda={lambda_val}:") + + mmr_results = vector_store.query( + query_obj, mmr=True, fetch_k=8, mmr_lambda=lambda_val + ) + + mmr_ids = mmr_results.ids or [] + + # Just show first 3 results + print(" Top 3 results:") + for i, node_id in enumerate(mmr_ids[:3], 1): + node = id_to_node.get(node_id) + text = node.text if node else "Unknown" + print(f" {i}. {text[:60]}...") + print() + +# ============================================================================= +# Example 4: Practical RAG Scenario +# ============================================================================= +print("=" * 70) +print("Example 4: Practical RAG Scenario") +print("=" * 70) +print() + +# Create a new store with diverse tech articles +rag_store = ZeusDBVectorStore(dim=1536, distance="cosine") + +articles = [ + # AI/ML cluster + "Machine learning models require large datasets for training accuracy.", + "Deep learning neural networks have revolutionized computer vision.", + "AI ethics is becoming increasingly important in model development.", + # Cloud/Infrastructure cluster + "Cloud computing provides scalable infrastructure for applications.", + "Kubernetes orchestrates containerized applications in production.", + "Serverless computing reduces operational overhead significantly.", + # Security cluster + "Cybersecurity threats are evolving with sophisticated attack vectors.", + "Zero-trust architecture is the modern approach to network security.", + "Encryption protects sensitive data from unauthorized access.", + # DevOps cluster + "CI/CD pipelines automate software deployment and testing.", + "Infrastructure as code enables reproducible deployments.", + "Monitoring and observability are critical for system reliability.", +] + +print(f"Creating knowledge base with {len(articles)} articles...") +rag_nodes = [] +for i, text in enumerate(articles): + node = TextNode(text=text, metadata={"doc_id": i, "source": "tech_articles"}) + node.embedding = embed_model.get_text_embedding(text) + rag_nodes.append(node) + +rag_store.add(rag_nodes) +print(f" βœ… Added {len(rag_nodes)} articles") +print() + +# Create ID to node mapping for RAG nodes +rag_id_to_node = {node.node_id: node for node in rag_nodes} + +# User question +question = "How do modern software systems ensure quality and security?" +print(f"Question: '{question}'") +print() + +query_embedding = embed_model.get_text_embedding(question) +query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=4) + +# Standard search - might get redundant results +print("πŸ“‹ Standard Search (Potential redundancy):") +standard_results = rag_store.query(query_obj) +standard_ids = standard_results.ids or [] + +for i, node_id in enumerate(standard_ids, 1): + node = rag_id_to_node.get(node_id) + text = node.text if node else "Unknown" + print(f" {i}. {text}") +print() + +# MMR search - ensures diverse perspectives +print("🎯 MMR Search (Diverse perspectives for comprehensive answer):") +mmr_results = rag_store.query( + query_obj, + mmr=True, + fetch_k=12, # Larger candidate pool + mmr_lambda=0.6, # Slightly favor relevance +) +mmr_ids = mmr_results.ids or [] + +for i, node_id in enumerate(mmr_ids, 1): + node = rag_id_to_node.get(node_id) + text = node.text if node else "Unknown" + print(f" {i}. {text}") +print() + +print("βœ… MMR provides diverse context covering:") +print(" β€’ Security aspects") +print(" β€’ Quality/testing practices") +print(" β€’ Infrastructure considerations") +print(" β€’ Monitoring and reliability") +print() + +# ============================================================================= +# Summary +# ============================================================================= +print("=" * 70) +print("Summary: When to Use MMR") +print("=" * 70) +print() +print("βœ… Use MMR when you want:") +print(" β€’ Diverse results instead of similar/redundant ones") +print(" β€’ Multiple perspectives on a topic") +print(" β€’ Comprehensive coverage for RAG applications") +print(" β€’ To avoid echo chamber effects in results") +print() +print("πŸ“Š Key Parameters:") +print(" β€’ mmr=True: Enable MMR re-ranking") +print(" β€’ fetch_k: Candidate pool size (default: 4*k, min: 20)") +print(" β€’ mmr_lambda: Relevance/diversity tradeoff (default: 0.7)") +print(" - 1.0 = pure relevance (like standard search)") +print(" - 0.5 = balanced") +print(" - 0.0 = maximum diversity") +print() +print("πŸ’‘ Best Practices:") +print(" β€’ Set fetch_k >> k (e.g., fetch_k=20 for k=5)") +print(" β€’ Start with lambda=0.7 (slightly favor relevance)") +print(" β€’ Lower lambda for more diverse results") +print(" β€’ Use with return_vector=True for better diversity calculation") +print() +print("⚠️ Trade-offs:") +print(" β€’ MMR is slower than standard search (needs re-ranking)") +print(" β€’ Requires fetch_k candidates (more initial results)") +print(" β€’ May return less relevant results for diversity") +print() +print("πŸ“š Common Use Cases:") +print(" β€’ RAG systems: Diverse context for comprehensive answers") +print(" β€’ Document exploration: Finding different perspectives") +print(" β€’ Research: Covering multiple aspects of a topic") +print(" β€’ Recommendation systems: Avoiding filter bubbles") +print() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/persistence_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/persistence_examples.py new file mode 100644 index 0000000000..8fb367eefc --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/persistence_examples.py @@ -0,0 +1,400 @@ +# examples/persistence_examples.py +""" +ZeusDB Persistence Examples for LlamaIndex + +Demonstrates how to save and load ZeusDB indexes with full state preservation, +including vectors, metadata, HNSW graph structure, and quantization models. +""" + +from pathlib import Path +import shutil + +from dotenv import load_dotenv + +from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex +from llama_index.core.vector_stores.types import VectorStoreQuery +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() + +# Configure OpenAI +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-4") + +print("=" * 70) +print("ZeusDB Persistence Examples") +print("=" * 70) +print() + +# ============================================================================= +# Example 1: Basic Save and Load +# ============================================================================= +print("=" * 70) +print("Example 1: Basic Save and Load") +print("=" * 70) +print() + +# Create sample documents +documents = [ + Document( + text=f"Document {i}: This is a sample document about technology " + f"and artificial intelligence.", + metadata={"doc_id": i, "category": "tech"}, + ) + for i in range(50) +] + +print(f"Creating index with {len(documents)} documents...") + +# Create vector store and index +vector_store = ZeusDBVectorStore(dim=1536, distance="cosine", index_type="hnsw") + +storage_context = StorageContext.from_defaults(vector_store=vector_store) +index = VectorStoreIndex.from_documents( + documents, storage_context=storage_context, show_progress=True +) + +print() +print("Original index:") +print(f" Vector count: {vector_store.get_vector_count()}") +print(f" Index info: {vector_store.info()}") +print() + +# Test search before saving +query_text = "artificial intelligence" +print(f"Testing search before save: '{query_text}'") + +# Get embedding for query - use get_text_embedding_batch for single query +embed_model = Settings.embed_model +query_embedding = embed_model.get_text_embedding(query_text) + +# Direct query to vector store +query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=3) +results = vector_store.query(query_obj) + +# βœ… Handle None values safely +result_ids = results.ids or [] +result_sims = results.similarities or [] + +print(f" Found {len(result_ids)} results") +for i, (node_id, similarity) in enumerate(zip(result_ids, result_sims), 1): + print(f" {i}. Score: {similarity:.4f}, ID: {node_id}") +print() + +# Save the index +save_path = "test_index.zdb" +print(f"Saving index to '{save_path}'...") +success = vector_store.save_index(save_path) +print(f" Save successful: {success}") +print() + +# Load the index +print(f"Loading index from '{save_path}'...") +loaded_vector_store = ZeusDBVectorStore.load_index(save_path) +print(" Load successful!") +print() + +# Verify loaded index +print("Loaded index:") +print(f" Vector count: {loaded_vector_store.get_vector_count()}") +print(f" Index info: {loaded_vector_store.info()}") +print() + +# Test search after loading - direct query +print(f"Testing search after load: '{query_text}'") +loaded_results = loaded_vector_store.query(query_obj) + +# βœ… Handle None values safely +loaded_ids = loaded_results.ids or [] +loaded_sims = loaded_results.similarities or [] + +print(f" Found {len(loaded_ids)} results") +for i, (node_id, similarity) in enumerate(zip(loaded_ids, loaded_sims), 1): + print(f" {i}. Score: {similarity:.4f}, ID: {node_id}") +print() + +# Compare results +print("Comparing search results:") +if len(result_ids) == len(loaded_ids): + print(f" βœ… Same number of results: {len(result_ids)}") + + if result_ids and loaded_ids: + # Compare top result IDs + if result_ids[0] == loaded_ids[0]: + print(f" βœ… Top result matches: ID={result_ids[0]}") + else: + print(f" ⚠️ Top results differ: {result_ids[0]} vs {loaded_ids[0]}") + + # Compare similarities (should be very close) + if result_sims and loaded_sims: + sim_diff = abs(result_sims[0] - loaded_sims[0]) + if sim_diff < 0.001: + print(f" βœ… Similarities match: {result_sims[0]:.4f}") + else: + print(f" ⚠️ Similarities differ by {sim_diff:.4f}") +else: + print(" ⚠️ Different number of results") +print() + +# ============================================================================= +# Example 2: Save and Load with Quantization +# ============================================================================= +print("=" * 70) +print("Example 2: Save and Load with Quantization") +print("=" * 70) +print() + +# Create index with quantization +print("Creating quantized index...") +quantization_config = { + "type": "pq", + "subvectors": 8, + "bits": 8, + "training_size": 1000, + "storage_mode": "quantized_only", +} + +quant_vector_store = ZeusDBVectorStore( + dim=1536, + distance="cosine", + index_type="hnsw", + quantization_config=quantization_config, +) + +quant_storage_context = StorageContext.from_defaults(vector_store=quant_vector_store) + +# Create more documents to trigger quantization +quant_documents = [ + Document( + text=f"Document {i}: Technology content for quantization testing.", + metadata={"doc_id": i, "category": "tech"}, + ) + for i in range(1100) # Enough to trigger training +] + +print(f"Adding {len(quant_documents)} documents...") +quant_index = VectorStoreIndex.from_documents( + quant_documents, storage_context=quant_storage_context, show_progress=False +) + +print() +print("Original quantized index:") +print(f" Vector count: {quant_vector_store.get_vector_count()}") +print(f" Is quantized: {quant_vector_store.is_quantized()}") +print(f" Training progress: {quant_vector_store.get_training_progress():.1f}%") +print(f" Storage mode: {quant_vector_store.get_storage_mode()}") + +qi = quant_vector_store.get_quantization_info() +if qi: + print(f" Compression: {qi.get('compression_ratio', 0):.1f}x") + print(f" Memory: {qi.get('memory_mb', 0):.2f} MB") + print(f" Trained: {qi.get('is_trained', False)}") +print() + +# Save quantized index +quant_save_path = "quantized_index.zdb" +print(f"Saving quantized index to '{quant_save_path}'...") +quant_success = quant_vector_store.save_index(quant_save_path) +print(f" Save successful: {quant_success}") +print() + +# Load quantized index +print(f"Loading quantized index from '{quant_save_path}'...") +loaded_quant_vs = ZeusDBVectorStore.load_index(quant_save_path) +print(" Load successful!") +print() + +# Verify quantization state preserved +print("Loaded quantized index:") +print(f" Vector count: {loaded_quant_vs.get_vector_count()}") +print(f" Is quantized: {loaded_quant_vs.is_quantized()}") +print(f" Training progress: {loaded_quant_vs.get_training_progress():.1f}%") +print(f" Storage mode: {loaded_quant_vs.get_storage_mode()}") + +loaded_qi = loaded_quant_vs.get_quantization_info() +if loaded_qi: + print(f" Compression: {loaded_qi.get('compression_ratio', 0):.1f}x") + print(f" Memory: {loaded_qi.get('memory_mb', 0):.2f} MB") + print(f" Trained: {loaded_qi.get('is_trained', False)}") +print() + +# Verify quantization state matches +print("Verifying quantization state preservation:") + +# Check original state +orig_quantized = quant_vector_store.is_quantized() +orig_storage = quant_vector_store.get_storage_mode() + +# Check loaded state +loaded_quantized = loaded_quant_vs.is_quantized() +loaded_storage = loaded_quant_vs.get_storage_mode() + +if orig_quantized == loaded_quantized: + print(f" βœ… Quantization state preserved: {loaded_quantized}") +else: + print(" ⚠️ Quantization state differs:") + print(f" Original: is_quantized={orig_quantized}, mode={orig_storage}") + print(f" Loaded: is_quantized={loaded_quantized}, mode={loaded_storage}") + print() + print(" ℹ️ Known Limitation (Current Release):") + print(" Quantized indexes load in raw vector mode.") + print(" Training state and config are preserved.") + print(" Search works perfectly, just without memory compression.") + print(" This will be fixed in the next ZeusDB release.") + print() + +if ( + quant_vector_store.get_training_progress() + == loaded_quant_vs.get_training_progress() +): + progress = loaded_quant_vs.get_training_progress() + print(f" βœ… Training progress preserved: {progress:.1f}%") +else: + print(" ⚠️ Training progress differs") + +if qi and loaded_qi: + if qi.get("compression_ratio") == loaded_qi.get("compression_ratio"): + ratio = loaded_qi.get("compression_ratio") + print(f" βœ… Compression ratio preserved: {ratio:.1f}x") + else: + print(" ⚠️ Compression ratio differs") + + if qi.get("is_trained") == loaded_qi.get("is_trained"): + trained = loaded_qi.get("is_trained") + print(f" βœ… Training status preserved: {trained}") + else: + print(" ⚠️ Training status differs") +print() + +# Test search on loaded quantized index +print("Testing search on loaded quantized index...") +quant_query_text = "technology and AI" +quant_query_embedding = embed_model.get_text_embedding(quant_query_text) +quant_query_obj = VectorStoreQuery( + query_embedding=quant_query_embedding, similarity_top_k=3 +) + +loaded_quant_results = loaded_quant_vs.query(quant_query_obj) + +# βœ… Handle None values safely +loaded_quant_ids = loaded_quant_results.ids or [] +loaded_quant_sims = loaded_quant_results.similarities or [] + +print(f" Found {len(loaded_quant_ids)} results") +for i, (node_id, similarity) in enumerate(zip(loaded_quant_ids, loaded_quant_sims), 1): + print(f" {i}. Score: {similarity:.4f}, ID: {node_id}") +print() + +# ============================================================================= +# Example 3: Multiple Save/Load Cycles +# ============================================================================= +print("=" * 70) +print("Example 3: Multiple Save/Load Cycles") +print("=" * 70) +print() + +# Create initial index +cycle_docs = [ + Document(text=f"Cycle document {i}", metadata={"doc_id": i, "version": 1}) + for i in range(100) +] + +print("Creating initial index...") +cycle_vs = ZeusDBVectorStore(dim=1536, distance="cosine") +cycle_sc = StorageContext.from_defaults(vector_store=cycle_vs) +cycle_idx = VectorStoreIndex.from_documents( + cycle_docs, storage_context=cycle_sc, show_progress=False +) + +initial_count = cycle_vs.get_vector_count() +print(f" Initial vector count: {initial_count}") +print() + +# Save and load multiple times +for cycle in range(1, 4): + save_path = f"cycle_{cycle}.zdb" + print(f"Cycle {cycle}: Save -> Load") + + # Save + print(f" Saving to '{save_path}'...") + cycle_vs.save_index(save_path) + + # Load + print(f" Loading from '{save_path}'...") + cycle_vs = ZeusDBVectorStore.load_index(save_path) + + # Verify + current_count = cycle_vs.get_vector_count() + current_info = cycle_vs.info() + print(f" Vector count after load: {current_count}") + print(f" Index info: {current_info}") + + if current_count == initial_count: + print(f" βœ… Vector count stable across cycle {cycle}") + else: + print(f" ⚠️ Vector count changed: {initial_count} -> {current_count}") + print() + +print("All cycles completed successfully!") +print() + +# ============================================================================= +# Example 4: Cleanup Test Files +# ============================================================================= +print("=" * 70) +print("Example 4: Cleanup") +print("=" * 70) +print() + +test_paths = [ + "test_index.zdb", + "quantized_index.zdb", + "cycle_1.zdb", + "cycle_2.zdb", + "cycle_3.zdb", +] + +print("Cleaning up test files:") +for path in test_paths: + if Path(path).exists(): + try: + shutil.rmtree(path) + print(f" βœ… Removed: {path}") + except Exception as e: + print(f" ⚠️ Failed to remove {path}: {e}") + else: + print(f" ℹ️ Not found: {path}") +print() + +# ============================================================================= +# Summary +# ============================================================================= +print("=" * 70) +print("Summary: Persistence Benefits") +print("=" * 70) +print() +print("βœ… Complete state preservation - vectors, metadata, HNSW graph") +print("βœ… Quantization support - PQ models and training state preserved") +print("βœ… Cross-platform compatibility - portable between systems") +print("βœ… Directory structure - organized file layout (.zdb directory)") +print("βœ… Direct querying - use vector_store.query() after loading") +print() +print("πŸ’‘ Key Points:") +print(" β€’ Save path creates a directory, not a single file") +print(" β€’ Quantized indexes preserve compression models and training state") +print(" β€’ Loaded indexes are ready for immediate use via .query()") +print(" β€’ Use save_index() and load_index() class method") +print() +print("⚠️ Note on LlamaIndex Integration:") +print(" β€’ ZeusDB doesn't store full text (only embeddings + metadata)") +print(" β€’ Use vector_store.query() directly instead of VectorStoreIndex") +print(" β€’ For full LlamaIndex integration, store documents separately") +print() +print("πŸ“š Persistence is essential for:") +print(" β€’ Production deployments with long-lived indexes") +print(" β€’ Sharing indexes between systems or team members") +print(" β€’ Implementing backup and recovery strategies") +print(" β€’ Avoiding re-indexing on application restart") +print() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quantization_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quantization_examples.py new file mode 100644 index 0000000000..08ee03295e --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quantization_examples.py @@ -0,0 +1,369 @@ +# examples/quantization_examples.py +""" +ZeusDB Product Quantization Examples for LlamaIndex + +Demonstrates how to use Product Quantization (PQ) for memory-efficient +vector storage with the LlamaIndex ZeusDB integration. +""" + +from dotenv import load_dotenv + +from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() + +# Configure OpenAI +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-4") # βœ… Fixed: Changed from gpt-5 to gpt-4 + +print("=" * 70) +print("ZeusDB Product Quantization Examples") +print("=" * 70) +print() + +# ============================================================================= +# Example 1: Basic Quantization Setup +# ============================================================================= +print("=" * 70) +print("Example 1: Basic Quantization Configuration") +print("=" * 70) +print() + +# Configure quantization for memory efficiency +quantization_config = { + "type": "pq", # Product Quantization + "subvectors": 8, # Divide 1536-dim into 8 subvectors + "bits": 8, # 256 centroids per subvector (2^8) + "training_size": 1000, # Minimum: 1000 (required by backend) + "storage_mode": "quantized_only", # Memory-optimized mode +} + +print("Quantization Config:") +for key, value in quantization_config.items(): + print(f" {key}: {value}") +print() + +# Create vector store with quantization +vector_store = ZeusDBVectorStore( + dim=1536, + distance="cosine", + index_type="hnsw", + quantization_config=quantization_config, +) + +storage_context = StorageContext.from_defaults(vector_store=vector_store) + +print("Initial State:") +print(f" Vector count: {vector_store.get_vector_count()}") +print(f" Is quantized: {vector_store.is_quantized()}") +print(f" Can use quantization: {vector_store.can_use_quantization()}") +print(f" Storage mode: {vector_store.get_storage_mode()}") +print(f" Training progress: {vector_store.get_training_progress():.1f}%") +print() + +# ============================================================================= +# Example 2: Adding Documents and Triggering Training +# ============================================================================= +print("=" * 70) +print("Example 2: Adding Documents (No Training in This Example)") +print("=" * 70) +print() + +# Create sample documents +documents = [ + Document( + text=f"Document {i}: This is a sample document about technology, " + f"artificial intelligence, and machine learning in the year {2020 + (i % 5)}.", + metadata={"doc_id": i, "category": "tech", "year": 2020 + (i % 5)}, + ) + for i in range(150) # Small sample - won't trigger training (need 1000) +] + +print(f"Adding {len(documents)} documents...") +print("⚠️ Note: Training requires 1000 documents, so it won't trigger here") +print(" (See Example 4 for actual training demonstration)") +print() + +# Build index - this will add all documents +index = VectorStoreIndex.from_documents( + documents, storage_context=storage_context, show_progress=True +) + +print() +print("After Adding Documents:") +print(f" Vector count: {vector_store.get_vector_count()}") +print(f" Is quantized: {vector_store.is_quantized()}") +print(f" Can use quantization: {vector_store.can_use_quantization()}") +print(f" Storage mode: {vector_store.get_storage_mode()}") +print(f" Training progress: {vector_store.get_training_progress():.1f}%") +print() + +# Get detailed quantization info +quant_info = vector_store.get_quantization_info() +if quant_info: + print("Quantization Details:") + for key, value in quant_info.items(): + if isinstance(value, float): + print(f" {key}: {value:.2f}") + else: + print(f" {key}: {value}") + print() +else: + print("ℹ️ Quantization info not available (training not yet triggered)") + print() + +# ============================================================================= +# Example 3: Querying with Index +# ============================================================================= +print("=" * 70) +print("Example 3: Searching (Without Quantization)") +print("=" * 70) +print() + +query_engine = index.as_query_engine(similarity_top_k=3) + +query = "Tell me about artificial intelligence and machine learning" +print(f"Query: {query}") +print() + +response = query_engine.query(query) +print(f"Response: {response}") +print() + +# Also try direct retrieval +retriever = index.as_retriever(similarity_top_k=5) +results = retriever.retrieve("technology and AI") + +print("Direct Retrieval Results (top 5):") +for i, node in enumerate(results, 1): + print(f"{i}. Score: {node.score:.4f}") + print(f" Text: {node.text[:80]}...") + print(f" Metadata: {node.metadata}") + print() + +# ============================================================================= +# Example 4: Comparing Different Quantization Configurations +# ============================================================================= +print("=" * 70) +print("Example 4: Comparing Quantization Configurations") +print("=" * 70) +print() + +configs = [ + { + "name": "High Compression (Memory Optimized)", + "config": { + "type": "pq", + "subvectors": 16, + "bits": 6, + "training_size": 1000, + "storage_mode": "quantized_only", + }, + "description": "~32x compression, lowest memory usage", + }, + { + "name": "Balanced (Recommended)", + "config": { + "type": "pq", + "subvectors": 8, + "bits": 8, + "training_size": 1000, + "storage_mode": "quantized_only", + }, + "description": "~16x compression, good accuracy/memory balance", + }, + { + "name": "High Accuracy (Keep Raw Vectors)", + "config": { + "type": "pq", + "subvectors": 4, + "bits": 8, + "training_size": 1000, + "storage_mode": "quantized_with_raw", + }, + "description": "~4x compression, keeps raw vectors", + }, +] + +# Create a test dataset - needs to be larger to trigger training +test_docs = [ + Document( + text=f"Test document {i} about various topics including science, " + f"technology, and research.", + metadata={"id": i, "category": "test"}, + ) + for i in range(1100) # Just over training threshold +] + +print("Testing different configurations with same dataset:") +print(f"(Using {len(test_docs)} documents for each configuration)") +print("⚠️ Note: This will take several minutes as it creates 3 indexes") +print(" and trains quantization for each one.") +print() + +for config_info in configs: + print(f"Configuration: {config_info['name']}") + print(f" Description: {config_info['description']}") + print(f" Settings: {config_info['config']}") + + # Create new vector store with this config + vs = ZeusDBVectorStore( + dim=1536, + distance="cosine", + index_type="hnsw", + quantization_config=config_info["config"], + ) + + sc = StorageContext.from_defaults(vector_store=vs) + + print(" Building index (this may take 1-2 minutes)...") + idx = VectorStoreIndex.from_documents( + test_docs, storage_context=sc, show_progress=False + ) + + print(" Results:") + print(f" Vector count: {vs.get_vector_count()}") + print(f" Is quantized: {vs.is_quantized()}") + print(f" Can use quantization: {vs.can_use_quantization()}") + print(f" Storage mode: {vs.get_storage_mode()}") + print(f" Training progress: {vs.get_training_progress():.1f}%") + + qi = vs.get_quantization_info() + if qi: + if "compression_ratio" in qi: + print(f" Compression ratio: {qi['compression_ratio']:.1f}x") + if "memory_mb" in qi: + print(f" Memory usage: {qi['memory_mb']:.2f} MB") + if "is_trained" in qi: + print(f" Training complete: {qi['is_trained']}") + + print() + +# ============================================================================= +# Example 5: Monitoring Training Progress +# ============================================================================= +print("=" * 70) +print("Example 5: Monitoring Training Progress") +print("=" * 70) +print() + +# Create fresh vector store +monitor_vs = ZeusDBVectorStore( + dim=1536, + distance="cosine", + index_type="hnsw", + quantization_config={ + "type": "pq", + "subvectors": 8, + "bits": 8, + "training_size": 1000, + "storage_mode": "quantized_only", + }, +) + +monitor_sc = StorageContext.from_defaults(vector_store=monitor_vs) + +print("Creating and adding documents to monitor training progress...") +print("Note: Training triggers at 1000 documents") +print() + +# Create all documents +print("Creating 1200 documents...") +all_docs = [ + Document( + text=f"Document {i}: Technology and AI content for training demonstration.", + metadata={"id": i, "category": "demo"}, + ) + for i in range(1200) +] + +print("Before adding:") +print(f" Vector count: {monitor_vs.get_vector_count()}") +print(f" Training progress: {monitor_vs.get_training_progress():.1f}%") +print(f" Is quantized: {monitor_vs.is_quantized()}") +print(f" Storage mode: {monitor_vs.get_storage_mode()}") +print() + +# βœ… Split documents into first batch and remaining +batch_size = 300 +first_batch = all_docs[:batch_size] +remaining_docs = all_docs[batch_size:] + +print(f"Adding {len(all_docs)} documents in batches of {batch_size}...") +print() + +# βœ… Create index with first batch (now monitor_idx is always defined) +print(f"Batch 1: Creating index with {len(first_batch)} documents...", end=" ") +monitor_idx = VectorStoreIndex.from_documents( + first_batch, storage_context=monitor_sc, show_progress=False +) + +count = monitor_vs.get_vector_count() +progress = monitor_vs.get_training_progress() +is_quantized = monitor_vs.is_quantized() + +print("Done!") +print(f" Vectors: {count}, Progress: {progress:.1f}%, Quantized: {is_quantized}") +print() + +# βœ… Add remaining documents in batches +remaining_batches = [ + remaining_docs[i : i + batch_size] + for i in range(0, len(remaining_docs), batch_size) +] + +for batch_num, batch in enumerate(remaining_batches, start=2): + print(f"Batch {batch_num}: Adding {len(batch)} documents...", end=" ") + for doc in batch: + monitor_idx.insert(doc) + + count = monitor_vs.get_vector_count() + progress = monitor_vs.get_training_progress() + is_quantized = monitor_vs.is_quantized() + + print("Done!") + print(f" Vectors: {count}, Progress: {progress:.1f}%, Quantized: {is_quantized}") + + if is_quantized and progress == 100.0: + print(" πŸŽ‰ Training completed!") + qi = monitor_vs.get_quantization_info() + if qi and "compression_ratio" in qi: + print(f" Compression: {qi['compression_ratio']:.1f}x") + print() + +print("Final state:") +print(f" Vector count: {monitor_vs.get_vector_count()}") +print(f" Training progress: {monitor_vs.get_training_progress():.1f}%") +print(f" Is quantized: {monitor_vs.is_quantized()}") +print(f" Storage mode: {monitor_vs.get_storage_mode()}") +print() + +# ============================================================================= +# Summary +# ============================================================================= +print("=" * 70) +print("Summary: Product Quantization Benefits") +print("=" * 70) +print() +print("βœ… Memory Reduction: 4x-256x depending on configuration") +print("βœ… Automatic Training: Triggers at configured threshold (minimum 1000)") +print("βœ… Transparent Search: Works seamlessly with quantized vectors") +print("βœ… Monitoring: Track training progress and quantization status") +print() +print("πŸ’‘ Recommendations:") +print(" β€’ Start with balanced config (subvectors=8, bits=8)") +print(" β€’ Use 'quantized_only' for maximum memory savings") +print(" β€’ Set training_size to 1000-10000 based on dataset size") +print(" β€’ Monitor training progress with get_training_progress()") +print() +print("⚠️ Important:") +print(" β€’ Minimum training_size is 1000 (enforced by backend)") +print(" β€’ Training occurs once when threshold is reached") +print(" β€’ Larger training sets generally produce better quantization") +print() +print("πŸ“š For large datasets (100k+ vectors), quantization provides") +print(" significant memory savings with minimal accuracy impact.") +print() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quickstart.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quickstart.py new file mode 100644 index 0000000000..2f82d514b9 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/quickstart.py @@ -0,0 +1,38 @@ +# quickstart.py +from dotenv import load_dotenv + +from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex +from llama_index.embeddings.openai import OpenAIEmbedding +from llama_index.llms.openai import OpenAI +from llama_index.vector_stores.zeusdb import ZeusDBVectorStore + +load_dotenv() # reads .env and sets OPENAI_API_KEY + +# Set up embedding model +Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small") +Settings.llm = OpenAI(model="gpt-5") + +# Create ZeusDB vector store +vector_store = ZeusDBVectorStore( + dim=1536, # OpenAI embedding dimension + distance="cosine", + index_type="hnsw", +) + +# Create storage context +storage_context = StorageContext.from_defaults(vector_store=vector_store) + +# Create documents +documents = [ + Document(text="ZeusDB is a high-performance vector database."), + Document(text="LlamaIndex provides RAG capabilities."), + Document(text="Vector search enables semantic similarity."), +] + +# Create index and store documents +index = VectorStoreIndex.from_documents(documents, storage_context=storage_context) + +# Query the index +query_engine = index.as_query_engine() +response = query_engine.query("What is ZeusDB?") +print(response) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/__init__.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/__init__.py new file mode 100644 index 0000000000..f604771f44 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/__init__.py @@ -0,0 +1,5 @@ +"""ZeusDB vector store integration for LlamaIndex.""" + +from llama_index.vector_stores.zeusdb.base import ZeusDBVectorStore + +__all__ = ["ZeusDBVectorStore"] diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py new file mode 100644 index 0000000000..3fa6dc57e1 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py @@ -0,0 +1,1177 @@ +# llama_index/vector_stores/zeusdb/base.py + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Iterator, MutableMapping, Sequence +from contextlib import contextmanager +from time import perf_counter +from typing import TYPE_CHECKING, Any, cast + +from llama_index.core.schema import BaseNode +from llama_index.core.vector_stores.types import ( + BasePydanticVectorStore, + FilterCondition, + FilterOperator, + MetadataFilter, + MetadataFilters, + VectorStoreQuery, + VectorStoreQueryResult, +) + +# ZeusDB runtime (umbrella package only) +from zeusdb import VectorDatabase # type: ignore + +if TYPE_CHECKING: + pass + +# ----------------------------------------------------------------------------- +# Enterprise Logging Integration with Safe Fallback: llamaindex-zeusdb package +# ----------------------------------------------------------------------------- +try: + from zeusdb.logging_config import ( # type: ignore[import] # noqa: I001 + get_logger as _get_logger, + ) + from zeusdb.logging_config import ( # type: ignore[import] + operation_context as _operation_context, + ) +except Exception: # fallback for OSS/dev environments + import logging + + class _StructuredAdapter(logging.LoggerAdapter): + """ + Adapter that moves arbitrary kwargs into 'extra' + for stdlib logging compatibility. + """ + + def process( + self, msg: str, kwargs: MutableMapping[str, Any] + ) -> tuple[str, MutableMapping[str, Any]]: + allowed = {"exc_info", "stack_info", "stacklevel", "extra"} + extra = kwargs.get("extra", {}) or {} + if not isinstance(extra, dict): + extra = {"_extra": repr(extra)} # defensive + fields = {k: kwargs.pop(k) for k in list(kwargs.keys()) if k not in allowed} + if fields: + extra.update(fields) + kwargs["extra"] = extra + return msg, kwargs + + def _get_logger(name: str) -> logging.LoggerAdapter: + base = logging.getLogger(name) + if not base.handlers: + base.addHandler(logging.NullHandler()) + return _StructuredAdapter(base, {}) + + @contextmanager + def _operation_context(operation_name: str, **context: Any) -> Iterator[None]: + logger.debug( + f"{operation_name} started", + operation=operation_name, + **context, + ) + start = perf_counter() + try: + yield + duration_ms = (perf_counter() - start) * 1000 + logger.info( + f"{operation_name} completed", + operation=operation_name, + duration_ms=duration_ms, + **context, + ) + except Exception as e: + duration_ms = (perf_counter() - start) * 1000 + logger.error( + f"{operation_name} failed", + operation=operation_name, + duration_ms=duration_ms, + error=str(e), + exc_info=True, + **context, + ) + raise + + +# Initialize module logger (central config owns handlers/format) +get_logger: Callable[[str], Any] = cast(Callable[[str], Any], _get_logger) +logger = get_logger("llamaindex_zeusdb") +operation_context = cast(Callable[..., Any], _operation_context) + +# ------------------------- +# Utilities & type helpers +# ------------------------- + +_DISTANCE_TO_SPACE = { + "cosine": "cosine", + "l2": "l2", + "euclidean": "l2", + "l1": "l1", + "manhattan": "l1", +} + + +def _infer_space(distance: str | None) -> str: + if not distance: + return "cosine" + key = distance.lower() + return _DISTANCE_TO_SPACE.get(key, "cosine") + + +def _similarity_from_distance(distance_value: float, space: str) -> float: + """ + Convert ZeusDB distance to a similarity score (higher = better). + - cosine: similarity = 1 - distance (assuming normalized embeddings). + - l2/l1: convert to negative distance so higher is better. + """ + if space == "cosine": + return 1.0 - float(distance_value) + return -float(distance_value) + + +def _extract_embedding(node: BaseNode) -> list[float] | None: + # LlamaIndex nodes typically have `embedding` populated before add() + emb = getattr(node, "embedding", None) + if emb is None and hasattr(node, "get_embedding"): + try: + emb = node.get_embedding() # type: ignore[attr-defined] + except Exception: + emb = None + return list(emb) if emb is not None else None + + +def _node_metadata(node: BaseNode) -> dict[str, Any]: + md = dict(node.metadata or {}) + # Propagate identifiers for delete-by-ref_doc_id and traceability. + if getattr(node, "ref_doc_id", None): + md["ref_doc_id"] = node.ref_doc_id + if getattr(node, "id_", None): + md["node_id"] = node.id_ + return md + + +def _translate_filter_op(op: FilterOperator) -> str: + """Map LlamaIndex FilterOperator -> ZeusDB operator keys.""" + mapping = { + FilterOperator.EQ: "eq", + FilterOperator.NE: "ne", + FilterOperator.GT: "gt", + FilterOperator.LT: "lt", + FilterOperator.GTE: "gte", + FilterOperator.LTE: "lte", + FilterOperator.IN: "in", + FilterOperator.NIN: "nin", + FilterOperator.ANY: "any", + FilterOperator.ALL: "all", + FilterOperator.CONTAINS: "contains", + FilterOperator.TEXT_MATCH: "text_match", + FilterOperator.TEXT_MATCH_INSENSITIVE: "text_match_insensitive", + FilterOperator.IS_EMPTY: "is_empty", + } + return mapping.get(op, "eq") + + +def _filters_to_zeusdb(filters: MetadataFilters | None) -> dict[str, Any] | None: + """ + Convert LlamaIndex MetadataFilters to ZeusDB flat format. + + ZeusDB expects flat dict with implicit AND: + {"key1": value, "key2": {"op": value}} + """ + if filters is None: + return None + + def _one(f: MetadataFilter | MetadataFilters) -> dict[str, Any]: + if isinstance(f, MetadataFilters): + cond = (f.condition or FilterCondition.AND).value.lower() + sub = [_one(sf) for sf in f.filters] + + if cond == "and": + # Merge into flat dict (implicit AND) + result = {} + for s in sub: + result.update(s) + return result + else: + # OR is NOT supported by Rust implementation + logger.warning( + "OR filters not supported by ZeusDB backend", + operation="filter_translation", + condition=cond, + ) + # Fallback: return first filter only + return sub[0] if sub else {} + + # Single filter + op_key = _translate_filter_op(f.operator) + + if op_key == "eq": + # Direct value for equality (matches Rust code) + return {f.key: f.value} + else: + # Operator wrapper for other ops + return {f.key: {op_key: f.value}} + + result = _one(filters) # Changed from 'z' to 'result' for consistency + + logger.debug("translated_filters", zeusdb_filter=result) + return result + + +# ------------------------- +# MMR helpers (opt-in only) +# ------------------------- + + +def _dot(a: list[float], b: list[float]) -> float: + return sum(x * y for x, y in zip(a, b)) + + +def _norm(a: list[float]) -> float: + return max(1e-12, sum(x * x for x in a) ** 0.5) + + +def _cosine_sim(a: list[float], b: list[float]) -> float: + return _dot(a, b) / (_norm(a) * _norm(b)) + + +def _mmr_select( + query_vec: list[float], + cand_vecs: list[list[float]], + k: int, + lamb: float, + precomputed_qs: list[float] | None = None, +) -> list[int]: + """ + Greedy Maximal Marginal Relevance. + Returns indices of the selected candidates. + """ + n = len(cand_vecs) + if n == 0 or k <= 0: + return [] + + rel = precomputed_qs or [_cosine_sim(query_vec, v) for v in cand_vecs] + selected: list[int] = [] + remaining = set(range(n)) + + # seed: most relevant + first = max(remaining, key=lambda i: rel[i]) + selected.append(first) + remaining.remove(first) + + while len(selected) < min(k, n) and remaining: + + def score(i: int) -> float: + # diversity term = max sim to any already-selected + max_div = max(_cosine_sim(cand_vecs[i], cand_vecs[j]) for j in selected) + return lamb * rel[i] - (1.0 - lamb) * max_div + + nxt = max(remaining, key=score) + selected.append(nxt) + remaining.remove(nxt) + + return selected + + +# ------------------------- +# ZeusDB Vector Store +# ------------------------- + + +class ZeusDBVectorStore(BasePydanticVectorStore): + """ + LlamaIndex VectorStore backed by ZeusDB (via the `zeusdb` umbrella package). + + Behaviors: + - Expects nodes with precomputed embeddings. + - Stores vectors + metadata; does not store full text (stores_text=False). + - Translates LlamaIndex MetadataFilters to ZeusDB filter dicts. + - Converts ZeusDB distances to similarity scores (higher = better). + - Supports opt-in MMR when the caller requests it. + - Provides async wrappers via thread offload. + + Persistence Note: Quantized indexes currently load in raw mode + """ + + stores_text: bool = False + is_embedding_query: bool = True + + def __init__( + self, + *, + dim: int | None = None, # Optional if using existing index + distance: str = "cosine", + index_type: str = "hnsw", + index_name: str = "default", + quantization_config: dict[str, Any] | None = None, + # ZeusDB tuning params (optional) + m: int | None = None, + ef_construction: int | None = None, + expected_size: int | None = None, + # Pre-existing ZeusDB index (optional) + zeusdb_index: Any | None = None, + # Extra kwargs forwarded to VectorDatabase.create() + **kwargs: Any, + ) -> None: + # super().__init__(stores_text=self.stores_text) + super().__init__(stores_text=False) # Use the literal value + + self._space = _infer_space(distance) + self._index_name = index_name + + if zeusdb_index is not None: + self._index = zeusdb_index + else: + if dim is None: + raise ValueError("dim is required when not providing zeusdb_index") + vdb = VectorDatabase() + create_kwargs: dict[str, Any] = { + "index_type": index_type, + "dim": dim, + "space": self._space, + } + if quantization_config is not None: + create_kwargs["quantization_config"] = quantization_config + if m is not None: + create_kwargs["m"] = m + if ef_construction is not None: + create_kwargs["ef_construction"] = ef_construction + if expected_size is not None: + create_kwargs["expected_size"] = expected_size + create_kwargs.update(kwargs) + with operation_context("create_index", space=self._space): + self._index = vdb.create(**create_kwargs) + + # ---- BasePydanticVectorStore API ---- + + @property + def client(self) -> Any: + return self._index + + def add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: + with operation_context( + "add_vectors", + requested=len(nodes), + overwrite=bool(kwargs.get("overwrite", True)), + ): + vectors: list[list[float]] = [] + metadatas: list[dict[str, Any]] = [] + ids: list[str] = [] + provided_count = 0 + + for n in nodes: + emb = _extract_embedding(n) + if emb is None: + continue + vectors.append(emb) + metadatas.append(_node_metadata(n)) + node_id = getattr(n, "id_", None) + if node_id is not None: + ids.append(str(node_id)) + provided_count += 1 + else: + ids.append("") # placeholder + + if not vectors: + logger.debug("add_vectors no-op (no embeddings)") + return [] + + payload: dict[str, Any] = { + "vectors": vectors, + "metadatas": metadatas, + } + + # All-or-nothing ID policy + if 0 < provided_count < len(ids): + logger.debug( + "partial_ids_ignored", + provided_count=provided_count, + total=len(ids), + ) + if provided_count == len(ids): + payload["ids"] = ids + + overwrite = bool(kwargs.get("overwrite", True)) + try: + result = self._index.add(payload, overwrite=overwrite) + except Exception as e: + logger.error( + "ZeusDB add operation failed", + operation="add_vectors", + node_count=len(nodes), + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + raise + + assigned_ids: list[str] = [] + if isinstance(result, dict) and "ids" in result: + assigned_ids = [str(x) for x in (result.get("ids") or [])] + elif hasattr(result, "ids"): + assigned_ids = [str(x) for x in (getattr(result, "ids") or [])] + + logger.debug( + "add_vectors summary", + requested=len(nodes), + inserted=len(assigned_ids) if assigned_ids else len(vectors), + had_all_ids=(provided_count == len(ids)), + ) + + # Return backend IDs if available; else fallback to provided ones + if assigned_ids: + return assigned_ids + return [i for i in ids if i] + + # ------------------------- + # Deletion & maintenance + # ------------------------- + + def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: + """ + Delete all nodes associated with a ref_doc_id. + + ⚠️ LIMITATION: This method is NOT SUPPORTED by ZeusDB's HNSW backend. + + The HNSW index only supports deletion by node ID via remove_point(). + There is no filter-based deletion or scalable way to find all node IDs + for a given ref_doc_id (list() doesn't work in QuantizedOnly mode). + + This method will raise NotImplementedError to be honest about the limitation. + + Alternative: Use delete_nodes(node_ids=[...]) if you have the node IDs. + """ + logger.error( + "delete() by ref_doc_id is not supported by ZeusDB HNSW backend", + operation="delete", + ref_doc_id=ref_doc_id, + ) + raise NotImplementedError( + "ZeusDB HNSW backend does not support deletion by ref_doc_id. " + "The backend only supports ID-based deletion via remove_point(). " + "Use delete_nodes(node_ids=[...]) instead if you have the node IDs." + ) + + def delete_nodes( + self, + node_ids: list[str] | None = None, + filters: MetadataFilters | None = None, + **delete_kwargs: Any, + ) -> None: + """ + Delete nodes by IDs. + + βœ… SUPPORTED: Deletion by explicit node IDs via remove_point(). + ❌ NOT SUPPORTED: Deletion by metadata filters. + + Args: + node_ids: List of node IDs to delete (supported) + filters: Metadata filters (NOT supported - will raise error if provided) + + Note: ZeusDB HNSW only supports direct ID-based deletion. + """ + if filters: + logger.error( + "delete_nodes() with filters is not supported by ZeusDB HNSW backend", + operation="delete_nodes", + has_filters=True, + ) + raise NotImplementedError( + "ZeusDB HNSW backend does not support filter-based deletion. " + "Only direct node ID deletion is supported." + ) + + if not node_ids: + logger.debug("delete_nodes called with no node_ids") + return + + with operation_context("delete_nodes", node_ids_count=len(node_ids)): + try: + success_count = 0 + failed_ids = [] + + for node_id in node_ids: + try: + result = self._index.remove_point(node_id) + if result: + success_count += 1 + else: + failed_ids.append(node_id) + except Exception as e: + failed_ids.append(node_id) + logger.warning( + "Failed to remove point", + operation="delete_nodes", + node_id=node_id, + error=str(e), + ) + + logger.info( + "Delete nodes completed", + operation="delete_nodes", + requested=len(node_ids), + deleted=success_count, + failed=len(failed_ids), + ) + + if failed_ids and len(failed_ids) < 10: + logger.debug( + "Failed node IDs", + operation="delete_nodes", + failed_ids=failed_ids, + ) + + except Exception as e: + logger.error( + "Delete nodes failed", + operation="delete_nodes", + node_ids_count=len(node_ids), + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + raise + + def clear(self) -> None: + """ + Clear all vectors from the index. + + ⚠️ LIMITATION: May not work correctly in QuantizedOnly mode. + + The clear() method may not properly clear quantized-only vectors. + """ + with operation_context("clear_index"): + try: + if hasattr(self._index, "clear"): + self._index.clear() + logger.info("Index cleared", operation="clear_index") + else: + logger.warning( + "clear() not available on index", + operation="clear_index", + ) + raise NotImplementedError( + "ZeusDB index does not expose clear() method" + ) + except Exception as e: + logger.error( + "Clear operation failed", + operation="clear_index", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + raise + + # ------------------------- + # Query (with optional MMR) + # ------------------------- + + def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult: + """ + Execute a vector search against ZeusDB. + + Kwargs understood by this adapter: + mmr (bool): Enable Maximal Marginal Relevance re-ranking. Default False. + mmr_lambda (float): Trade-off [0..1]. 1=relevance, 0=diversity. + Default 0.7 (or the provided `query.mmr_threshold`). + fetch_k (int): Candidate pool when MMR is on. Default max(20, 4*k). + ef_search (int): HNSW runtime search breadth; forwarded to ZeusDB. + return_vector (bool): Ask backend to return raw vectors. Auto-enabled + when MMR is requested. + auto_fallback (bool): If results < k with no filters, retry once with + a broader search. Default True. + """ + with operation_context("query", has_embedding=bool(query.query_embedding)): + if not query.query_embedding: + return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) + + # Detect explicit MMR requests + want_mmr = False + mode = getattr(query, "mode", None) + if mode and str(mode).lower().endswith("mmr"): + want_mmr = True + if kwargs.get("mmr", False): + want_mmr = True + mmr_threshold = getattr(query, "mmr_threshold", None) + if mmr_threshold is not None: + want_mmr = True + + # Common query prep + k = int(query.hybrid_top_k or query.similarity_top_k or 1) + zfilter = _filters_to_zeusdb(query.filters) + ef_search = kwargs.get("ef_search") + + fetch_k = k if not want_mmr else int(kwargs.get("fetch_k", max(20, 4 * k))) + fetch_k = max(fetch_k, k) + + default_lambda = 0.7 if mmr_threshold is None else mmr_threshold + mmr_lambda = float(kwargs.get("mmr_lambda", default_lambda)) + if mmr_lambda < 0.0: + mmr_lambda = 0.0 + elif mmr_lambda > 1.0: + mmr_lambda = 1.0 + + return_vector = bool(kwargs.get("return_vector", False) or want_mmr) + + logger.debug( + "query parameters", + k=k, + fetch_k=fetch_k, + mmr=want_mmr, + mmr_lambda=mmr_lambda if want_mmr else None, + has_filters=zfilter is not None, + ef_search=ef_search, + return_vector=return_vector, + ) + + search_kwargs: dict[str, Any] = { + "vector": list(query.query_embedding), + "top_k": fetch_k, + } + if zfilter is not None: + search_kwargs["filter"] = zfilter + if ef_search is not None: + search_kwargs["ef_search"] = ef_search + if return_vector: + search_kwargs["return_vector"] = True + + # Execute search with timing and error context + t0 = perf_counter() + try: + res = self._index.search(**search_kwargs) + except Exception as e: + logger.error( + "ZeusDB search failed", + operation="query", + error=str(e), + error_type=type(e).__name__, + params=search_kwargs, + exc_info=True, + ) + raise + search_ms = (perf_counter() - t0) * 1000 + + # Normalize hits + hits: list[dict[str, Any]] = [] + if isinstance(res, dict) and "results" in res: + hits = res.get("results") or [] + elif isinstance(res, list): + hits = res + + cand_ids: list[str] = [] + cand_dists: list[float] = [] + cand_vecs: list[list[float]] = [] + + for h in hits: + _id = str(h.get("id")) if "id" in h else str(h.get("node_id", "")) + cand_ids.append(_id) + + if "distance" in h: + cand_dists.append(float(h["distance"])) + elif "score" in h: + cand_dists.append(1.0 - float(h["score"])) + else: + cand_dists.append(1.0) + + if return_vector: + v = h.get("vector") + cand_vecs.append( + [float(x) for x in v] if isinstance(v, list) else [] + ) + + # Broadened search fallback (default on) + fallback_used = False + if len(cand_ids) < k and not zfilter and kwargs.get("auto_fallback", True): + logger.debug( + "broadening_search_retry", + initial_results=len(cand_ids), + requested_k=k, + ) + try: + broader_res = self._index.search( + vector=list(query.query_embedding), + top_k=max(k, fetch_k), + ef_search=max(500, max(k, fetch_k) * 10), + return_vector=return_vector, + ) + if isinstance(broader_res, dict) and "results" in broader_res: + broader_hits = broader_res.get("results") or [] + elif isinstance(broader_res, list): + broader_hits = broader_res + else: + broader_hits = [] + + if len(broader_hits) > len(hits): + hits = broader_hits + cand_ids, cand_dists, cand_vecs = [], [], [] + for h in hits: + _id = ( + str(h.get("id")) + if "id" in h + else str(h.get("node_id", "")) + ) + cand_ids.append(_id) + if "distance" in h: + cand_dists.append(float(h["distance"])) + elif "score" in h: + cand_dists.append(1.0 - float(h["score"])) + else: + cand_dists.append(1.0) + if return_vector: + v = h.get("vector") + cand_vecs.append( + [float(x) for x in v] if isinstance(v, list) else [] + ) + fallback_used = True + logger.info( + "broadened_search_applied", + gained=len(cand_ids), + ) + except Exception as e: + logger.debug( + "broadened_search_failed", + error=str(e), + error_type=type(e).__name__, + ) + + # Optional MMR rerank (opt-in only) + mmr_ms = 0.0 + if want_mmr: + if ( + cand_vecs + and all(cand_vecs) + and isinstance(query.query_embedding, list) + ): + t1 = perf_counter() + qv = list(query.query_embedding) + rel_q = [_cosine_sim(qv, v) for v in cand_vecs] + sel_idx = _mmr_select( + qv, + cand_vecs, + k=k, + lamb=mmr_lambda, + precomputed_qs=rel_q, + ) + mmr_ms = (perf_counter() - t1) * 1000 + sel_ids = [cand_ids[i] for i in sel_idx] + sel_sims = [rel_q[i] for i in sel_idx] + logger.info( + "mmr_rerank_applied", + selected=len(sel_ids), + fetch_k=fetch_k, + mmr_lambda=mmr_lambda, + search_ms=search_ms, + rerank_ms=mmr_ms, + space=self._space, + fallback_used=fallback_used, + ) + return VectorStoreQueryResult( + nodes=None, similarities=sel_sims, ids=sel_ids + ) + # If vectors missing, fall through to dense ranking + + # Default: dense similarity ranking + ids: list[str] = [] + sims: list[float] = [] + for _id, dist in zip(cand_ids, cand_dists): + ids.append(_id) + sims.append(_similarity_from_distance(dist, self._space)) + + logger.info( + "Query completed", + operation="query", + search_ms=search_ms, + mmr_ms=mmr_ms, + results_count=len(hits), + final_count=len(ids), + k=k, + mmr=want_mmr, + space=self._space, + fallback_used=fallback_used, + has_filters=zfilter is not None, + ) + return VectorStoreQueryResult( + nodes=None, + similarities=sims, + ids=ids, + ) + + def persist(self, persist_path: str, fs: Any | None = None) -> None: + with operation_context("persist_index", path=persist_path): + try: + if hasattr(self._index, "save"): + self._index.save(persist_path) # type: ignore[attr-defined] + except Exception as e: + logger.error( + "ZeusDB persist failed", + operation="persist_index", + path=persist_path, + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + raise + + # ------------------------- + # Async wrappers (thread offload) + # ------------------------- + + # async def aadd(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: + async def async_add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: + """Thread-offloaded async variant of add().""" + return await asyncio.to_thread(self.add, nodes, **kwargs) + + async def aquery( + self, query: VectorStoreQuery, **kwargs: Any + ) -> VectorStoreQueryResult: + """Thread-offloaded async variant of query().""" + return await asyncio.to_thread(self.query, query, **kwargs) + + async def adelete(self, ref_doc_id: str, **kwargs: Any) -> None: + """Thread-offloaded async variant of delete().""" + return await asyncio.to_thread(self.delete, ref_doc_id, **kwargs) + + async def adelete_nodes( + self, + node_ids: list[str] | None = None, + filters: MetadataFilters | None = None, + **kwargs: Any, + ) -> None: + """Thread-offloaded async variant of delete_nodes().""" + return await asyncio.to_thread(self.delete_nodes, node_ids, filters, **kwargs) + + async def aclear(self) -> None: + """Thread-offloaded async variant of clear().""" + return await asyncio.to_thread(self.clear) + + # ------------------------- + # Factory methods and convenience utilities + # ------------------------- + + @classmethod + def from_nodes( + cls, + nodes: list[BaseNode], + *, + dim: int | None = None, + distance: str = "cosine", + index_type: str = "hnsw", + **kwargs: Any, + ) -> ZeusDBVectorStore: + """Create ZeusDBVectorStore from nodes with embeddings.""" + if not nodes: + raise ValueError("Cannot create store from empty nodes list") + + # Infer dimension from first node if not provided + if dim is None: + first_emb = _extract_embedding(nodes[0]) + if first_emb is None: + raise ValueError("First node has no embedding to infer dimension") + dim = len(first_emb) + + store = cls( + dim=dim, + distance=distance, + index_type=index_type, + **kwargs, + ) + store.add(nodes) + return store + + @classmethod + def load_index( + cls, + path: str, + **kwargs: Any, + ) -> ZeusDBVectorStore: + """ + Load ZeusDB index from disk. + + Quantized indexes will load in raw mode. + The quantization model and training state are preserved, but quantized + search will not be active until the next ZeusDB release. + + The index will function correctly using raw vectors, + with full search accuracy but without memory compression benefits. + """ + with operation_context("load_index", path=path): + vdb = VectorDatabase() + zeusdb_index = vdb.load(path) + + store = cls(zeusdb_index=zeusdb_index, **kwargs) + + # Detect and warn about quantization state + try: + can_use = store.can_use_quantization() + is_active = store.is_quantized() + storage_mode = store.get_storage_mode() + + if can_use and not is_active: + logger.warning( + "Quantized index loaded in raw mode", + operation="load_index", + storage_mode=storage_mode, + can_use_quantization=can_use, + is_quantized=is_active, + ) + + quant_info = store.get_quantization_info() + if quant_info: + logger.info( + "Quantization config preserved but not active", + operation="load_index", + compression_ratio=quant_info.get( + "compression_ratio", "N/A" + ), + subvectors=quant_info.get("subvectors", "N/A"), + bits=quant_info.get("bits", "N/A"), + ) + + logger.info( + "Index will use raw vectors. Search accuracy preserved. " + "Memory compression unavailable until next release.", + operation="load_index", + ) + + except Exception as e: + logger.debug( + "Could not check quantization status", + operation="load_index", + error=str(e), + error_type=type(e).__name__, + ) + + return store + + def get_vector_count(self) -> int: + """Return total vectors in the index (best-effort).""" + try: + if hasattr(self._index, "get_vector_count"): + return int(self._index.get_vector_count()) # type: ignore + except Exception as e: + logger.error( + "get_vector_count failed", + error=str(e), + error_type=type(e).__name__, + ) + return 0 + + def get_zeusdb_stats(self) -> dict[str, Any]: + """Return ZeusDB stats (best-effort).""" + try: + if hasattr(self._index, "get_stats"): + stats = self._index.get_stats() # type: ignore + return dict(stats) if isinstance(stats, dict) else {} + except Exception as e: + logger.error( + "get_zeusdb_stats failed", + error=str(e), + error_type=type(e).__name__, + ) + return {} + + def save_index(self, path: str) -> bool: + """Save index to disk (best-effort wrapper).""" + try: + if hasattr(self._index, "save"): + self._index.save(path) # type: ignore[attr-defined] + return True + except Exception as e: + logger.error( + "save_index failed", + path=path, + error=str(e), + error_type=type(e).__name__, + ) + return False + + def info(self) -> str: + """ + Get a human-readable info string about the index. + + Example: + >>> print(vector_store.info()) + HNSWIndex(dim=1536, space=cosine, vectors=1200, quantized=True, ...) + """ + try: + info_str = self._index.info() + logger.debug( + "Retrieved index info", operation="info", info_length=len(info_str) + ) + return info_str + except Exception as e: + logger.error( + "Failed to get index info", + operation="info", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + return f"ZeusDBVectorStore(error: {e})" + + # ------------------------------------------------------------------------- + # Quantization Methods + # ------------------------------------------------------------------------- + + def get_training_progress(self) -> float: + """ + Get quantization training progress percentage. + + Returns: + float: Training progress as percentage (0.0 to 100.0). + Returns 0.0 if quantization is not configured or on error. + + Example: + >>> progress = vector_store.get_training_progress() + >>> print(f"Training: {progress:.1f}% complete") + """ + try: + progress = self._index.get_training_progress() + logger.debug( + "Retrieved training progress", + operation="get_training_progress", + progress_percent=progress, + ) + return progress + except Exception as e: + logger.error( + "Failed to get training progress", + operation="get_training_progress", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + return 0.0 + + def is_quantized(self) -> bool: + """ + Check whether quantized search is currently active. + + Returns: + bool: True if index is using quantized vectors for search, + False otherwise or on error. + + Example: + >>> if vector_store.is_quantized(): + ... print("Using quantized search") + """ + try: + quantized = self._index.is_quantized() + logger.debug( + "Retrieved quantization status", + operation="is_quantized", + is_quantized=quantized, + ) + return quantized + except Exception as e: + logger.error( + "Failed to check quantization status", + operation="is_quantized", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + return False + + def can_use_quantization(self) -> bool: + """ + Check whether quantization is available (e.g., PQ training completed). + + Returns: + bool: True if quantization is trained and ready to use, + False otherwise or on error. + + Example: + >>> if vector_store.can_use_quantization(): + ... print("Quantization ready") + """ + try: + available = self._index.can_use_quantization() + logger.debug( + "Retrieved quantization availability", + operation="can_use_quantization", + can_use_quantization=available, + ) + return available + except Exception as e: + logger.error( + "Failed to check quantization availability", + operation="can_use_quantization", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + return False + + def get_storage_mode(self) -> str: + """ + Get current storage mode. + + Returns: + str: Storage mode string. Possible values: + - 'raw_only': Only raw vectors stored + - 'quantized_only': Only quantized vectors (memory optimized) + - 'quantized_with_raw': Both quantized and raw vectors + - 'quantized_active': Quantization is active + - 'unknown': On error or unable to determine + + Example: + >>> mode = vector_store.get_storage_mode() + >>> print(f"Storage mode: {mode}") + """ + try: + mode = self._index.get_storage_mode() + logger.debug( + "Retrieved storage mode", + operation="get_storage_mode", + storage_mode=mode, + ) + return mode + except Exception as e: + logger.error( + "Failed to get storage mode", + operation="get_storage_mode", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + return "unknown" + + def get_quantization_info(self) -> dict[str, Any] | None: + """ + Get detailed quantization information. + + Returns: + Optional[Dict]: Dictionary containing quantization details: + - compression_ratio: Memory compression factor (e.g., 16.0 for 16x) + - memory_mb: Estimated memory usage in megabytes + - subvectors: Number of subvectors used + - bits: Bits per quantized code + - trained: Whether training is complete + - training_size: Number of vectors used for training + Returns None if quantization is not configured/trained or on error. + + Example: + >>> info = vector_store.get_quantization_info() + >>> if info: + ... print(f"Compression: {info['compression_ratio']:.1f}x") + ... print(f"Memory: {info['memory_mb']:.2f} MB") + """ + try: + info = self._index.get_quantization_info() + logger.debug( + "Retrieved quantization info", + operation="get_quantization_info", + has_quantization=info is not None, + ) + return info + except Exception as e: + logger.error( + "Failed to get quantization info", + operation="get_quantization_info", + error=str(e), + error_type=type(e).__name__, + exc_info=True, + ) + return None diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/py.typed b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml new file mode 100644 index 0000000000..0d053d4ed7 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml @@ -0,0 +1,127 @@ +[build-system] +requires = ["hatchling>=1.18"] +build-backend = "hatchling.build" + +[project] +name = "llama-index-vector-stores-zeusdb" +version = "0.1.5" +description = "LlamaIndex integration for ZeusDB vector database. Enterprise-grade RAG with high-performance vector search." +readme = "README.md" +requires-python = ">=3.10,<4.0" +license = { text = "MIT" } +authors = [{ name = "ZeusDB", email = "contact@zeusdb.com" }] +keywords = [ + "vector-database", + "embeddings", + "semantic-search", + "rag", + "llama-index", + "zeusdb", + "hnsw", +] +classifiers = [ + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Software Development :: Libraries :: Python Modules", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +# Runtime dependencies +dependencies = [ + "llama-index-core>=0.14.6", # Latest release as of Oct 26, 2025 + "zeusdb>=0.0.8", # Megapackage version +] + +[project.optional-dependencies] +# Install with: uv pip install -e ".[dev]" +dev = [ + "pytest>=8.4", + "pytest-asyncio>=1.2.0", + "pytest-mock>=3.15", + "ruff>=0.12", + "black>=25.9", + "mypy>=1.17", + "pre-commit>=4.3", + "codespell[toml]>=2.2.6", + # "ipython>=9.5", + "ipython>=8.10.0", #need support for python 3.10 + "jupyter>=1.0.0", + "pylint>=3.3", # Latest + # "tree-sitter-languages>=1.10.2", # Not compatible with 3.13 + "types-requests>=2.32", # Latest KEEP + # "types-setuptools>=80.9", # Latest + "llama-index-embeddings-openai>=0.5.1", + "llama-index-llms-openai>=0.6.6", + "python-dotenv>=1.1.1", # For examples and environment variable loading +] + +[project.urls] +Homepage = "https://github.com/zeusdb/llama-index-vector-stores-zeusdb" +Repository = "https://github.com/zeusdb/llama-index-vector-stores-zeusdb" +Issues = "https://github.com/zeusdb/llama-index-vector-stores-zeusdb/issues" + +# --- LlamaHub metadata (REQUIRED for the PR) --- +[tool.llamahub] +contains_example = true +import_path = "llama_index.vector_stores.zeusdb" + +[tool.llamahub.class_authors] +ZeusDBVectorStore = "zeusdb" +# ---------------------------------------------------------- + +[tool.codespell] +check-filenames = true +check-hidden = true +skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb" + +[tool.mypy] +disallow_untyped_defs = true +exclude = ["_static", "build", "examples", "notebooks", "venv", ".venv"] +ignore_missing_imports = true +python_version = "3.10" + +[tool.hatch.build.targets.wheel] +packages = ["llama_index"] +include = ["llama_index/vector_stores/zeusdb/py.typed"] + +[tool.hatch.build.targets.sdist] +include = [ + "llama_index", + "README.md", + "LICENSE", + "pyproject.toml", + "llama_index/vector_stores/zeusdb/py.typed", +] + +[tool.ruff] +line-length = 88 +target-version = "py310" + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "F", # pyflakes + "I", # isort + "UP", # pyupgrade +] +ignore = [] + +[tool.black] +line-length = 88 +target-version = ['py310', 'py311', 'py312'] +include = '\.pyi?$' + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = "test_*.py" +python_classes = "Test*" +python_functions = "test_*" +asyncio_mode = "auto" + +[tool.ruff.lint.isort] +known-first-party = ["llama_index"] +section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"] +force-sort-within-sections = true diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/__init__.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py new file mode 100644 index 0000000000..7bb8c8bcce --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py @@ -0,0 +1,408 @@ +# tests/test_vector_stores_zeusdb.py +from dataclasses import dataclass +import importlib +import math +import sys +import types + +import pytest + + +# ----------------------------- +# Minimal in-memory ZeusDB fake +# ----------------------------- +def _cosine_distance(a, b): + dot = sum(x * y for x, y in zip(a, b)) + na = math.sqrt(sum(x * x for x in a)) or 1e-12 + nb = math.sqrt(sum(x * x for x in b)) or 1e-12 + sim = dot / (na * nb) + return 1.0 - sim # adapter expects distance; later converts back to similarity + + +class _FakeIndex: + def __init__(self, dim, space="cosine", **kwargs): + self.dim = dim + self.space = space + self.records = [] # list of dict: {id, vector, metadata} + self._next_id = 1 + self._saved_path = None + + def add(self, payload, overwrite=True): + vectors = payload.get("vectors", []) + metadatas = payload.get("metadatas", []) + ids = payload.get("ids", None) + + assigned = [] + for i, v in enumerate(vectors): + meta = dict(metadatas[i]) if i < len(metadatas) else {} + if ids and ids[i]: + rid = str(ids[i]) + else: + rid = str(self._next_id) + self._next_id += 1 + # overwrite means replace if existing id matches + if overwrite: + self.records = [r for r in self.records if r["id"] != rid] + self.records.append({"id": rid, "vector": list(v), "metadata": meta}) + assigned.append(rid) + return {"ids": assigned} + + def remove_point(self, node_id): + """Remove a point by ID - ADDED FOR DELETE SUPPORT""" + initial_count = len(self.records) + self.records = [r for r in self.records if r["id"] != str(node_id)] + removed = initial_count > len(self.records) + return removed # Return True if something was removed + + def _match_filter(self, meta, zfilter): + """Match filter in FLAT format (like actual Rust code).""" + if not zfilter: + return True + + # Iterate over flat filter dict + for field, condition in zfilter.items(): + field_value = meta.get(field) + + # Direct value = equality check + if isinstance(condition, str | int | float | bool | type(None)): + if field_value != condition: + return False + # Operator dict + elif isinstance(condition, dict): + for op, target_value in condition.items(): + if op == "eq" and field_value != target_value: + return False + elif op == "ne" and field_value == target_value: + return False + elif op == "gt" and not ( + isinstance(field_value, int | float) + and field_value > target_value + ): + return False + elif op == "gte" and not ( + isinstance(field_value, int | float) + and field_value >= target_value + ): + return False + elif op == "lt" and not ( + isinstance(field_value, int | float) + and field_value < target_value + ): + return False + elif op == "lte" and not ( + isinstance(field_value, int | float) + and field_value <= target_value + ): + return False + elif op == "in" and field_value not in target_value: + return False + elif op == "contains": + if ( + isinstance(field_value, list) + and target_value not in field_value + ): + return False + elif ( + isinstance(field_value, str) + and str(target_value) not in field_value + ): + return False + else: + return False + + return True + + def search(self, vector, top_k=5, filter=None, ef_search=None, return_vector=False): + # filter records + recs = [r for r in self.records if self._match_filter(r["metadata"], filter)] + # score + scored = [] + for r in recs: + if self.space == "cosine": + dist = _cosine_distance(vector, r["vector"]) + else: + # simple L2 + dist = math.sqrt(sum((x - y) ** 2 for x, y in zip(vector, r["vector"]))) + res = {"id": r["id"], "distance": dist} + if return_vector: + res["vector"] = list(r["vector"]) + scored.append(res) + scored.sort(key=lambda x: x["distance"]) + return {"results": scored[:top_k]} + + def get_records(self, ids, return_vector=False): + """Get records by IDs - needed for node reconstruction.""" + if isinstance(ids, str): + ids = [ids] + + results = [] + id_set = set(str(i) for i in ids) + + for record in self.records: + if str(record["id"]) in id_set: + result = { + "id": record["id"], + "metadata": record["metadata"].copy(), + "score": 0.0, + } + if return_vector: + result["vector"] = record["vector"] + results.append(result) + + return results + + def clear(self): + self.records.clear() + + def save(self, path): + self._saved_path = path # track that save was called + + def get_vector_count(self): + return len(self.records) + + def get_stats(self): + return {"dim": self.dim, "space": self.space, "count": len(self.records)} + + +class _FakeVectorDatabase: + def create(self, index_type="hnsw", dim=None, space="cosine", **kwargs): + if dim is None: + raise ValueError("dim required") + return _FakeIndex(dim=dim, space=space, **kwargs) + + def load(self, path): + # In a "real" load we'd deserialize. For tests, hand back an empty index. + return _FakeIndex(dim=8, space="cosine") + + +@pytest.fixture(autouse=True) +def fake_zeusdb_module(monkeypatch): + """ + Install a fake `zeusdb` module *before* importing the adapter. + """ + mod = types.ModuleType("zeusdb") + mod.VectorDatabase = _FakeVectorDatabase # type: ignore[attr-defined] + # optional logging_config fallback (adapter handles absence) + sys.modules["zeusdb"] = mod + yield + sys.modules.pop("zeusdb", None) + + +# ------------------------------------- +# Import the adapter after faking ZeusDB +# ------------------------------------- +@pytest.fixture +def ZeusDBVectorStore(): + # reload our package to bind to the fake module + import llama_index.vector_stores.zeusdb.base as zeusdb_base + + importlib.reload(zeusdb_base) + return zeusdb_base.ZeusDBVectorStore + + +# -------------------- +# Minimal node helper +# -------------------- +@dataclass +class FakeNode: + text: str + embedding: list[float] + metadata: dict + id_: str | None = None + ref_doc_id: str | None = None + + +# ---------- +# Test cases +# ---------- +def test_add_and_query_basic(ZeusDBVectorStore): + store = ZeusDBVectorStore(dim=3, distance="cosine") + n1 = FakeNode("a", [1.0, 0.0, 0.0], {"category": "x"}, id_="1", ref_doc_id="docA") + n2 = FakeNode("b", [0.9, 0.1, 0.0], {"category": "y"}, id_="2", ref_doc_id="docA") + out_ids = store.add([n1, n2]) + assert set(out_ids) == {"1", "2"} + + q = [1.0, 0.0, 0.0] + from llama_index.core.vector_stores.types import VectorStoreQuery + + res = store.query(VectorStoreQuery(query_embedding=q, similarity_top_k=1)) + assert res.ids and res.similarities + assert res.ids[0] in {"1", "2"} + # nearest to [1,0,0] should be id "1" + assert res.ids[0] == "1" + + +def test_filters_and_delete_nodes(ZeusDBVectorStore): + store = ZeusDBVectorStore(dim=3) + nodes = [ + FakeNode("a", [1.0, 0.0, 0.0], {"category": "x"}, id_="1", ref_doc_id="docA"), + FakeNode("b", [0.0, 1.0, 0.0], {"category": "y"}, id_="2", ref_doc_id="docB"), + FakeNode("c", [0.0, 0.0, 1.0], {"category": "x"}, id_="3", ref_doc_id="docC"), + ] + store.add(nodes) + + from llama_index.core.vector_stores.types import ( + FilterCondition, + FilterOperator, + MetadataFilters, + VectorStoreQuery, + ) + + mf = MetadataFilters.from_dicts( + [ + {"key": "category", "value": "x", "operator": FilterOperator.EQ}, + ], + condition=FilterCondition.AND, + ) + res = store.query( + VectorStoreQuery( + query_embedding=[1.0, 0.0, 0.0], + similarity_top_k=10, + filters=mf, + ) + ) + assert set(res.ids) == {"1", "3"} # filter applied + + # Now delete node id "3" via delete_nodes + store.delete_nodes(node_ids=["3"]) + res2 = store.query( + VectorStoreQuery( + query_embedding=[1.0, 0.0, 0.0], + similarity_top_k=10, + filters=mf, + ) + ) + assert set(res2.ids) == {"1"} + + +def test_delete_by_ref_doc_id_not_supported(ZeusDBVectorStore): + """Test that delete by ref_doc_id raises NotImplementedError.""" + store = ZeusDBVectorStore(dim=3) + nodes = [ + FakeNode("a", [1.0, 0.0, 0.0], {"k": 1}, id_="1", ref_doc_id="docA"), + FakeNode("b", [0.0, 1.0, 0.0], {"k": 2}, id_="2", ref_doc_id="docA"), + FakeNode("c", [0.0, 0.0, 1.0], {"k": 3}, id_="3", ref_doc_id="docB"), + ] + store.add(nodes) + + # Should raise NotImplementedError + with pytest.raises(NotImplementedError) as exc_info: + store.delete("docA") + + assert "ref_doc_id" in str(exc_info.value).lower() + assert "remove_point" in str(exc_info.value).lower() + + +def test_delete_nodes_by_id_works(ZeusDBVectorStore): + """Test that delete_nodes by ID works correctly.""" + store = ZeusDBVectorStore(dim=3) + nodes = [ + FakeNode("a", [1.0, 0.0, 0.0], {"k": 1}, id_="1", ref_doc_id="docA"), + FakeNode("b", [0.0, 1.0, 0.0], {"k": 2}, id_="2", ref_doc_id="docA"), + FakeNode("c", [0.0, 0.0, 1.0], {"k": 3}, id_="3", ref_doc_id="docB"), + ] + store.add(nodes) + + assert store.get_vector_count() == 3 + + # Delete by node IDs + store.delete_nodes(node_ids=["1", "2"]) + + assert store.get_vector_count() == 1 + + from llama_index.core.vector_stores.types import VectorStoreQuery + + res = store.query( + VectorStoreQuery(query_embedding=[1.0, 0.0, 0.0], similarity_top_k=10) + ) + assert set(res.ids) == {"3"} # only docB remains + + +def test_clear(ZeusDBVectorStore): + store = ZeusDBVectorStore(dim=3) + store.add([FakeNode("a", [1.0, 0.0, 0.0], {}, id_="1", ref_doc_id="docA")]) + assert store.get_vector_count() == 1 + store.clear() + assert store.get_vector_count() == 0 + + +def test_persist_and_load(ZeusDBVectorStore, tmp_path): + store = ZeusDBVectorStore(dim=3) + store.add([FakeNode("a", [1.0, 0.0, 0.0], {}, id_="1", ref_doc_id="docA")]) + out = tmp_path / "index.zeusdb" + store.persist(str(out)) + + # load_index should construct via VectorDatabase.load(...) + loaded = ZeusDBVectorStore.load_index(str(out)) + assert loaded is not None + assert hasattr(loaded, "client") + + +def test_from_nodes_infers_dim(ZeusDBVectorStore): + nodes = [ + FakeNode("a", [1.0, 0.0, 0.0, 0.0], {}, id_="1", ref_doc_id="docA"), + FakeNode("b", [0.0, 1.0, 0.0, 0.0], {}, id_="2", ref_doc_id="docA"), + ] + store = ZeusDBVectorStore.from_nodes(nodes) + assert store.get_vector_count() == 2 + + +def test_getters(ZeusDBVectorStore): + store = ZeusDBVectorStore(dim=2) + store.add([FakeNode("a", [1.0, 0.0], {"tier": 1}, id_="1", ref_doc_id="d")]) + assert store.get_vector_count() == 1 + stats = store.get_zeusdb_stats() + assert "dim" in stats and "space" in stats and "count" in stats + + +@pytest.mark.asyncio +async def test_async_paths(ZeusDBVectorStore): + store = ZeusDBVectorStore(dim=3) + + # inherited BasePydanticVectorStore.async_add -> calls add() sync by default + ids = await store.async_add( + [FakeNode("a", [1.0, 0.0, 0.0], {}, id_="1", ref_doc_id="d")] + ) + assert ids == ["1"] + + q = [1.0, 0.0, 0.0] + from llama_index.core.vector_stores.types import VectorStoreQuery + + res = await store.aquery(VectorStoreQuery(query_embedding=q, similarity_top_k=1)) + assert res.ids and res.ids[0] == "1" + + # adelete by ref_doc_id should raise NotImplementedError + with pytest.raises(NotImplementedError): + await store.adelete(ref_doc_id="d") + + # But adelete_nodes by ID should work + await store.adelete_nodes(node_ids=["1"]) + res2 = await store.aquery(VectorStoreQuery(query_embedding=q, similarity_top_k=1)) + assert res2.ids == [] or res2.ids is None + + # clear via aclear should not error + await store.aclear() + + +def test_query_with_mmr(ZeusDBVectorStore): + store = ZeusDBVectorStore(dim=3) + # three "near-ish" points around e1 + nodes = [ + FakeNode("n1", [1.0, 0.0, 0.0], {}, id_="1", ref_doc_id="d"), + FakeNode("n2", [0.9, 0.1, 0.0], {}, id_="2", ref_doc_id="d"), + FakeNode("n3", [0.8, 0.2, 0.0], {}, id_="3", ref_doc_id="d"), + ] + store.add(nodes) + + from llama_index.core.vector_stores.types import VectorStoreQuery + + # request k=2 with MMR enabled; fetch_k larger to force rerank variety + res = store.query( + VectorStoreQuery(query_embedding=[1.0, 0.0, 0.0], similarity_top_k=2), + mmr=True, + fetch_k=5, + mmr_lambda=0.7, + return_vector=True, + ) + assert len(res.ids) == 2 + assert set(res.ids).issubset({"1", "2", "3"}) From 398556fa9d2d2d38d8ee29ca248767488d9999f9 Mon Sep 17 00:00:00 2001 From: doubleinfinity <6169958+doubleinfinity@users.noreply.github.com> Date: Fri, 31 Oct 2025 23:18:54 +1100 Subject: [PATCH 2/6] fix(vector-stores): linting and formatting issues --- .../vector_stores/ZeusDBIndexDemo.ipynb | 89 ++++++++----------- .../CHANGELOG.md | 10 +-- .../README.md | 55 +++++++++--- .../examples/async_examples.py | 6 -- .../llama_index/vector_stores/zeusdb/base.py | 1 - .../pyproject.toml | 16 ++-- 6 files changed, 94 insertions(+), 83 deletions(-) diff --git a/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb b/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb index 8270d2aa2d..77cd40312f 100644 --- a/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb +++ b/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb @@ -183,7 +183,7 @@ "import os\n", "import getpass\n", "\n", - "os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')" + "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")" ] }, { @@ -290,9 +290,7 @@ "source": [ "# Create ZeusDB vector store\n", "vector_store = ZeusDBVectorStore(\n", - " dim=1536, # OpenAI embedding dimension\n", - " distance=\"cosine\",\n", - " index_type=\"hnsw\"\n", + " dim=1536, distance=\"cosine\", index_type=\"hnsw\" # OpenAI embedding dimension\n", ")\n", "\n", "# Create storage context\n", @@ -302,14 +300,11 @@ "documents = [\n", " Document(text=\"ZeusDB is a high-performance vector database.\"),\n", " Document(text=\"LlamaIndex provides RAG capabilities.\"),\n", - " Document(text=\"Vector search enables semantic similarity.\")\n", + " Document(text=\"Vector search enables semantic similarity.\"),\n", "]\n", "\n", "# Create index and store documents\n", - "index = VectorStoreIndex.from_documents(\n", - " documents,\n", - " storage_context=storage_context\n", - ")\n", + "index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)\n", "\n", "# Query the index\n", "query_engine = index.as_query_engine()\n", @@ -346,10 +341,7 @@ "embed_model = Settings.embed_model\n", "query_embedding = embed_model.get_text_embedding(\"machine learning\")\n", "\n", - "query_obj = VectorStoreQuery(\n", - " query_embedding=query_embedding,\n", - " similarity_top_k=2\n", - ")\n", + "query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=2)\n", "\n", "# Execute query\n", "results = vector_store.query(query_obj)\n", @@ -388,7 +380,7 @@ " query_obj,\n", " mmr=True,\n", " fetch_k=10,\n", - " mmr_lambda=0.7 # 0.0=max diversity, 1.0=pure relevance\n", + " mmr_lambda=0.7, # 0.0=max diversity, 1.0=pure relevance\n", ")\n", "\n", "print(f\"MMR Results: {len(mmr_results.ids or [])} items (with diversity)\")" @@ -420,7 +412,7 @@ "from llama_index.core.vector_stores.types import (\n", " MetadataFilters,\n", " FilterOperator,\n", - " FilterCondition\n", + " FilterCondition,\n", ")\n", "\n", "# Create a fresh vector store for this example\n", @@ -431,29 +423,31 @@ "documents_with_meta = [\n", " Document(\n", " text=\"Python is great for data science\",\n", - " metadata={\"category\": \"tech\", \"year\": 2024}\n", + " metadata={\"category\": \"tech\", \"year\": 2024},\n", " ),\n", " Document(\n", " text=\"JavaScript is for web development\",\n", - " metadata={\"category\": \"tech\", \"year\": 2023}\n", + " metadata={\"category\": \"tech\", \"year\": 2023},\n", " ),\n", " Document(\n", " text=\"Climate change impacts ecosystems\",\n", - " metadata={\"category\": \"environment\", \"year\": 2024}\n", + " metadata={\"category\": \"environment\", \"year\": 2024},\n", " ),\n", "]\n", "\n", "# Build index with metadata\n", "index = VectorStoreIndex.from_documents(\n", - " documents_with_meta,\n", - " storage_context=storage_context\n", + " documents_with_meta, storage_context=storage_context\n", ")\n", "\n", "# Create metadata filter\n", - "filters = MetadataFilters.from_dicts([\n", - " {\"key\": \"category\", \"value\": \"tech\", \"operator\": FilterOperator.EQ},\n", - " {\"key\": \"year\", \"value\": 2024, \"operator\": FilterOperator.GTE}\n", - "], condition=FilterCondition.AND)\n", + "filters = MetadataFilters.from_dicts(\n", + " [\n", + " {\"key\": \"category\", \"value\": \"tech\", \"operator\": FilterOperator.EQ},\n", + " {\"key\": \"year\", \"value\": 2024, \"operator\": FilterOperator.GTE},\n", + " ],\n", + " condition=FilterCondition.AND,\n", + ")\n", "\n", "# Use the retriever with filters (recommended approach)\n", "retriever = index.as_retriever(similarity_top_k=5, filters=filters)\n", @@ -533,18 +527,18 @@ "source": [ "# Create quantized vector store for memory efficiency\n", "quantization_config = {\n", - " 'type': 'pq',\n", - " 'subvectors': 8,\n", - " 'bits': 8,\n", - " 'training_size': 1000,\n", - " 'storage_mode': 'quantized_only'\n", + " \"type\": \"pq\",\n", + " \"subvectors\": 8,\n", + " \"bits\": 8,\n", + " \"training_size\": 1000,\n", + " \"storage_mode\": \"quantized_only\",\n", "}\n", "\n", "vector_store = ZeusDBVectorStore(\n", " dim=1536,\n", " distance=\"cosine\",\n", " index_type=\"hnsw\",\n", - " quantization_config=quantization_config\n", + " quantization_config=quantization_config,\n", ")\n", "\n", "# Check quantization status\n", @@ -585,16 +579,10 @@ "delete_sc = StorageContext.from_defaults(vector_store=delete_vs)\n", "\n", "# Create documents\n", - "delete_docs = [\n", - " Document(text=f\"Document {i}\", metadata={\"doc_id\": i})\n", - " for i in range(5)\n", - "]\n", + "delete_docs = [Document(text=f\"Document {i}\", metadata={\"doc_id\": i}) for i in range(5)]\n", "\n", "# Build index\n", - "delete_index = VectorStoreIndex.from_documents(\n", - " delete_docs,\n", - " storage_context=delete_sc\n", - ")\n", + "delete_index = VectorStoreIndex.from_documents(delete_docs, storage_context=delete_sc)\n", "\n", "print(f\"Before delete: {delete_vs.get_vector_count()} vectors\")\n", "\n", @@ -606,7 +594,7 @@ " # Extract node IDs from results\n", " node_ids_to_delete = [result.node.node_id for result in results[:2]]\n", " print(f\"Deleting node IDs: {node_ids_to_delete[0][:8]}...\")\n", - " \n", + "\n", " # Delete by node IDs\n", " delete_vs.delete_nodes(node_ids=node_ids_to_delete)\n", " print(f\"After delete: {delete_vs.get_vector_count()} vectors\")\n", @@ -650,40 +638,37 @@ "# In Jupyter, use nest_asyncio to handle event loops\n", "try:\n", " import nest_asyncio\n", + "\n", " nest_asyncio.apply()\n", "except ImportError:\n", " pass\n", "\n", + "\n", "async def async_operations():\n", " # Create nodes\n", - " nodes = [\n", - " TextNode(text=f\"Document {i}\", metadata={\"doc_id\": i})\n", - " for i in range(10)\n", - " ]\n", - " \n", + " nodes = [TextNode(text=f\"Document {i}\", metadata={\"doc_id\": i}) for i in range(10)]\n", + "\n", " # Generate embeddings (required before adding)\n", " embed_model = Settings.embed_model\n", " for node in nodes:\n", " node.embedding = embed_model.get_text_embedding(node.text)\n", - " \n", + "\n", " # Add nodes asynchronously\n", " node_ids = await vector_store.async_add(nodes)\n", " print(f\"Added {len(node_ids)} nodes\")\n", - " \n", + "\n", " # Query asynchronously\n", " query_embedding = embed_model.get_text_embedding(\"document\")\n", - " query_obj = VectorStoreQuery(\n", - " query_embedding=query_embedding,\n", - " similarity_top_k=3\n", - " )\n", - " \n", + " query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=3)\n", + "\n", " results = await vector_store.aquery(query_obj)\n", " print(f\"Found {len(results.ids or [])} results\")\n", - " \n", + "\n", " # Delete asynchronously\n", " await vector_store.adelete_nodes(node_ids=node_ids[:2])\n", " print(f\"Deleted 2 nodes, {vector_store.get_vector_count()} remaining\")\n", "\n", + "\n", "# Run async function\n", "await async_operations() # In Jupyter\n", "# asyncio.run(async_operations()) # In regular Python scripts" diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md index 000063c69b..76d3634fc6 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md @@ -26,9 +26,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed -- **Async API naming**: Renamed `aadd()` to `async_add()` to align with LlamaIndex standard async method naming conventions. The adapter now uses the official LlamaIndex async API (`async_add`, `aquery`, `adelete`, `adelete_nodes`, `aclear`) for consistency across the ecosystem. +- **Async API naming**: Renamed `add()` to `async_add()` to align with LlamaIndex standard async method naming conventions. The adapter now uses the official LlamaIndex async API (`async_add`, `aquery`, `adelete`, `adelete_nodes`, `aclear`) for consistency across the ecosystem. - Updated async examples to use standard LlamaIndex async method names instead of custom aliases. -- Updated documentation strings to reference `async_add()` instead of `aadd()`. +- Updated documentation strings to reference `async_add()` instead of `add()`. --- @@ -38,7 +38,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Delete examples - Added `examples/delete_examples.py` demonstrating supported deletion operations and workarounds. - Test coverage - Added tests for ID-based deletion and proper error handling for unsupported operations. -- Addeed Product **Quantization (PQ) support** - Full support for memory-efficient vector compression with automatic training +- Added Product **Quantization (PQ) support** - Full support for memory-efficient vector compression with automatic training - **Persistence Support**: Complete save/load functionality for ZeusDB indexes - Save indexes to disk with `save_index(path)` - Load indexes from disk with `load_index(path)` @@ -47,8 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Cross-platform compatibility for sharing indexes between systems - Added comprehensive persistence examples (`examples/persistence_examples.py`) - **Async Support**: Full asynchronous operation support for non-blocking workflows - - `aadd()` - Add nodes asynchronously - - `aquery()` - Query asynchronously + - `add()` - Add nodes asynchronously + - `aquery()` - Query asynchronously - `adelete_nodes()` - Delete nodes by IDs asynchronously - Thread-offloaded async wrappers using `asyncio.to_thread()` - **MMR (Maximal Marginal Relevance) Search**: Diversity-focused retrieval for comprehensive results diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md index 8eb7977b4c..7dc74d6b85 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md @@ -9,7 +9,7 @@ ZeusDB vector database integration for LlamaIndex. Connect LlamaIndex's RAG fram - **Advanced Filtering**: Comprehensive metadata filtering with complex operators - **MMR Support**: Maximal Marginal Relevance for diverse, non-redundant results - **Quantization**: Product Quantization (PQ) for memory-efficient vector storage -- **Async Support**: Async wrappers for non-blocking operations (`aadd`, `aquery`, `adelete_nodes`) +- **Async Support**: Async methods for non-blocking operations (`async_add`, `aquery`, `adelete_nodes`) ## Installation @@ -117,15 +117,48 @@ Non-blocking operations for web servers and concurrent workflows: ```python import asyncio - -# Async add -node_ids = await vector_store.aadd(nodes) - -# Async query -results = await vector_store.aquery(query_obj) - -# Async delete -await vector_store.adelete_nodes(node_ids=["id1", "id2"]) +from llama_index.core.schema import TextNode + +# In Jupyter, use nest_asyncio to handle event loops +try: + import nest_asyncio + nest_asyncio.apply() +except ImportError: + pass + +async def async_operations(): + # Create nodes + nodes = [ + TextNode(text=f"Document {i}", metadata={"doc_id": i}) + for i in range(10) + ] + + # Generate embeddings (required before adding) + embed_model = Settings.embed_model + for node in nodes: + node.embedding = embed_model.get_text_embedding(node.text) + + # Add nodes asynchronously + node_ids = await vector_store.async_add(nodes) + print(f"Added {len(node_ids)} nodes") + + # Query asynchronously + query_embedding = embed_model.get_text_embedding("document") + query_obj = VectorStoreQuery( + query_embedding=query_embedding, + similarity_top_k=3 + ) + + results = await vector_store.aquery(query_obj) + print(f"Found {len(results.ids or [])} results") + + # Delete asynchronously + await vector_store.adelete_nodes(node_ids=node_ids[:2]) + print(f"Deleted 2 nodes, {vector_store.get_vector_count()} remaining") + +# Run async function +await async_operations() # In Jupyter +# asyncio.run(async_operations()) # In regular Python scripts ``` ### Metadata Filtering @@ -155,7 +188,7 @@ results = vector_store.query( ) ``` -**Supported operators**: EQ, NE, GT, GTE, LT, LTE, IN, NIN, ANY, ALL, CONTAINS, TEXT_MATCH, TEXT_MATCH_INSENSITIVE +**Supported operators**: EQ, NE, GT, GTE, LT, LTE, IN, ANY, ALL, CONTAINS, TEXT_MATCH, TEXT_MATCH_INSENSITIVE ## Configuration diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py index 4cd8f0e2dd..94a58de98a 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/examples/async_examples.py @@ -66,7 +66,6 @@ async def example_basic_async(): # Async add print(f"Adding {len(nodes)} nodes asynchronously...") start = time.perf_counter() - # node_ids = await vector_store.aadd(nodes) node_ids = await vector_store.async_add(nodes) add_time = (time.perf_counter() - start) * 1000 print(f" βœ… Added {len(node_ids)} nodes in {add_time:.2f}ms") @@ -137,7 +136,6 @@ async def example_concurrent_queries(): node.embedding = embed_model.get_text_embedding(text) nodes.append(node) - # await vector_store.aadd(nodes) await vector_store.async_add(nodes) print(f" βœ… Added {len(nodes)} documents") print() @@ -240,7 +238,6 @@ async def process_batch(batch: list[str], batch_id: int): node.embedding = embed_model.get_text_embedding(text) nodes.append(node) - # node_ids = await vector_store.aadd(nodes) node_ids = await vector_store.async_add(nodes) elapsed = (time.perf_counter() - start) * 1000 @@ -286,7 +283,6 @@ async def example_error_handling(): node.embedding = embed_model.get_text_embedding(node.text) nodes.append(node) - # node_ids = await vector_store.aadd(nodes) node_ids = await vector_store.async_add(nodes) print(f"Added {len(node_ids)} documents") print() @@ -372,7 +368,6 @@ async def example_timeouts(): print("Adding documents with timeout...") try: node_ids = await asyncio.wait_for( - # vector_store.aadd(nodes), vector_store.async_add(nodes), timeout=10.0, # 10 second timeout ) @@ -436,7 +431,6 @@ async def main(): print(" β€’ Quick prototypes and experiments") print() print("πŸ“š Available async methods:") - # print(" β€’ aadd(nodes) - Add nodes asynchronously") print(" β€’ async_add(nodes) - Add nodes asynchronously") print(" β€’ aquery(query) - Query asynchronously") print(" β€’ adelete_nodes(node_ids) - Delete by IDs asynchronously") diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py index 3fa6dc57e1..eb2c4e3da0 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py @@ -817,7 +817,6 @@ def persist(self, persist_path: str, fs: Any | None = None) -> None: # Async wrappers (thread offload) # ------------------------- - # async def aadd(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: async def async_add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: """Thread-offloaded async variant of add().""" return await asyncio.to_thread(self.add, nodes, **kwargs) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml index 0d053d4ed7..e75c0237e1 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml @@ -38,15 +38,15 @@ dependencies = [ [project.optional-dependencies] # Install with: uv pip install -e ".[dev]" dev = [ - "pytest>=8.4", - "pytest-asyncio>=1.2.0", - "pytest-mock>=3.15", - "ruff>=0.12", - "black>=25.9", - "mypy>=1.17", + "pytest>=8.4", + "pytest-asyncio>=1.2.0", + "pytest-mock>=3.15", + "ruff>=0.12", + "black>=25.9", + "mypy>=1.17", "pre-commit>=4.3", "codespell[toml]>=2.2.6", - # "ipython>=9.5", + # "ipython>=9.5", "ipython>=8.10.0", #need support for python 3.10 "jupyter>=1.0.0", "pylint>=3.3", # Latest @@ -55,7 +55,7 @@ dev = [ # "types-setuptools>=80.9", # Latest "llama-index-embeddings-openai>=0.5.1", "llama-index-llms-openai>=0.6.6", - "python-dotenv>=1.1.1", # For examples and environment variable loading + "python-dotenv>=1.1.1", # For examples and environment variable loading ] [project.urls] From c4f7842bc41c8cf5634aee082691c77e61d997dc Mon Sep 17 00:00:00 2001 From: doubleinfinity <6169958+doubleinfinity@users.noreply.github.com> Date: Tue, 4 Nov 2025 22:47:22 +1100 Subject: [PATCH 3/6] fix(vector-stores): apply pre-commit linting fixes --- .../vector_stores/ZeusDBIndexDemo.ipynb | 36 ++++++--- .../community/integrations/vector_stores.md | 7 +- .../CHANGELOG.md | 11 ++- .../README.md | 72 +++++++++--------- .../pyproject.toml | 76 +++++++++---------- 5 files changed, 112 insertions(+), 90 deletions(-) diff --git a/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb b/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb index 77cd40312f..ead6298e75 100644 --- a/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb +++ b/docs/examples/vector_stores/ZeusDBIndexDemo.ipynb @@ -290,7 +290,9 @@ "source": [ "# Create ZeusDB vector store\n", "vector_store = ZeusDBVectorStore(\n", - " dim=1536, distance=\"cosine\", index_type=\"hnsw\" # OpenAI embedding dimension\n", + " dim=1536,\n", + " distance=\"cosine\",\n", + " index_type=\"hnsw\", # OpenAI embedding dimension\n", ")\n", "\n", "# Create storage context\n", @@ -304,7 +306,9 @@ "]\n", "\n", "# Create index and store documents\n", - "index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)\n", + "index = VectorStoreIndex.from_documents(\n", + " documents, storage_context=storage_context\n", + ")\n", "\n", "# Query the index\n", "query_engine = index.as_query_engine()\n", @@ -341,7 +345,9 @@ "embed_model = Settings.embed_model\n", "query_embedding = embed_model.get_text_embedding(\"machine learning\")\n", "\n", - "query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=2)\n", + "query_obj = VectorStoreQuery(\n", + " query_embedding=query_embedding, similarity_top_k=2\n", + ")\n", "\n", "# Execute query\n", "results = vector_store.query(query_obj)\n", @@ -579,10 +585,14 @@ "delete_sc = StorageContext.from_defaults(vector_store=delete_vs)\n", "\n", "# Create documents\n", - "delete_docs = [Document(text=f\"Document {i}\", metadata={\"doc_id\": i}) for i in range(5)]\n", + "delete_docs = [\n", + " Document(text=f\"Document {i}\", metadata={\"doc_id\": i}) for i in range(5)\n", + "]\n", "\n", "# Build index\n", - "delete_index = VectorStoreIndex.from_documents(delete_docs, storage_context=delete_sc)\n", + "delete_index = VectorStoreIndex.from_documents(\n", + " delete_docs, storage_context=delete_sc\n", + ")\n", "\n", "print(f\"Before delete: {delete_vs.get_vector_count()} vectors\")\n", "\n", @@ -646,7 +656,10 @@ "\n", "async def async_operations():\n", " # Create nodes\n", - " nodes = [TextNode(text=f\"Document {i}\", metadata={\"doc_id\": i}) for i in range(10)]\n", + " nodes = [\n", + " TextNode(text=f\"Document {i}\", metadata={\"doc_id\": i})\n", + " for i in range(10)\n", + " ]\n", "\n", " # Generate embeddings (required before adding)\n", " embed_model = Settings.embed_model\n", @@ -659,7 +672,9 @@ "\n", " # Query asynchronously\n", " query_embedding = embed_model.get_text_embedding(\"document\")\n", - " query_obj = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=3)\n", + " query_obj = VectorStoreQuery(\n", + " query_embedding=query_embedding, similarity_top_k=3\n", + " )\n", "\n", " results = await vector_store.aquery(query_obj)\n", " print(f\"Found {len(results.ids or [])} results\")\n", @@ -699,7 +714,9 @@ "source": [ "# Get index statistics\n", "stats = vector_store.get_zeusdb_stats()\n", - "print(f\"Key stats: vectors={stats.get('total_vectors')}, space={stats.get('space')}\")\n", + "print(\n", + " f\"Key stats: vectors={stats.get('total_vectors')}, space={stats.get('space')}\"\n", + ")\n", "\n", "# Get vector count\n", "count = vector_store.get_vector_count()\n", @@ -735,8 +752,7 @@ "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.13.3" + "pygments_lexer": "ipython3" } }, "nbformat": 4, diff --git a/docs/src/content/docs/framework/community/integrations/vector_stores.md b/docs/src/content/docs/framework/community/integrations/vector_stores.md index 05602f3fe4..c60d6f60d1 100644 --- a/docs/src/content/docs/framework/community/integrations/vector_stores.md +++ b/docs/src/content/docs/framework/community/integrations/vector_stores.md @@ -976,7 +976,7 @@ Settings.llm = OpenAI(model="gpt-5") vector_store = ZeusDBVectorStore( dim=1536, # OpenAI embedding dimension distance="cosine", - index_type="hnsw" + index_type="hnsw", ) # Create storage context @@ -986,13 +986,12 @@ storage_context = StorageContext.from_defaults(vector_store=vector_store) documents = [ Document(text="ZeusDB is a high-performance vector database."), Document(text="LlamaIndex provides RAG capabilities."), - Document(text="Vector search enables semantic similarity.") + Document(text="Vector search enables semantic similarity."), ] # Create index and store documents index = VectorStoreIndex.from_documents( - documents, - storage_context=storage_context + documents, storage_context=storage_context ) # Query the index diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md index 76d3634fc6..c3c11181ee 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/CHANGELOG.md @@ -1,4 +1,5 @@ + # Changelog All notable changes to this project will be documented in this file. @@ -61,12 +62,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed -- Filter format alignment β€” Updated _filters_to_zeusdb() to produce a flat dictionary with implicit AND (e.g., { "key": value, "other": { "op": value } }) instead of a nested {"and": [...]} structure, matching the Rust implementation. -- Test infrastructure β€” Updated _match_filter() in the test fake from the nested format to the flat format to reflect production behavior. +- Filter format alignment β€” Updated \_filters_to_zeusdb() to produce a flat dictionary with implicit AND (e.g., { "key": value, "other": { "op": value } }) instead of a nested {"and": [...]} structure, matching the Rust implementation. +- Test infrastructure β€” Updated \_match_filter() in the test fake from the nested format to the flat format to reflect production behavior. ### Fixed -- Filter translation for metadata queries - _filters_to_zeusdb() now emits the flat format expected by the ZeusDB backend. Single filters and AND combinations are handled correctly. The previous nested format could cause filtered queries to return zero results. +- Filter translation for metadata queries - \_filters_to_zeusdb() now emits the flat format expected by the ZeusDB backend. Single filters and AND combinations are handled correctly. The previous nested format could cause filtered queries to return zero results. - Deletion behavior - Correctly implemented ID-based deletion using `remove_point()`. Delete operations now properly remove vectors from the index and update vector counts. --- @@ -86,15 +87,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added + ### Changed + ### Fixed + ### Removed + --- diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md index 7dc74d6b85..3b100c1c48 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/README.md @@ -34,7 +34,7 @@ Settings.llm = OpenAI(model="gpt-5") vector_store = ZeusDBVectorStore( dim=1536, # OpenAI embedding dimension distance="cosine", - index_type="hnsw" + index_type="hnsw", ) # Create storage context @@ -44,13 +44,12 @@ storage_context = StorageContext.from_defaults(vector_store=vector_store) documents = [ Document(text="ZeusDB is a high-performance vector database."), Document(text="LlamaIndex provides RAG capabilities."), - Document(text="Vector search enables semantic similarity.") + Document(text="Vector search enables semantic similarity."), ] # Create index and store documents index = VectorStoreIndex.from_documents( - documents, - storage_context=storage_context + documents, storage_context=storage_context ) # Query the index @@ -86,7 +85,7 @@ results = vector_store.query( VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=5), mmr=True, fetch_k=20, - mmr_lambda=0.7 # 0.0=max diversity, 1.0=pure relevance + mmr_lambda=0.7, # 0.0=max diversity, 1.0=pure relevance ) # Note: MMR automatically enables return_vector=True for diversity calculation @@ -102,12 +101,12 @@ vector_store = ZeusDBVectorStore( dim=1536, distance="cosine", quantization_config={ - 'type': 'pq', - 'subvectors': 8, - 'bits': 8, - 'training_size': 1000, - 'storage_mode': 'quantized_only' - } + "type": "pq", + "subvectors": 8, + "bits": 8, + "training_size": 1000, + "storage_mode": "quantized_only", + }, ) ``` @@ -122,40 +121,42 @@ from llama_index.core.schema import TextNode # In Jupyter, use nest_asyncio to handle event loops try: import nest_asyncio + nest_asyncio.apply() except ImportError: pass + async def async_operations(): # Create nodes nodes = [ TextNode(text=f"Document {i}", metadata={"doc_id": i}) for i in range(10) ] - + # Generate embeddings (required before adding) embed_model = Settings.embed_model for node in nodes: node.embedding = embed_model.get_text_embedding(node.text) - + # Add nodes asynchronously node_ids = await vector_store.async_add(nodes) print(f"Added {len(node_ids)} nodes") - + # Query asynchronously query_embedding = embed_model.get_text_embedding("document") query_obj = VectorStoreQuery( - query_embedding=query_embedding, - similarity_top_k=3 + query_embedding=query_embedding, similarity_top_k=3 ) - + results = await vector_store.aquery(query_obj) print(f"Found {len(results.ids or [])} results") - + # Delete asynchronously await vector_store.adelete_nodes(node_ids=node_ids[:2]) print(f"Deleted 2 nodes, {vector_store.get_vector_count()} remaining") + # Run async function await async_operations() # In Jupyter # asyncio.run(async_operations()) # In regular Python scripts @@ -169,21 +170,22 @@ Filter results by metadata: from llama_index.core.vector_stores.types import ( MetadataFilters, FilterOperator, - FilterCondition + FilterCondition, ) # Create metadata filter -filters = MetadataFilters.from_dicts([ - {"key": "category", "value": "tech", "operator": FilterOperator.EQ}, - {"key": "year", "value": 2024, "operator": FilterOperator.GTE} -], condition=FilterCondition.AND) +filters = MetadataFilters.from_dicts( + [ + {"key": "category", "value": "tech", "operator": FilterOperator.EQ}, + {"key": "year", "value": 2024, "operator": FilterOperator.GTE}, + ], + condition=FilterCondition.AND, +) # Query with filters results = vector_store.query( VectorStoreQuery( - query_embedding=query_embedding, - similarity_top_k=5, - filters=filters + query_embedding=query_embedding, similarity_top_k=5, filters=filters ) ) ``` @@ -192,15 +194,15 @@ results = vector_store.query( ## Configuration -| Parameter | Description | Default | -|-----------|-------------|---------| -| `dim` | Vector dimension | Required | -| `distance` | Distance metric (`cosine`, `l2`, `l1`) | `cosine` | -| `index_type` | Index type (`hnsw`) | `hnsw` | -| `m` | HNSW connectivity parameter | 16 | -| `ef_construction` | HNSW build-time search depth | 200 | -| `expected_size` | Expected number of vectors | 10000 | -| `quantization_config` | PQ quantization settings | None | +| Parameter | Description | Default | +| --------------------- | -------------------------------------- | -------- | +| `dim` | Vector dimension | Required | +| `distance` | Distance metric (`cosine`, `l2`, `l1`) | `cosine` | +| `index_type` | Index type (`hnsw`) | `hnsw` | +| `m` | HNSW connectivity parameter | 16 | +| `ef_construction` | HNSW build-time search depth | 200 | +| `expected_size` | Expected number of vectors | 10000 | +| `quantization_config` | PQ quantization settings | None | ## License diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml index e75c0237e1..8aa01a3916 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml @@ -8,8 +8,8 @@ version = "0.1.5" description = "LlamaIndex integration for ZeusDB vector database. Enterprise-grade RAG with high-performance vector search." readme = "README.md" requires-python = ">=3.10,<4.0" -license = { text = "MIT" } -authors = [{ name = "ZeusDB", email = "contact@zeusdb.com" }] +license = {text = "MIT"} +authors = [{name = "ZeusDB", email = "contact@zeusdb.com"}] keywords = [ "vector-database", "embeddings", @@ -47,11 +47,11 @@ dev = [ "pre-commit>=4.3", "codespell[toml]>=2.2.6", # "ipython>=9.5", - "ipython>=8.10.0", #need support for python 3.10 + "ipython>=8.10.0", # need support for python 3.10 "jupyter>=1.0.0", - "pylint>=3.3", # Latest + "pylint>=3.3", # Latest # "tree-sitter-languages>=1.10.2", # Not compatible with 3.13 - "types-requests>=2.32", # Latest KEEP + "types-requests>=2.32", # Latest KEEP # "types-setuptools>=80.9", # Latest "llama-index-embeddings-openai>=0.5.1", "llama-index-llms-openai>=0.6.6", @@ -63,29 +63,16 @@ Homepage = "https://github.com/zeusdb/llama-index-vector-stores-zeusdb" Repository = "https://github.com/zeusdb/llama-index-vector-stores-zeusdb" Issues = "https://github.com/zeusdb/llama-index-vector-stores-zeusdb/issues" -# --- LlamaHub metadata (REQUIRED for the PR) --- -[tool.llamahub] -contains_example = true -import_path = "llama_index.vector_stores.zeusdb" - -[tool.llamahub.class_authors] -ZeusDBVectorStore = "zeusdb" -# ---------------------------------------------------------- +[tool.black] +line-length = 88 +target-version = ['py310', 'py311', 'py312'] +include = '\.pyi?$' [tool.codespell] check-filenames = true check-hidden = true skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb" - -[tool.mypy] -disallow_untyped_defs = true -exclude = ["_static", "build", "examples", "notebooks", "venv", ".venv"] -ignore_missing_imports = true -python_version = "3.10" - -[tool.hatch.build.targets.wheel] -packages = ["llama_index"] -include = ["llama_index/vector_stores/zeusdb/py.typed"] +[tool.hatch] [tool.hatch.build.targets.sdist] include = [ @@ -96,23 +83,23 @@ include = [ "llama_index/vector_stores/zeusdb/py.typed", ] -[tool.ruff] -line-length = 88 -target-version = "py310" +[tool.hatch.build.targets.wheel] +packages = ["llama_index"] +include = ["llama_index/vector_stores/zeusdb/py.typed"] -[tool.ruff.lint] -select = [ - "E", # pycodestyle errors - "F", # pyflakes - "I", # isort - "UP", # pyupgrade -] -ignore = [] +# --- LlamaHub metadata (REQUIRED for the PR) --- +[tool.llamahub] +contains_example = true +import_path = "llama_index.vector_stores.zeusdb" -[tool.black] -line-length = 88 -target-version = ['py310', 'py311', 'py312'] -include = '\.pyi?$' +[tool.llamahub.class_authors] +ZeusDBVectorStore = "zeusdb" + +[tool.mypy] +disallow_untyped_defs = true +exclude = ["_static", "build", "examples", "notebooks", "venv", ".venv"] +ignore_missing_imports = true +python_version = "3.10" [tool.pytest.ini_options] testpaths = ["tests"] @@ -121,6 +108,19 @@ python_classes = "Test*" python_functions = "test_*" asyncio_mode = "auto" +[tool.ruff] +line-length = 88 +target-version = "py310" + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "F", # pyflakes + "I", # isort + "UP", # pyupgrade +] +ignore = [] + [tool.ruff.lint.isort] known-first-party = ["llama_index"] section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"] From 595c3021c4c054be6589a405243004733235c11e Mon Sep 17 00:00:00 2001 From: doubleinfinity <6169958+doubleinfinity@users.noreply.github.com> Date: Wed, 5 Nov 2025 01:05:54 +1100 Subject: [PATCH 4/6] fix(vector-stores): align llama-index-core version range with other integrations --- .../llama-index-vector-stores-zeusdb/pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml index 8aa01a3916..6dddc298d6 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/pyproject.toml @@ -31,7 +31,7 @@ classifiers = [ ] # Runtime dependencies dependencies = [ - "llama-index-core>=0.14.6", # Latest release as of Oct 26, 2025 + "llama-index-core>=0.13.0,<0.15", "zeusdb>=0.0.8", # Megapackage version ] @@ -72,6 +72,7 @@ include = '\.pyi?$' check-filenames = true check-hidden = true skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb" + [tool.hatch] [tool.hatch.build.targets.sdist] From 6845af1cf9fa1f581f06300b4d390fa9ee91a635 Mon Sep 17 00:00:00 2001 From: doubleinfinity <6169958+doubleinfinity@users.noreply.github.com> Date: Thu, 6 Nov 2025 18:46:45 +1100 Subject: [PATCH 5/6] fix(vector-stores): implement lazy imports for CI compatibility --- .../llama_index/vector_stores/zeusdb/base.py | 511 ++++++++---------- 1 file changed, 230 insertions(+), 281 deletions(-) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py index eb2c4e3da0..3d1c575894 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/llama_index/vector_stores/zeusdb/base.py @@ -8,49 +8,53 @@ from time import perf_counter from typing import TYPE_CHECKING, Any, cast -from llama_index.core.schema import BaseNode -from llama_index.core.vector_stores.types import ( - BasePydanticVectorStore, - FilterCondition, - FilterOperator, - MetadataFilter, - MetadataFilters, - VectorStoreQuery, - VectorStoreQueryResult, -) - -# ZeusDB runtime (umbrella package only) -from zeusdb import VectorDatabase # type: ignore - +# Use only TYPE_CHECKING imports to avoid runtime hard deps if TYPE_CHECKING: - pass + from llama_index.core.vector_stores.types import ( # type: ignore + VectorStoreQuery, + VectorStoreQueryResult, + ) + +# Try to import the real base, else provide a minimal stub +try: + from llama_index.core.vector_stores.types import ( # type: ignore + BasePydanticVectorStore as _BasePydanticVectorStore, # type: ignore[assignment] + ) +except Exception: # pragma: no cover -# ----------------------------------------------------------------------------- -# Enterprise Logging Integration with Safe Fallback: llamaindex-zeusdb package -# ----------------------------------------------------------------------------- + class _BasePydanticVectorStore: # type: ignore[no-redef] + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + +# ----------------------------------------------------------------------- +# Enterprise Logging Integration with Safe Fallback +# Works without zeusdb or llama_index.core present +# ----------------------------------------------------------------------- try: - from zeusdb.logging_config import ( # type: ignore[import] # noqa: I001 + from zeusdb.logging_config import ( # type: ignore[import] get_logger as _get_logger, ) from zeusdb.logging_config import ( # type: ignore[import] operation_context as _operation_context, ) -except Exception: # fallback for OSS/dev environments +except Exception: # fallback for OSS or dev environments import logging class _StructuredAdapter(logging.LoggerAdapter): """ - Adapter that moves arbitrary kwargs into 'extra' - for stdlib logging compatibility. + Move arbitrary kwargs into 'extra' for stdlib logging compatibility. """ def process( - self, msg: str, kwargs: MutableMapping[str, Any] + self, + msg: str, + kwargs: MutableMapping[str, Any], ) -> tuple[str, MutableMapping[str, Any]]: allowed = {"exc_info", "stack_info", "stacklevel", "extra"} extra = kwargs.get("extra", {}) or {} if not isinstance(extra, dict): - extra = {"_extra": repr(extra)} # defensive + extra = {"_extra": repr(extra)} fields = {k: kwargs.pop(k) for k in list(kwargs.keys()) if k not in allowed} if fields: extra.update(fields) @@ -93,13 +97,14 @@ def _operation_context(operation_name: str, **context: Any) -> Iterator[None]: raise -# Initialize module logger (central config owns handlers/format) +# Initialize module logger through central configuration get_logger: Callable[[str], Any] = cast(Callable[[str], Any], _get_logger) logger = get_logger("llamaindex_zeusdb") operation_context = cast(Callable[..., Any], _operation_context) + # ------------------------- -# Utilities & type helpers +# Utilities and type helpers # ------------------------- _DISTANCE_TO_SPACE = { @@ -120,17 +125,16 @@ def _infer_space(distance: str | None) -> str: def _similarity_from_distance(distance_value: float, space: str) -> float: """ - Convert ZeusDB distance to a similarity score (higher = better). - - cosine: similarity = 1 - distance (assuming normalized embeddings). - - l2/l1: convert to negative distance so higher is better. + Convert ZeusDB distance to a similarity score, higher is better. + cosine: similarity = 1 - distance (assumes normalized embeddings). + l2 or l1: return negative distance to match higher is better. """ if space == "cosine": return 1.0 - float(distance_value) return -float(distance_value) -def _extract_embedding(node: BaseNode) -> list[float] | None: - # LlamaIndex nodes typically have `embedding` populated before add() +def _extract_embedding(node: Any) -> list[float] | None: emb = getattr(node, "embedding", None) if emb is None and hasattr(node, "get_embedding"): try: @@ -140,9 +144,8 @@ def _extract_embedding(node: BaseNode) -> list[float] | None: return list(emb) if emb is not None else None -def _node_metadata(node: BaseNode) -> dict[str, Any]: - md = dict(node.metadata or {}) - # Propagate identifiers for delete-by-ref_doc_id and traceability. +def _node_metadata(node: Any) -> dict[str, Any]: + md = dict(getattr(node, "metadata", {}) or {}) if getattr(node, "ref_doc_id", None): md["ref_doc_id"] = node.ref_doc_id if getattr(node, "id_", None): @@ -150,76 +153,118 @@ def _node_metadata(node: BaseNode) -> dict[str, Any]: return md -def _translate_filter_op(op: FilterOperator) -> str: - """Map LlamaIndex FilterOperator -> ZeusDB operator keys.""" +# Safe enum/filter imports with stubs for offline mode +try: + from llama_index.core.vector_stores.types import ( # type: ignore + FilterCondition, # type: ignore[assignment] + FilterOperator, # type: ignore[assignment] + MetadataFilter, # type: ignore[assignment] + MetadataFilters, # type: ignore[assignment] + ) +except Exception: # pragma: no cover + # Minimal safe stubs to keep static analyzers quiet + class FilterCondition: # type: ignore[no-redef] + AND = type("V", (), {"value": "and"}) + + class FilterOperator: # type: ignore[no-redef] + EQ = "eq" + NE = "ne" + GT = "gt" + LT = "lt" + GTE = "gte" + LTE = "lte" + IN = "in" + NIN = "nin" + ANY = "any" + ALL = "all" + CONTAINS = "contains" + TEXT_MATCH = "text_match" + TEXT_MATCH_INSENSITIVE = "text_match_insensitive" + IS_EMPTY = "is_empty" + + class MetadataFilter: # type: ignore[no-redef] + key: str + value: Any + operator: Any + + class MetadataFilters: # type: ignore[no-redef] + filters: list[Any] + condition: Any + + +def _translate_filter_op(op: Any) -> str: + """ + Map FilterOperator to ZeusDB operator keys. + Works with actual enums or fallback string values. + """ mapping = { - FilterOperator.EQ: "eq", - FilterOperator.NE: "ne", - FilterOperator.GT: "gt", - FilterOperator.LT: "lt", - FilterOperator.GTE: "gte", - FilterOperator.LTE: "lte", - FilterOperator.IN: "in", - FilterOperator.NIN: "nin", - FilterOperator.ANY: "any", - FilterOperator.ALL: "all", - FilterOperator.CONTAINS: "contains", - FilterOperator.TEXT_MATCH: "text_match", - FilterOperator.TEXT_MATCH_INSENSITIVE: "text_match_insensitive", - FilterOperator.IS_EMPTY: "is_empty", + getattr(FilterOperator, "EQ", "eq"): "eq", + getattr(FilterOperator, "NE", "ne"): "ne", + getattr(FilterOperator, "GT", "gt"): "gt", + getattr(FilterOperator, "LT", "lt"): "lt", + getattr(FilterOperator, "GTE", "gte"): "gte", + getattr(FilterOperator, "LTE", "lte"): "lte", + getattr(FilterOperator, "IN", "in"): "in", + getattr(FilterOperator, "NIN", "nin"): "nin", + getattr(FilterOperator, "ANY", "any"): "any", + getattr(FilterOperator, "ALL", "all"): "all", + getattr(FilterOperator, "CONTAINS", "contains"): "contains", + getattr(FilterOperator, "TEXT_MATCH", "text_match"): "text_match", + getattr( + FilterOperator, + "TEXT_MATCH_INSENSITIVE", + "text_match_insensitive", + ): "text_match_insensitive", + getattr(FilterOperator, "IS_EMPTY", "is_empty"): "is_empty", } return mapping.get(op, "eq") -def _filters_to_zeusdb(filters: MetadataFilters | None) -> dict[str, Any] | None: +def _filters_to_zeusdb( + filters: MetadataFilters | None, +) -> dict[str, Any] | None: """ Convert LlamaIndex MetadataFilters to ZeusDB flat format. - ZeusDB expects flat dict with implicit AND: - {"key1": value, "key2": {"op": value}} + ZeusDB expects a flat dict with implicit AND: + {"key1": value, "key2": {"op": value}} """ if filters is None: return None def _one(f: MetadataFilter | MetadataFilters) -> dict[str, Any]: if isinstance(f, MetadataFilters): - cond = (f.condition or FilterCondition.AND).value.lower() - sub = [_one(sf) for sf in f.filters] + cond_val = getattr(getattr(f, "condition", None), "value", "and") + cond = str(cond_val).lower() if cond_val else "and" + sub = [_one(sf) for sf in getattr(f, "filters", [])] if cond == "and": - # Merge into flat dict (implicit AND) - result = {} + merged: dict[str, Any] = {} for s in sub: - result.update(s) - return result - else: - # OR is NOT supported by Rust implementation - logger.warning( - "OR filters not supported by ZeusDB backend", - operation="filter_translation", - condition=cond, - ) - # Fallback: return first filter only - return sub[0] if sub else {} + merged.update(s) + return merged + logger.warning( + "OR filters not supported by ZeusDB backend", + operation="filter_translation", + condition=cond, + ) + return sub[0] if sub else {} - # Single filter - op_key = _translate_filter_op(f.operator) + op_key = _translate_filter_op(getattr(f, "operator", "eq")) + key = getattr(f, "key", "") + val = getattr(f, "value", None) if op_key == "eq": - # Direct value for equality (matches Rust code) - return {f.key: f.value} - else: - # Operator wrapper for other ops - return {f.key: {op_key: f.value}} - - result = _one(filters) # Changed from 'z' to 'result' for consistency + return {key: val} + return {key: {op_key: val}} + result = _one(filters) logger.debug("translated_filters", zeusdb_filter=result) return result # ------------------------- -# MMR helpers (opt-in only) +# MMR helpers (opt in only) # ------------------------- @@ -244,7 +289,7 @@ def _mmr_select( ) -> list[int]: """ Greedy Maximal Marginal Relevance. - Returns indices of the selected candidates. + Returns indices of selected candidates. """ n = len(cand_vecs) if n == 0 or k <= 0: @@ -254,7 +299,6 @@ def _mmr_select( selected: list[int] = [] remaining = set(range(n)) - # seed: most relevant first = max(remaining, key=lambda i: rel[i]) selected.append(first) remaining.remove(first) @@ -262,7 +306,6 @@ def _mmr_select( while len(selected) < min(k, n) and remaining: def score(i: int) -> float: - # diversity term = max sim to any already-selected max_div = max(_cosine_sim(cand_vecs[i], cand_vecs[j]) for j in selected) return lamb * rel[i] - (1.0 - lamb) * max_div @@ -278,19 +321,17 @@ def score(i: int) -> float: # ------------------------- -class ZeusDBVectorStore(BasePydanticVectorStore): +class ZeusDBVectorStore(_BasePydanticVectorStore): # type: ignore[misc] """ - LlamaIndex VectorStore backed by ZeusDB (via the `zeusdb` umbrella package). + LlamaIndex VectorStore backed by ZeusDB (umbrella package). Behaviors: - - Expects nodes with precomputed embeddings. - - Stores vectors + metadata; does not store full text (stores_text=False). - - Translates LlamaIndex MetadataFilters to ZeusDB filter dicts. - - Converts ZeusDB distances to similarity scores (higher = better). - - Supports opt-in MMR when the caller requests it. - - Provides async wrappers via thread offload. - - Persistence Note: Quantized indexes currently load in raw mode + - Expects nodes with precomputed embeddings + - Stores vectors and metadata only (stores_text=False) + - Translates MetadataFilters to ZeusDB flat filters + - Converts distances to similarity scores (higher=better) + - Supports optional MMR when requested + - Provides async wrappers via thread offload """ stores_text: bool = False @@ -299,22 +340,18 @@ class ZeusDBVectorStore(BasePydanticVectorStore): def __init__( self, *, - dim: int | None = None, # Optional if using existing index + dim: int | None = None, distance: str = "cosine", index_type: str = "hnsw", index_name: str = "default", quantization_config: dict[str, Any] | None = None, - # ZeusDB tuning params (optional) m: int | None = None, ef_construction: int | None = None, expected_size: int | None = None, - # Pre-existing ZeusDB index (optional) zeusdb_index: Any | None = None, - # Extra kwargs forwarded to VectorDatabase.create() **kwargs: Any, ) -> None: - # super().__init__(stores_text=self.stores_text) - super().__init__(stores_text=False) # Use the literal value + super().__init__(stores_text=False) self._space = _infer_space(distance) self._index_name = index_name @@ -324,6 +361,9 @@ def __init__( else: if dim is None: raise ValueError("dim is required when not providing zeusdb_index") + # Defer zeusdb import to runtime + from zeusdb import VectorDatabase # type: ignore + vdb = VectorDatabase() create_kwargs: dict[str, Any] = { "index_type": index_type, @@ -348,7 +388,7 @@ def __init__( def client(self) -> Any: return self._index - def add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: + def add(self, nodes: Sequence[Any], **kwargs: Any) -> list[str]: with operation_context( "add_vectors", requested=len(nodes), @@ -370,10 +410,13 @@ def add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: ids.append(str(node_id)) provided_count += 1 else: - ids.append("") # placeholder + ids.append("") if not vectors: - logger.debug("add_vectors no-op (no embeddings)") + logger.debug( + "add_vectors no-op", + reason="no embeddings", + ) return [] payload: dict[str, Any] = { @@ -381,7 +424,6 @@ def add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: "metadatas": metadatas, } - # All-or-nothing ID policy if 0 < provided_count < len(ids): logger.debug( "partial_ids_ignored", @@ -414,42 +456,28 @@ def add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: logger.debug( "add_vectors summary", requested=len(nodes), - inserted=len(assigned_ids) if assigned_ids else len(vectors), + inserted=(len(assigned_ids) if assigned_ids else len(vectors)), had_all_ids=(provided_count == len(ids)), ) - # Return backend IDs if available; else fallback to provided ones if assigned_ids: return assigned_ids return [i for i in ids if i] # ------------------------- - # Deletion & maintenance + # Deletion and maintenance # ------------------------- def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: - """ - Delete all nodes associated with a ref_doc_id. - - ⚠️ LIMITATION: This method is NOT SUPPORTED by ZeusDB's HNSW backend. - - The HNSW index only supports deletion by node ID via remove_point(). - There is no filter-based deletion or scalable way to find all node IDs - for a given ref_doc_id (list() doesn't work in QuantizedOnly mode). - - This method will raise NotImplementedError to be honest about the limitation. - - Alternative: Use delete_nodes(node_ids=[...]) if you have the node IDs. - """ logger.error( - "delete() by ref_doc_id is not supported by ZeusDB HNSW backend", + "delete by ref_doc_id not supported by ZeusDB HNSW backend", operation="delete", ref_doc_id=ref_doc_id, ) raise NotImplementedError( - "ZeusDB HNSW backend does not support deletion by ref_doc_id. " - "The backend only supports ID-based deletion via remove_point(). " - "Use delete_nodes(node_ids=[...]) instead if you have the node IDs." + "ZeusDB HNSW backend does not support deletion by " + "ref_doc_id. Only remove_point() is supported. " + "Use delete_nodes(node_ids=[...]) instead." ) def delete_nodes( @@ -458,27 +486,15 @@ def delete_nodes( filters: MetadataFilters | None = None, **delete_kwargs: Any, ) -> None: - """ - Delete nodes by IDs. - - βœ… SUPPORTED: Deletion by explicit node IDs via remove_point(). - ❌ NOT SUPPORTED: Deletion by metadata filters. - - Args: - node_ids: List of node IDs to delete (supported) - filters: Metadata filters (NOT supported - will raise error if provided) - - Note: ZeusDB HNSW only supports direct ID-based deletion. - """ if filters: logger.error( - "delete_nodes() with filters is not supported by ZeusDB HNSW backend", + "delete_nodes with filters not supported by ZeusDB HNSW backend", operation="delete_nodes", has_filters=True, ) raise NotImplementedError( - "ZeusDB HNSW backend does not support filter-based deletion. " - "Only direct node ID deletion is supported." + "ZeusDB HNSW backend does not support filter based " + "deletion. Only direct node ID deletion is supported." ) if not node_ids: @@ -488,7 +504,7 @@ def delete_nodes( with operation_context("delete_nodes", node_ids_count=len(node_ids)): try: success_count = 0 - failed_ids = [] + failed_ids: list[str] = [] for node_id in node_ids: try: @@ -533,26 +549,20 @@ def delete_nodes( raise def clear(self) -> None: - """ - Clear all vectors from the index. - - ⚠️ LIMITATION: May not work correctly in QuantizedOnly mode. - - The clear() method may not properly clear quantized-only vectors. - """ with operation_context("clear_index"): try: if hasattr(self._index, "clear"): self._index.clear() - logger.info("Index cleared", operation="clear_index") + logger.info( + "Index cleared", + operation="clear_index", + ) else: logger.warning( - "clear() not available on index", + "clear not available on index", operation="clear_index", ) - raise NotImplementedError( - "ZeusDB index does not expose clear() method" - ) + raise NotImplementedError("ZeusDB index does not expose clear") except Exception as e: logger.error( "Clear operation failed", @@ -571,22 +581,25 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul """ Execute a vector search against ZeusDB. - Kwargs understood by this adapter: - mmr (bool): Enable Maximal Marginal Relevance re-ranking. Default False. - mmr_lambda (float): Trade-off [0..1]. 1=relevance, 0=diversity. - Default 0.7 (or the provided `query.mmr_threshold`). - fetch_k (int): Candidate pool when MMR is on. Default max(20, 4*k). - ef_search (int): HNSW runtime search breadth; forwarded to ZeusDB. - return_vector (bool): Ask backend to return raw vectors. Auto-enabled - when MMR is requested. - auto_fallback (bool): If results < k with no filters, retry once with - a broader search. Default True. + Kwargs understood: + mmr (bool): enable MMR reranking, default False + mmr_lambda (float): trade-off [0,1], default 0.7 + fetch_k (int): candidate pool when MMR on + ef_search (int): HNSW search breadth + return_vector (bool): request raw vectors + auto_fallback (bool): broaden search if results < k """ - with operation_context("query", has_embedding=bool(query.query_embedding)): + with operation_context( + "query", + has_embedding=bool(query.query_embedding), + ): if not query.query_embedding: + from llama_index.core.vector_stores.types import ( # type: ignore + VectorStoreQueryResult, + ) + return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) - # Detect explicit MMR requests want_mmr = False mode = getattr(query, "mode", None) if mode and str(mode).lower().endswith("mmr"): @@ -597,9 +610,8 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul if mmr_threshold is not None: want_mmr = True - # Common query prep k = int(query.hybrid_top_k or query.similarity_top_k or 1) - zfilter = _filters_to_zeusdb(query.filters) + zfilter = _filters_to_zeusdb(getattr(query, "filters", None)) ef_search = kwargs.get("ef_search") fetch_k = k if not want_mmr else int(kwargs.get("fetch_k", max(20, 4 * k))) @@ -636,7 +648,6 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul if return_vector: search_kwargs["return_vector"] = True - # Execute search with timing and error context t0 = perf_counter() try: res = self._index.search(**search_kwargs) @@ -652,7 +663,6 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul raise search_ms = (perf_counter() - t0) * 1000 - # Normalize hits hits: list[dict[str, Any]] = [] if isinstance(res, dict) and "results" in res: hits = res.get("results") or [] @@ -680,7 +690,6 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul [float(x) for x in v] if isinstance(v, list) else [] ) - # Broadened search fallback (default on) fallback_used = False if len(cand_ids) < k and not zfilter and kwargs.get("auto_fallback", True): logger.debug( @@ -692,10 +701,13 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul broader_res = self._index.search( vector=list(query.query_embedding), top_k=max(k, fetch_k), - ef_search=max(500, max(k, fetch_k) * 10), + ef_search=max( + 500, + max(k, fetch_k) * 10, + ), return_vector=return_vector, ) - if isinstance(broader_res, dict) and "results" in broader_res: + if isinstance(broader_res, dict) and ("results" in broader_res): broader_hits = broader_res.get("results") or [] elif isinstance(broader_res, list): broader_hits = broader_res @@ -735,7 +747,6 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul error_type=type(e).__name__, ) - # Optional MMR rerank (opt-in only) mmr_ms = 0.0 if want_mmr: if ( @@ -756,6 +767,10 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul mmr_ms = (perf_counter() - t1) * 1000 sel_ids = [cand_ids[i] for i in sel_idx] sel_sims = [rel_q[i] for i in sel_idx] + from llama_index.core.vector_stores.types import ( # type: ignore + VectorStoreQueryResult, + ) + logger.info( "mmr_rerank_applied", selected=len(sel_ids), @@ -767,17 +782,21 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul fallback_used=fallback_used, ) return VectorStoreQueryResult( - nodes=None, similarities=sel_sims, ids=sel_ids + nodes=None, + similarities=sel_sims, + ids=sel_ids, ) - # If vectors missing, fall through to dense ranking - # Default: dense similarity ranking ids: list[str] = [] sims: list[float] = [] for _id, dist in zip(cand_ids, cand_dists): ids.append(_id) sims.append(_similarity_from_distance(dist, self._space)) + from llama_index.core.vector_stores.types import ( # type: ignore + VectorStoreQueryResult, + ) + logger.info( "Query completed", operation="query", @@ -801,7 +820,9 @@ def persist(self, persist_path: str, fs: Any | None = None) -> None: with operation_context("persist_index", path=persist_path): try: if hasattr(self._index, "save"): - self._index.save(persist_path) # type: ignore[attr-defined] + self._index.save( # type: ignore[attr-defined] + persist_path + ) except Exception as e: logger.error( "ZeusDB persist failed", @@ -814,21 +835,20 @@ def persist(self, persist_path: str, fs: Any | None = None) -> None: raise # ------------------------- - # Async wrappers (thread offload) + # Async wrappers # ------------------------- - async def async_add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]: - """Thread-offloaded async variant of add().""" + async def async_add(self, nodes: Sequence[Any], **kwargs: Any) -> list[str]: return await asyncio.to_thread(self.add, nodes, **kwargs) async def aquery( - self, query: VectorStoreQuery, **kwargs: Any + self, + query: VectorStoreQuery, + **kwargs: Any, ) -> VectorStoreQueryResult: - """Thread-offloaded async variant of query().""" return await asyncio.to_thread(self.query, query, **kwargs) async def adelete(self, ref_doc_id: str, **kwargs: Any) -> None: - """Thread-offloaded async variant of delete().""" return await asyncio.to_thread(self.delete, ref_doc_id, **kwargs) async def adelete_nodes( @@ -837,32 +857,28 @@ async def adelete_nodes( filters: MetadataFilters | None = None, **kwargs: Any, ) -> None: - """Thread-offloaded async variant of delete_nodes().""" return await asyncio.to_thread(self.delete_nodes, node_ids, filters, **kwargs) async def aclear(self) -> None: - """Thread-offloaded async variant of clear().""" return await asyncio.to_thread(self.clear) # ------------------------- - # Factory methods and convenience utilities + # Factory methods and utils # ------------------------- @classmethod def from_nodes( cls, - nodes: list[BaseNode], + nodes: list[Any], *, dim: int | None = None, distance: str = "cosine", index_type: str = "hnsw", **kwargs: Any, ) -> ZeusDBVectorStore: - """Create ZeusDBVectorStore from nodes with embeddings.""" if not nodes: raise ValueError("Cannot create store from empty nodes list") - # Infer dimension from first node if not provided if dim is None: first_emb = _extract_embedding(nodes[0]) if first_emb is None: @@ -884,23 +900,14 @@ def load_index( path: str, **kwargs: Any, ) -> ZeusDBVectorStore: - """ - Load ZeusDB index from disk. - - Quantized indexes will load in raw mode. - The quantization model and training state are preserved, but quantized - search will not be active until the next ZeusDB release. - - The index will function correctly using raw vectors, - with full search accuracy but without memory compression benefits. - """ with operation_context("load_index", path=path): + from zeusdb import VectorDatabase # type: ignore + vdb = VectorDatabase() zeusdb_index = vdb.load(path) store = cls(zeusdb_index=zeusdb_index, **kwargs) - # Detect and warn about quantization state try: can_use = store.can_use_quantization() is_active = store.is_quantized() @@ -921,15 +928,17 @@ def load_index( "Quantization config preserved but not active", operation="load_index", compression_ratio=quant_info.get( - "compression_ratio", "N/A" + "compression_ratio", + "N/A", ), subvectors=quant_info.get("subvectors", "N/A"), bits=quant_info.get("bits", "N/A"), ) logger.info( - "Index will use raw vectors. Search accuracy preserved. " - "Memory compression unavailable until next release.", + "Index will use raw vectors. Search " + "accuracy preserved. Memory compression " + "unavailable until next release.", operation="load_index", ) @@ -944,10 +953,12 @@ def load_index( return store def get_vector_count(self) -> int: - """Return total vectors in the index (best-effort).""" + """Return total vectors in index (best-effort).""" try: if hasattr(self._index, "get_vector_count"): - return int(self._index.get_vector_count()) # type: ignore + return int( + self._index.get_vector_count() # type: ignore + ) except Exception as e: logger.error( "get_vector_count failed", @@ -971,10 +982,12 @@ def get_zeusdb_stats(self) -> dict[str, Any]: return {} def save_index(self, path: str) -> bool: - """Save index to disk (best-effort wrapper).""" + """Save index to disk (best-effort).""" try: if hasattr(self._index, "save"): - self._index.save(path) # type: ignore[attr-defined] + self._index.save( # type: ignore[attr-defined] + path + ) return True except Exception as e: logger.error( @@ -986,17 +999,13 @@ def save_index(self, path: str) -> bool: return False def info(self) -> str: - """ - Get a human-readable info string about the index. - - Example: - >>> print(vector_store.info()) - HNSWIndex(dim=1536, space=cosine, vectors=1200, quantized=True, ...) - """ + """Get human-readable info string about index.""" try: info_str = self._index.info() logger.debug( - "Retrieved index info", operation="info", info_length=len(info_str) + "Retrieved index info", + operation="info", + info_length=len(info_str), ) return info_str except Exception as e: @@ -1009,22 +1018,12 @@ def info(self) -> str: ) return f"ZeusDBVectorStore(error: {e})" - # ------------------------------------------------------------------------- + # --------------------------------------------------------------- # Quantization Methods - # ------------------------------------------------------------------------- + # --------------------------------------------------------------- def get_training_progress(self) -> float: - """ - Get quantization training progress percentage. - - Returns: - float: Training progress as percentage (0.0 to 100.0). - Returns 0.0 if quantization is not configured or on error. - - Example: - >>> progress = vector_store.get_training_progress() - >>> print(f"Training: {progress:.1f}% complete") - """ + """Get quantization training progress percentage.""" try: progress = self._index.get_training_progress() logger.debug( @@ -1044,17 +1043,7 @@ def get_training_progress(self) -> float: return 0.0 def is_quantized(self) -> bool: - """ - Check whether quantized search is currently active. - - Returns: - bool: True if index is using quantized vectors for search, - False otherwise or on error. - - Example: - >>> if vector_store.is_quantized(): - ... print("Using quantized search") - """ + """Check whether quantized search is active.""" try: quantized = self._index.is_quantized() logger.debug( @@ -1074,17 +1063,7 @@ def is_quantized(self) -> bool: return False def can_use_quantization(self) -> bool: - """ - Check whether quantization is available (e.g., PQ training completed). - - Returns: - bool: True if quantization is trained and ready to use, - False otherwise or on error. - - Example: - >>> if vector_store.can_use_quantization(): - ... print("Quantization ready") - """ + """Check whether quantization is available.""" try: available = self._index.can_use_quantization() logger.debug( @@ -1104,21 +1083,7 @@ def can_use_quantization(self) -> bool: return False def get_storage_mode(self) -> str: - """ - Get current storage mode. - - Returns: - str: Storage mode string. Possible values: - - 'raw_only': Only raw vectors stored - - 'quantized_only': Only quantized vectors (memory optimized) - - 'quantized_with_raw': Both quantized and raw vectors - - 'quantized_active': Quantization is active - - 'unknown': On error or unable to determine - - Example: - >>> mode = vector_store.get_storage_mode() - >>> print(f"Storage mode: {mode}") - """ + """Get current storage mode.""" try: mode = self._index.get_storage_mode() logger.debug( @@ -1137,26 +1102,10 @@ def get_storage_mode(self) -> str: ) return "unknown" - def get_quantization_info(self) -> dict[str, Any] | None: - """ - Get detailed quantization information. - - Returns: - Optional[Dict]: Dictionary containing quantization details: - - compression_ratio: Memory compression factor (e.g., 16.0 for 16x) - - memory_mb: Estimated memory usage in megabytes - - subvectors: Number of subvectors used - - bits: Bits per quantized code - - trained: Whether training is complete - - training_size: Number of vectors used for training - Returns None if quantization is not configured/trained or on error. - - Example: - >>> info = vector_store.get_quantization_info() - >>> if info: - ... print(f"Compression: {info['compression_ratio']:.1f}x") - ... print(f"Memory: {info['memory_mb']:.2f} MB") - """ + def get_quantization_info( + self, + ) -> dict[str, Any] | None: + """Get detailed quantization information.""" try: info = self._index.get_quantization_info() logger.debug( From 96b09fad9f48d935cbcff2b1349128b2eaf1803f Mon Sep 17 00:00:00 2001 From: doubleinfinity <6169958+doubleinfinity@users.noreply.github.com> Date: Fri, 7 Nov 2025 22:26:47 +1100 Subject: [PATCH 6/6] fix(vector-stores): add conftest.py with graceful mock fallback for CI --- .../tests/conftest.py | 343 ++++++++++++++++++ .../tests/test_vector_stores_zeusdb.py | 29 +- 2 files changed, 354 insertions(+), 18 deletions(-) create mode 100644 llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/conftest.py diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/conftest.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/conftest.py new file mode 100644 index 0000000000..2394aa9c5b --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/conftest.py @@ -0,0 +1,343 @@ +""" +Pytest bootstrap for ZeusDBVectorStore tests. + +Prefers real deps, falls back to lightweight mocks when missing. +""" + +from contextlib import contextmanager +import sys +from types import ModuleType + +# ---------------------------------------- +# Try real llama_index.core first +# ---------------------------------------- +try: + import llama_index.core.vector_stores.types # type: ignore # noqa: F401 + + _HAS_CORE = True +except Exception: + _HAS_CORE = False + +# ---------------------------------------- +# Try real zeusdb first +# ---------------------------------------- +try: + import zeusdb # type: ignore # noqa: F401 + + _HAS_ZEUSDB = True +except Exception: + _HAS_ZEUSDB = False + +# ---------------------------------------- +# Mock llama_index.core types if missing +# ---------------------------------------- +if not _HAS_CORE: + + class MockVectorStoreQuery: + """Mock VectorStoreQuery.""" + + def __init__( + self, + query_embedding=None, + similarity_top_k=1, + hybrid_top_k=None, + filters=None, + mode=None, + mmr_threshold=None, + query_str=None, + **kwargs, + ): + self.query_embedding = query_embedding + self.similarity_top_k = similarity_top_k + self.hybrid_top_k = hybrid_top_k + self.filters = filters + self.mode = mode + self.mmr_threshold = mmr_threshold + self.query_str = query_str + + class MockVectorStoreQueryResult: + """Mock VectorStoreQueryResult.""" + + def __init__(self, nodes=None, similarities=None, ids=None): + self.nodes = nodes or [] + self.similarities = similarities or [] + self.ids = ids or [] + + def __iter__(self): + """Match construction style used by code under test.""" + yield from () + + class MockFilterOperator: + """Mock FilterOperator enum.""" + + EQ = "eq" + NE = "ne" + GT = "gt" + LT = "lt" + GTE = "gte" + LTE = "lte" + IN = "in" + NIN = "nin" + ANY = "any" + ALL = "all" + CONTAINS = "contains" + TEXT_MATCH = "text_match" + TEXT_MATCH_INSENSITIVE = "text_match_insensitive" + IS_EMPTY = "is_empty" + + class MockFilterCondition: + """Mock FilterCondition enum.""" + + class AND: + value = "and" + + class MockMetadataFilter: + """Mock MetadataFilter.""" + + def __init__(self, key, value, operator): + self.key = key + self.value = value + self.operator = operator + + class MockMetadataFilters: + """Mock MetadataFilters.""" + + def __init__(self, filters=None, condition=None): + self.filters = filters or [] + self.condition = condition or MockFilterCondition.AND + + @classmethod + def from_dicts(cls, items, condition=None) -> "MockMetadataFilters": + """Build MetadataFilters from list of dict specs.""" + fl = [] + for it in items or []: + fl.append( + MockMetadataFilter( + key=it.get("key"), + value=it.get("value"), + operator=it.get("operator", MockFilterOperator.EQ), + ) + ) + return cls(filters=fl, condition=condition or MockFilterCondition.AND) + + class MockBasePydanticVectorStore: + """Mock BasePydanticVectorStore base class.""" + + stores_text: bool = False + is_embedding_query: bool = True + + def __init__(self, stores_text=False, **kwargs): + self.stores_text = stores_text + + # Build the module tree + from typing import Any, cast + + core_types: ModuleType = ModuleType("llama_index.core.vector_stores.types") + ct = cast(Any, core_types) + ct.VectorStoreQuery = MockVectorStoreQuery + ct.VectorStoreQueryResult = MockVectorStoreQueryResult + ct.FilterOperator = MockFilterOperator + ct.FilterCondition = MockFilterCondition + ct.MetadataFilter = MockMetadataFilter + ct.MetadataFilters = MockMetadataFilters + ct.BasePydanticVectorStore = MockBasePydanticVectorStore + + sys.modules.setdefault("llama_index", ModuleType("llama_index")) + sys.modules.setdefault("llama_index.core", ModuleType("llama_index.core")) + sys.modules["llama_index.core.vector_stores"] = ModuleType( + "llama_index.core.vector_stores" + ) + sys.modules["llama_index.core.vector_stores.types"] = core_types + +# ---------------------------------------- +# Mock zeusdb if missing, with minimal behavior +# ---------------------------------------- +if not _HAS_ZEUSDB: + + class _InMemoryIndex: + """Minimal in-memory vector index for CI testing.""" + + def __init__(self, *, dim, space, **kwargs): + self.dim = dim + self.space = space + self._vectors: list[list[float]] = [] + self._metadatas: list[dict] = [] + self._ids: list[str] = [] + self._id_counter = 0 + + def add(self, payload, overwrite=True): + """Add vectors to index.""" + vectors = payload.get("vectors") or [] + metadatas = payload.get("metadatas") or [{}] * len(vectors) + ids = payload.get("ids") + out_ids = [] + for i, vec in enumerate(vectors): + if ids and ids[i]: + _id = str(ids[i]) + else: + self._id_counter += 1 + _id = str(self._id_counter) + self._vectors.append(list(vec)) + self._metadatas.append(dict(metadatas[i] if i < len(metadatas) else {})) + self._ids.append(_id) + out_ids.append(_id) + return {"ids": out_ids} + + def _dist(self, a, b): + """Calculate distance between vectors.""" + if self.space == "cosine": + import math + + def _dot(x, y): + return sum(xx * yy for xx, yy in zip(x, y)) + + def _norm(x): + return math.sqrt(max(1e-12, sum(xx * xx for xx in x))) + + return 1.0 - (_dot(a, b) / (_norm(a) * _norm(b))) + # default l2 + import math + + return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b))) + + def search( + self, + *, + vector, + top_k, + filter=None, + ef_search=None, + return_vector=False, + ): + """Search for similar vectors.""" + results = [] + for _id, v, md in zip(self._ids, self._vectors, self._metadatas): + # Simple filter: exact match on key: value + if filter: + ok = True + for k, val in filter.items(): + if isinstance(val, dict): + if "eq" in val: + if md.get(k) != val["eq"]: + ok = False + break + else: + if md.get(k) != val: + ok = False + break + if not ok: + continue + d = self._dist(vector, v) + item = {"id": _id, "distance": d} + if return_vector: + item["vector"] = list(v) + results.append(item) + results.sort(key=lambda x: x["distance"]) + return results[: int(top_k)] + + def remove_point(self, node_id): + """Remove a point by ID.""" + try: + idx = self._ids.index(str(node_id)) + except ValueError: + return False + for arr in (self._ids, self._vectors, self._metadatas): + arr.pop(idx) + return True + + def clear(self): + """Clear all vectors.""" + self._vectors.clear() + self._metadatas.clear() + self._ids.clear() + + def save(self, path): + """Save index (no-op for tests).""" + return True + + def get_vector_count(self): + """Get number of vectors.""" + return len(self._ids) + + def get_stats(self): + """Get index statistics.""" + return {"count": len(self._ids)} + + def info(self): + """Get index info string.""" + return ( + f"HNSWIndex(dim={self.dim}, " + f"space={self.space}, vectors={len(self._ids)})" + ) + + def get_training_progress(self): + """Get quantization training progress.""" + return 0.0 + + def is_quantized(self): + """Check if quantized.""" + return False + + def can_use_quantization(self): + """Check if quantization available.""" + return False + + def get_storage_mode(self): + """Get storage mode.""" + return "raw_only" + + def get_quantization_info(self): + """Get quantization info.""" + return None + + def load(self, path): + """Load index (not used on instance).""" + return self + + class VectorDatabase: + """Mock VectorDatabase.""" + + def create(self, **kwargs): + """Create index.""" + return _InMemoryIndex(**kwargs) + + def load(self, path): + """Load index from path.""" + return _InMemoryIndex(dim=1, space="cosine") + + # Minimal logging_config mock + logging_config: ModuleType = ModuleType("zeusdb.logging_config") + + def get_logger(name: str): + """Get logger.""" + + class _Dummy: + def debug(self, *a, **k): + pass + + def info(self, *a, **k): + pass + + def warning(self, *a, **k): + pass + + def error(self, *a, **k): + pass + + return _Dummy() + + @contextmanager + def operation_context(operation_name: str, **context): + """Operation context manager.""" + yield + + logging_config.get_logger = get_logger # type: ignore[attr-defined] + logging_config.operation_context = ( # type: ignore[attr-defined] + operation_context + ) + + zeusdb_mod: ModuleType = ModuleType("zeusdb") + zeusdb_mod.VectorDatabase = VectorDatabase # type: ignore[attr-defined] + + sys.modules["zeusdb"] = zeusdb_mod + sys.modules["zeusdb.logging_config"] = logging_config diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py index 7bb8c8bcce..9bdbd86d09 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-zeusdb/tests/test_vector_stores_zeusdb.py @@ -7,10 +7,18 @@ import pytest +# Import llama_index.core types at module level so conftest can mock them FIRST +from llama_index.core.vector_stores.types import ( + FilterCondition, + FilterOperator, + MetadataFilters, + VectorStoreQuery, +) -# ----------------------------- + +# --------- -------- ----------- # Minimal in-memory ZeusDB fake -# ----------------------------- +# ------- -------- ------------ def _cosine_distance(a, b): dot = sum(x * y for x, y in zip(a, b)) na = math.sqrt(sum(x * x for x in a)) or 1e-12 @@ -223,8 +231,6 @@ def test_add_and_query_basic(ZeusDBVectorStore): assert set(out_ids) == {"1", "2"} q = [1.0, 0.0, 0.0] - from llama_index.core.vector_stores.types import VectorStoreQuery - res = store.query(VectorStoreQuery(query_embedding=q, similarity_top_k=1)) assert res.ids and res.similarities assert res.ids[0] in {"1", "2"} @@ -241,14 +247,7 @@ def test_filters_and_delete_nodes(ZeusDBVectorStore): ] store.add(nodes) - from llama_index.core.vector_stores.types import ( - FilterCondition, - FilterOperator, - MetadataFilters, - VectorStoreQuery, - ) - - mf = MetadataFilters.from_dicts( + mf = MetadataFilters.from_dicts( # type: ignore[misc] [ {"key": "category", "value": "x", "operator": FilterOperator.EQ}, ], @@ -310,8 +309,6 @@ def test_delete_nodes_by_id_works(ZeusDBVectorStore): assert store.get_vector_count() == 1 - from llama_index.core.vector_stores.types import VectorStoreQuery - res = store.query( VectorStoreQuery(query_embedding=[1.0, 0.0, 0.0], similarity_top_k=10) ) @@ -366,8 +363,6 @@ async def test_async_paths(ZeusDBVectorStore): assert ids == ["1"] q = [1.0, 0.0, 0.0] - from llama_index.core.vector_stores.types import VectorStoreQuery - res = await store.aquery(VectorStoreQuery(query_embedding=q, similarity_top_k=1)) assert res.ids and res.ids[0] == "1" @@ -394,8 +389,6 @@ def test_query_with_mmr(ZeusDBVectorStore): ] store.add(nodes) - from llama_index.core.vector_stores.types import VectorStoreQuery - # request k=2 with MMR enabled; fetch_k larger to force rerank variety res = store.query( VectorStoreQuery(query_embedding=[1.0, 0.0, 0.0], similarity_top_k=2),