|
| 1 | +# Adding a New Vector Database |
| 2 | + |
| 3 | +This guide will walk you through the process of adding a new vector database to Llama Stack. |
| 4 | + |
| 5 | +> **_NOTE:_** Here's an example Pull Request of the [Milvus Vector Database Provider](https://github.com/meta-llama/llama-stack/pull/1467). |
| 6 | +
|
| 7 | +Vector Database providers are used to store and retrieve vector embeddings. Vector databases are not limited to vector |
| 8 | +search but can support keyword and hybrid search. Additionally, vector database can also support operations like |
| 9 | +filtering, sorting, and aggregating vectors. |
| 10 | + |
| 11 | +## Steps to Add a New Vector Database Provider |
| 12 | +1. **Choose the Database Type**: Determine if your vector database is a remote service, inline, or both. |
| 13 | + - Remote databases make requests to external services, while inline databases execute locally. Some providers support both. |
| 14 | +2. **Implement the Provider**: Create a new provider class that inherits from `VectorDatabaseProvider` and implements the required methods. |
| 15 | + - Implement methods for vector storage, retrieval, search, and any additional features your database supports. |
| 16 | + - You will need to implement the following methods for `YourVectorIndex`: |
| 17 | + - `YourVectorIndex.create()` |
| 18 | + - `YourVectorIndex.initialize()` |
| 19 | + - `YourVectorIndex.add_chunks()` |
| 20 | + - `YourVectorIndex.delete_chunk()` |
| 21 | + - `YourVectorIndex.query_vector()` |
| 22 | + - `YourVectorIndex.query_keyword()` |
| 23 | + - `YourVectorIndex.query_hybrid()` |
| 24 | + - You will need to implement the following methods for `YourVectorIOAdapter`: |
| 25 | + - `YourVectorIOAdapter.initialize()` |
| 26 | + - `YourVectorIOAdapter.shutdown()` |
| 27 | + - `YourVectorIOAdapter.list_vector_dbs()` |
| 28 | + - `YourVectorIOAdapter.register_vector_db()` |
| 29 | + - `YourVectorIOAdapter.unregister_vector_db()` |
| 30 | + - `YourVectorIOAdapter.insert_chunks()` |
| 31 | + - `YourVectorIOAdapter.query_chunks()` |
| 32 | + - `YourVectorIOAdapter.delete_chunks()` |
| 33 | +3. **Add to Registry**: Register your provider in the appropriate registry file. |
| 34 | + - Update {repopath}`llama_stack/providers/registry/vector_io.py` to include your new provider. |
| 35 | +```python |
| 36 | +from llama_stack.providers.registry.specs import InlineProviderSpec |
| 37 | +from llama_stack.providers.registry.api import Api |
| 38 | + |
| 39 | +InlineProviderSpec( |
| 40 | + api=Api.vector_io, |
| 41 | + provider_type="inline::milvus", |
| 42 | + pip_packages=["pymilvus>=2.4.10"], |
| 43 | + module="llama_stack.providers.inline.vector_io.milvus", |
| 44 | + config_class="llama_stack.providers.inline.vector_io.milvus.MilvusVectorIOConfig", |
| 45 | + api_dependencies=[Api.inference], |
| 46 | + optional_api_dependencies=[Api.files], |
| 47 | + description="", |
| 48 | +), |
| 49 | +``` |
| 50 | +4. **Add Tests**: Create unit tests and integration tests for your provider in the `tests/` directory. |
| 51 | + - Unit Tests |
| 52 | + - By following the structure of the class methods, you will be able to easily run unit and integration tests for your database. |
| 53 | + 1. You have to configure the tests for your provide in `/tests/unit/providers/vector_io/conftest.py`. |
| 54 | + 2. Update the `vector_provider` fixture to include your provider if they are an inline provider. |
| 55 | + 3. Create a `your_vectorprovider_index` fixture that initializes your vector index. |
| 56 | + 4. Create a `your_vectorprovider_adapter` fixture that initializes your vector adapter. |
| 57 | + 5. Add your provider to the `vector_io_providers` fixture dictionary. |
| 58 | + - Please follow the naming convention of `your_vectorprovider_index` and `your_vectorprovider_adapter` as the tests require this to execute properly. |
| 59 | + - Integration Tests |
| 60 | + - Integration tests are located in {repopath}`tests/integration`. These tests use the python client-SDK APIs (from the `llama_stack_client` package) to test functionality. |
| 61 | + - The two set of integration tests are: |
| 62 | + - `tests/integration/vector_io/test_vector_io.py`: This file tests registration, insertion, and retrieval. |
| 63 | + - `tests/integration/vector_io/test_openai_vector_stores.py`: These tests are for OpenAI-compatible vector stores and test the OpenAI API compatibility. |
| 64 | + - You will need to update `skip_if_provider_doesnt_support_openai_vector_stores` to include your provider as well as `skip_if_provider_doesnt_support_openai_vector_stores_search` to test the appropriate search functionality. |
| 65 | + - Running the tests in the GitHub CI |
| 66 | + - You will need to update the `.github/workflows/integration-vector-io-tests.yml` file to include your provider. |
| 67 | + - If your provider is a remote provider, you will also have to add a container to spin up and run it in the action. |
| 68 | + - Updating the pyproject.yml |
| 69 | + - If you are adding tests for the `inline` provider you will have to update the `unit` group. |
| 70 | + - `uv add new_pip_package --group unit` |
| 71 | + - If you are adding tests for the `remote` provider you will have to update the `test` group, which is used in the GitHub CI for integration tests. |
| 72 | + - `uv add new_pip_package --group test` |
| 73 | +5. **Update Documentation**: Please update the documentation for end users |
| 74 | + - Generate the provider documentation by running {repopath}`./scripts/provider_codegen.py`. |
| 75 | + - Update the autogenerated content in the registry/vector_io.py file with information about your provider. Please see other providers for examples. |
0 commit comments