Adding a New Graph Database to cognee

"""Adapter for <engine_name> graph database."""
 
import json
import asyncio
from typing import Dict, Any, List, Optional, Tuple
 
from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 
logger = get_logger()
 
class <EngineName>Adapter(GraphDBInterface):
    """Adapter for <engine_name> graph database operations."""
 
    def __init__(
        self,
        graph_database_url: str = "",
        graph_database_username: Optional[str] = None,
        graph_database_password: Optional[str] = None,
    ):
        self.graph_database_url = graph_database_url
        self.graph_database_username = graph_database_username
        self.graph_database_password = graph_database_password
        self.connection = None
 
    async def query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Tuple]:
        """Execute an async query.
        
        If your graph database library provides an async SDK, call it directly here.
        If it only provides a synchronous client, you can run the call via
        `loop.run_in_executor()` or a similar technique to avoid blocking the event loop.
        """
 
        loop = asyncio.get_running_loop()
        params = params or {}
 
        def blocking_query():
            try:
                # Example usage with your driver
                # cursor = self.connection.execute(query, params)
                # results = cursor.fetchall()
                return []
            except Exception as e:
                logger.error(f"<engine_name> query execution failed: {e}")
                raise
 
        return await loop.run_in_executor(self.executor, blocking_query)
 
    # -- Example: Add a node
    async def add_node(self, node_data: Any) -> None:
        """Add a single node to <engine_name>."""
        # Implement logic:
        # 1. Extract relevant fields (id, text, type, properties, etc.).
        # 2. Construct a CREATE query.
        # 3. Call self.query(query_str, params).
        pass
 
    # -- Example: Retrieve all nodes/edges
    async def get_graph_data(self) -> Tuple[List, List]:
        """Retrieve all nodes and edges from <engine_name>."""
        # Return (nodes, edges) where each node is `(node_id, properties_dict)`
        # and each edge is `(source_id, target_id, relationship_label, properties_dict)`.
        return ([], [])
 
    # -- Additional methods (delete_node, add_edge, etc.) ...

import sys
import asyncio
import pathlib
from os import path
 
# NOTE: Importing the register module we let cognee know it can use the Networkx adapter
import packages.graph.networkx.register
 
async def main():
    from cognee import config, prune, add, cognify, search, SearchType
 
    system_path = pathlib.Path(__file__).parent
    config.system_root_directory(path.join(system_path, ".cognee-system"))
    config.data_root_directory(path.join(system_path, ".cognee-data"))
 
    config.set_graph_db_config({
        "graph_database_provider": "engine-name",
    })
 
    await prune.prune_data()
    await prune.prune_system()
 
    text = """
    Natural language processing (NLP) is an interdisciplinary
    subfield of computer science and information retrieval.
    """
 
    await add(text)
 
    await cognify()
 
    query_text = "Tell me about NLP"
 
    search_results = await search(query_type=SearchType.GRAPH_COMPLETION, query_text=query_text)
 
    for result_text in search_results:
        print(result_text)
 
if __name__ == "__main__":
    asyncio.run(main())

name: test | <engine_name>
 
on:
  workflow_dispatch:
  pull_request:
    types: [labeled, synchronize]
 
concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true
 
env:
  RUNTIME__LOG_LEVEL: ERROR
 
jobs:
  run_<engine_name>_integration_test:
    name: test
    runs-on: ubuntu-22.04
 
    defaults:
      run:
        shell: bash
 
    steps:
      - name: Check out
        uses: actions/checkout@master
 
      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11.x'
 
      - name: Install Poetry
        uses: snok/install-poetry@v1.4.1
        with:
          virtualenvs-create: true
          virtualenvs-in-project: true
          installer-parallel: true
 
      - name: Install dependencies
        # If your pyproject.toml has an extra named '<engine_name>', use:
        run: poetry install -E <engine_name> --no-interaction
 
      - name: Run <engine_name> 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 }}
          GRAPH_DATABASE_PROVIDER: ''
          GRAPH_DATABASE_URL: ''
          GRAPH_DATABASE_PASSWORD: ''
        run: poetry run python ./cognee_community_graph_adapter_engine_name/examples/example.py
        working-directory: ./packages/graph/engine_name

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