Vector Database Integration

Add a new vector store adapter to Cognee.

Adding a New Vector Store to cognee

This guide describes how to integrate a new vector database engine into cognee, following the same pattern used for existing engines (e.g., Weaviate, Qdrant, Milvus, PGVector, LanceDB).

Repository Options

Cognee used both the core and the community repositories to host vector-database adapters.

🚨 From now on every vector-database adapter – except LanceDB – will live in the cognee-community repository.

Qdrant have already been migrated from the core library and the remaining adapters will follow shortly. You can find Redis, OpenSearch, and Azure AI Search integrations available in the community repo.

Therefore all new adapter contributions must target the community repository.
The core repository will keep only the built-in LanceDB integration.

For Community Repository

To add a new adapter to cognee-community:

1. Fork and clone the cognee-community repository, and branch from main — unlike the core Cognee repo, cognee-community has no dev branch.
2. Create your adapter in packages/vector/<engine_name>/cognee_community_vector_adapter_<engine_name>/
3. Inside that directory add __init__.py, <engine_name>_adapter.py, and register.py (see the Redis implementation as an example).
4. At the package root packages/<engine_name>/ add __init__.py, pyproject.toml, and README.md.
5. Run the shared vector conformance suite in packages/shared/contract_suite/vector_contract.py against your adapter.
6. Submit a pull request to the community repository

Below are the recommended steps in more detail.

---

Why cognee-community?

cognee-community is the extension hub for Cognee.
Anything that is not part of the core lives here—adapters for third-party databases, pipelines, community contributed additional tasks, etc.
Placing your adapter in this repository means:

• Your code is released under the community license and can evolve independently of the core.
• It can be installed with pip install cognee-community-vector-adapter-(engine_name) without pulling in heavyweight drivers for users who don’t need them. For example, for Qdrant it is pip install cognee-community-vector-adapter-qdrant
• These packages can be called with the cognee core package using the registration step described below.

If you are unfamiliar with the layout, have a look at the existing folders under packages/* in the community repo—each sub-folder represents a separate provider implemented in exactly the way you are about to do.

---

1. Implement the Adapter

File: packages/vector/<engine_name>/cognee_community_vector_adapter_<engine_name>/<engine_name>_adapter.py

Your adapter must implement the VectorDBInterface protocol, implementing all required methods for collection management, data point operations, and search functionality. Here is a sample skeleton with placeholders:

"""Adapter for <engine_name> vector database."""

import asyncio
from typing import List, Optional, Dict, Any
from concurrent.futures import ThreadPoolExecutor

from cognee.exceptions import InvalidValueError
from cognee.infrastructure.engine import DataPoint
from cognee.infrastructure.engine.utils import parse_id
from cognee.shared.logging_utils import get_logger

from ..embeddings.EmbeddingEngine import EmbeddingEngine
from ..models.ScoredResult import ScoredResult
from ..utils import normalize_distances
from ..vector_db_interface import VectorDBInterface
from ..exceptions import CollectionNotFoundError

logger = get_logger()

class IndexSchema(DataPoint):
    """
    Define a schema for indexing data points with a text field.
    
    This class inherits from the DataPoint class and specifies the structure of a single
    data point that includes a text attribute.
    """
    text: str
    metadata: dict = {"index_fields": ["text"]}

class <EngineName>Adapter(VectorDBInterface):
    """Adapter for <engine_name> vector database operations."""
    
    name = "<EngineName>"
    
    def __init__(
        self,
        url: str,
        api_key: Optional[str] = None,
        embedding_engine: EmbeddingEngine = None,
        database_name: Optional[str] = None,
        vector_db_host: str = "",
        vector_db_port: str = "",
        vector_db_username: str = "",
        vector_db_password: str = "",
        **kwargs,
    ):
        self.url = url
        self.api_key = api_key
        self.embedding_engine = embedding_engine
        self.database_name = database_name
        # Connection details for a self-hosted store on a non-default
        # host/port, or one that requires credentials.
        self.host = vector_db_host
        self.port = vector_db_port
        self.username = vector_db_username
        self.password = vector_db_password
        self.client = None
        self.executor = ThreadPoolExecutor()
        self._initialize_connection()
    
    def _initialize_connection(self) -> None:
        """Establish connection to <engine_name>."""
        try:
            # Initialize your vector database client here
            # Example: self.client = YourVectorDBClient(url=self.url, api_key=self.api_key)
            logger.debug(f"Successfully connected to <engine_name> at {self.url}")
        except Exception as e:
            logger.error(f"Failed to initialize <engine_name> connection: {e}")
            raise
    
    async def has_collection(self, collection_name: str) -> bool:
        """Check if a specified collection exists."""
        try:
            # Implement collection existence check
            # Example: return await self.client.collection_exists(collection_name)
            return False
        except Exception as e:
            logger.error(f"Error checking collection existence: {e}")
            return False
    
    async def create_collection(
        self,
        collection_name: str,
        payload_schema: Optional[object] = None,
    ):
        """Create a new collection with an optional payload schema."""
        if await self.has_collection(collection_name):
            return
        
        try:
            vector_size = self.embedding_engine.get_vector_size()
            # Implement collection creation logic
            # Example:
            # await self.client.create_collection(
            #     name=collection_name,
            #     vector_size=vector_size,
            #     distance_metric="cosine"
            # )
            logger.debug(f"Created collection: {collection_name}")
        except Exception as e:
            logger.error(f"Failed to create collection {collection_name}: {e}")
            raise
    
    async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
        """Insert new data points into the specified collection."""
        if not await self.has_collection(collection_name):
            await self.create_collection(collection_name)
        
        # Generate embeddings for data points
        embeddings = await self.embedding_engine.embed_text(
            [DataPoint.get_embeddable_data(data_point) for data_point in data_points]
        )
        
        try:
            # Implement data point insertion logic
            # Example:
            # formatted_points = [
            #     {
            #         "id": str(data_point.id),
            #         "vector": embeddings[i],
            #         "payload": data_point.model_dump()
            #     }
            #     for i, data_point in enumerate(data_points)
            # ]
            # await self.client.upsert(collection_name, formatted_points)
            logger.debug(f"Inserted {len(data_points)} data points into {collection_name}")
        except Exception as e:
            logger.error(f"Failed to insert data points: {e}")
            raise
    
    async def retrieve(self, collection_name: str, data_point_ids: List[str]):
        """Retrieve data points from a collection using their IDs."""
        if not data_point_ids:
            # No ids requested. Return [] instead of building an empty
            # "id IN ()" filter, which some backends reject as a parse error.
            # The built-in LanceDB and PGVector adapters return [] here.
            return []
        try:
            # Implement data point retrieval logic
            # Example:
            # results = await self.client.retrieve(collection_name, data_point_ids)
            # return [
            #     ScoredResult(
            #         id=parse_id(result["id"]),
            #         payload=result["payload"],
            #         score=0
            #     )
            #     for result in results
            # ]
            return []
        except Exception as e:
            logger.error(f"Failed to retrieve data points: {e}")
            return []
    
    async def search(
        self,
        collection_name: str,
        query_text: Optional[str] = None,
        query_vector: Optional[List[float]] = None,
        limit: int = 15,
        with_vector: bool = False,
    ) -> List[ScoredResult]:
        """Perform a search in the specified collection."""
        if query_text is None and query_vector is None:
            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
        
        if not await self.has_collection(collection_name):
            logger.warning(f"Collection '{collection_name}' not found; returning [].")
            return []
        
        if query_vector is None:
            query_vector = (await self.embedding_engine.embed_text([query_text]))[0]
        
        try:
            # Implement search logic
            # Example:
            # results = await self.client.search(
            #     collection_name=collection_name,
            #     query_vector=query_vector,
            #     limit=limit,
            #     include_vector=with_vector
            # )
            # 
            # return [
            #     ScoredResult(
            #         id=parse_id(result["id"]),
            #         payload=result["payload"],
            #         score=result["score"]
            #     )
            #     for result in results
            # ]
            return []
        except Exception as e:
            logger.error(f"Error searching collection '{collection_name}': {e}")
            return []
    
    async def batch_search(
        self,
        collection_name: str,
        query_texts: List[str],
        limit: int,
        with_vectors: bool = False,
    ):
        """Perform a batch search using multiple text queries."""
        query_vectors = await self.embedding_engine.embed_text(query_texts)
        
        return await asyncio.gather(
            *[
                self.search(
                    collection_name=collection_name,
                    query_vector=query_vector,
                    limit=limit,
                    with_vector=with_vectors,
                )
                for query_vector in query_vectors
            ]
        )
    
    async def delete_data_points(self, collection_name: str, data_point_ids: List[str]):
        """Delete specified data points from a collection."""
        try:
            # Implement deletion logic
            # Example:
            # result = await self.client.delete(collection_name, data_point_ids)
            # return result
            logger.debug(f"Deleted {len(data_point_ids)} data points from {collection_name}")
        except Exception as e:
            logger.error(f"Failed to delete data points: {e}")
            raise
    
    async def create_vector_index(self, index_name: str, index_property_name: str):
        """Create a vector index based on an index name and property name."""
        return await self.create_collection(f"{index_name}_{index_property_name}")
    
    async def index_data_points(
        self, index_name: str, index_property_name: str, data_points: List[DataPoint]
    ):
        """Index a list of data points by creating an associated vector index collection."""
        formatted_data_points = [
            IndexSchema(
                id=data_point.id,
                text=DataPoint.get_embeddable_data(data_point),
            )
            for data_point in data_points
        ]
        
        return await self.create_data_points(
            f"{index_name}_{index_property_name}",
            formatted_data_points,
        )
    
    async def prune(self):
        """Remove all data from the vector database."""
        try:
            # Implement pruning logic - delete all collections or data
            # Example:
            # collections = await self.client.list_collections()
            # for collection in collections:
            #     await self.client.delete_collection(collection.name)
            logger.debug("Pruned all data from <engine_name>")
        except Exception as e:
            logger.error(f"Failed to prune data: {e}")
            raise

Keep the method signatures consistent with VectorDBInterface. Reference the existing adapters like QDrantAdapter for more comprehensive examples.

What Cognee passes your constructor

Cognee never instantiates your adapter directly — the vector engine factory (create_vector_engine) does, resolving your class out of the registry by the provider name you passed to use_vector_adapter. It calls your __init__ with these keyword arguments, all by keyword, never positionally:

Keyword	Source	Notes
url	vector_db_url
api_key	vector_db_key
embedding_engine	resolved by Cognee	an EmbeddingEngine instance
database_name	vector_db_name
vector_db_host	vector_db_host	connection details — these let your adapter reach a store on a non-default host/port, or one that requires credentials
vector_db_port	vector_db_port	forwarded as a string (see below)
vector_db_username	vector_db_username
vector_db_password	vector_db_password

Apart from embedding_engine, every value above comes from user configuration — set through config.set_vector_db_config({...}) or the matching VECTOR_DB_* environment variables. See the community adapters overview.

Two normalization details to code against:

• vector_db_port arrives as a string, not an int. Numeric values are stringified before forwarding. It also defaults to 1234 in VectorConfig, so you receive vector_db_port="1234" even when the user never configured a port — don’t treat a truthy port as proof the user chose one.
• None never reaches you for these parameters. Any of them left as None is replaced by the factory’s own default, an empty string. So an unset host, username, or password arrives as "".

Accept **kwargs. Cognee may forward additional keywords to registry adapters over time, and an adapter with a closed signature fails at construction with TypeError: __init__() got an unexpected keyword argument. A trailing **kwargs (as in the skeleton above) keeps your adapter working across Cognee versions — ignore what you don’t need rather than declaring it.

If you store the credentials on the instance, keep them out of your logs and error messages — see Security below.

2. Test with a Dedicated Script

File: packages/vector/engine_name/examples/example.py

Create a script that loads cognee, configures it to use your new <engine_name> provider, and runs basic usage checks. For example:

import os
import pathlib
import cognee
from cognee import SearchType
from cognee.modules.search.operations import get_history
from cognee.modules.users.methods import get_default_user
from cognee.shared.logging_utils import get_logger

logger = get_logger()

async def main():
    cognee.config.set_vector_db_provider("<engine_name>")
    data_directory_path = str(
        pathlib.Path(
            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_<engine_name>")
        ).resolve()
    )
    cognee.config.data_root_directory(data_directory_path)
    cognee_directory_path = str(
        pathlib.Path(
            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_<engine_name>")
        ).resolve()
    )
    cognee.config.system_root_directory(cognee_directory_path)

    await cognee.prune.prune_data()
    await cognee.prune.prune_system(metadata=True)

    dataset_name = "test_dataset"

    # Add some test data
    test_text = """
    This is a test document for the <engine_name> vector database integration.
    It contains information about artificial intelligence and machine learning.
    Vector databases are essential for semantic search and retrieval applications.
    """

    await cognee.remember(test_text, dataset_name=dataset_name)

    # Test vector search functionality
    from cognee.infrastructure.databases.vector import get_vector_engine_async

    vector_engine = await get_vector_engine_async()
    search_results = await vector_engine.search("Entity_name", "artificial intelligence")

    assert len(search_results) != 0, "The search results list is empty."
    print(f"Found {len(search_results)} search results")

    # Test cognee recall functionality
    search_results = await cognee.recall(
        "machine learning",
        SearchType.GRAPH_COMPLETION,
        datasets=[dataset_name],
    )
    assert len(search_results) != 0, "The search results list is empty."

    search_results = await cognee.recall(
        "vector databases",
        SearchType.CHUNKS,
        datasets=[dataset_name],
    )
    assert len(search_results) != 0, "The chunks search results list is empty."

    # Test search history
    user = await get_default_user()
    history = await get_history(user.id)
    assert len(history) > 0, "Search history is empty."

    # Test cleanup
    await cognee.prune.prune_data()
    assert not os.path.isdir(data_directory_path), "Local data files are not deleted"

    await cognee.prune.prune_system(metadata=True)
    
    # Verify database is empty after pruning
    # Add specific checks for your vector database here
    # Example:
    # vector_engine = await get_vector_engine_async()
    # collections = await vector_engine.client.list_collections()
    # assert len(collections) == 0, f"<EngineName> vector database is not empty"

    print("All tests passed successfully!")

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

This script ensures a basic end-to-end test, verifying:

• Your new adapter can be selected by cognee.
• Data can be remembered into the vector database flow.
• Recall functionality works correctly.
• The database is empty after a prune operation.

3. Create a Test Workflow

File: .github/workflows/vector_db_tests.yml

Add a new job to the existing vector database tests workflow:

  run-<engine_name>-tests:
    name: <EngineName> Tests
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, '<engine_name>') }}
    
    # Add services if your database needs to run as a container
    # services:
    #   <engine_name>:
    #     image: <your_db_image>:<version>
    #     ports:
    #       - <port>:<port>
    #     env:
    #       # Add any required environment variables
    
    steps:
      - name: Check out
        uses: actions/checkout@v4
        with:
          fetch-depth: 0

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: ${{ inputs.python-version }}

      - name: Install specific db dependency
        run: |
          poetry install -E <engine_name>

      # Add any setup steps if needed
      # - name: Wait for <EngineName> to be healthy
      #   run: |
      #     # Add health check commands

      - name: Run <EngineName> Tests
        env:
          ENV: 'dev'
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
          EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
          EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
          VECTOR_DB_URL: ${{ secrets.<ENGINE_NAME>_URL }}
          VECTOR_DB_KEY: ${{ secrets.<ENGINE_NAME>_API_KEY }}
        run: poetry run python ./cognee_community_vector_adapter_engine_name/examples/example.py
        working-directory: ./packages/vector/engine_name

5. Update pyproject.toml

Add your vector database client dependencies to the optional dependencies:

pyproject.toml:

[tool.poetry.dependencies]
python = "^3.11"
cognee
your-vector-db-client

[tool.poetry.extras]
...

## 6. Final Checklist

- [ ] Implement your `<EngineName>Adapter` in `packages/vector/<engine_name>/cognee_community_vector_adapter_<engine_name>/<engine_name>_adapter.py`.
- [ ] Add a register helper (`register.py`) and call it before configuring Cognee:

    ```python
    from cognee.infrastructure.databases.vector import use_vector_adapter

    from .<engine_name>_adapter import <EngineName>Adapter

    use_vector_adapter("engine_name", <EngineName>Adapter)

• Run the shared conformance suite in packages/shared/contract_suite/vector_contract.py against your adapter — it is the common contract every community vector adapter is expected to satisfy.

• If your backend can isolate storage per user + dataset, register a dataset-database handler from the same register.py:

  from cognee.infrastructure.databases.dataset_database_handler import (
      use_dataset_database_handler,
  )
  
  use_dataset_database_handler("engine_name", <EngineName>DatasetDatabaseHandler, "engine_name")
  

  Without a handler, your adapter can only be used with ENABLE_BACKEND_ACCESS_CONTROL=false.

• Create a test or example script test_<engine_name>.py or example.py that you can use in your test workflow.

• Create** a test workflow: .github/workflows/engine_name/test_<engine_name>.yml.

• Add required dependencies to pyproject.toml optional dependencies.

• Open a PR to verify that your new integration passes CI.

Additional Considerations

Error Handling

• Implement proper error handling for connection failures, timeouts, and API errors.
• Log meaningful error messages that help with debugging.
• Handle graceful degradation when the vector database is unavailable.

Performance

• Design for instance reuse: Cognee’s vector engine factory caches adapter instances keyed by their configuration parameters. Multiple calls with identical settings return the same adapter object. Design your adapter to be safe for reuse — avoid per-instance mutable state that cannot be safely shared across concurrent calls, and prefer lazy or thread-safe initialization patterns. When an entry is evicted (e.g. cache eviction or cache_clear), the factory calls your adapter’s close(). For capacity eviction the close is deferred until every leased reference to that instance is released, but the dataset-queue teardown and the idle reaper force-close immediately even while idle references are still held (holders transparently re-resolve to a fresh instance on next use) — so implement close() idempotently and safe to run while stale references remain; if it raises, the error is logged and swallowed rather than propagated to the caller.
• Consider implementing connection pooling if your vector database supports it.
• Never check out a second pooled connection while holding a session. If your adapter is pool-backed, resolve tables and metadata before opening the session that uses them — helpers like get_table() open their own connection, so calling one inside a session pins two connections per in-flight call and deadlocks a bounded pool once concurrency reaches its ceiling. Cognee’s own PGVectorAdapter follows this ordering in retrieve(), search(), and delete_data_points().
• Add proper async/await patterns to avoid blocking operations.
• Implement batch operations efficiently where possible.

Security

• Never log sensitive information like API keys or connection strings.
• Validate inputs to prevent injection attacks.
• Follow your vector database’s security best practices.

Documentation

• Add docstrings to all public methods.
• Document any specific configuration requirements.
• Include examples of how to use your vector database with cognee.

That’s all! This approach keeps cognee’s architecture flexible, allowing you to swap in any vector database provider easily. If you need more advanced functionality (e.g., custom indexes, filters, or advanced search capabilities), simply implement them in your adapter class following the same patterns.

Join the Conversation!

Have questions about creating custom tasks? Join our community to discuss implementation strategies and best practices!

Join the community