cognee/.github/workflows/search_db_tests.yml

name: Reusable Search DB Tests

permissions:
  contents: read

on:
  workflow_call:
    inputs:
      databases:
        required: false
        type: string
        default: "all"
        description: "Which vector databases to test (comma-separated list or 'all')"
    secrets:
      WEAVIATE_API_URL:
        required: false
      WEAVIATE_API_KEY:
        required: false

jobs:
  run-kuzu-lance-sqlite-search-tests:
    name: Search test for Kuzu/LanceDB/Sqlite
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'kuzu/lance/sqlite') }}
    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

      - name: Run Kuzu search 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: 'kuzu'
          VECTOR_DB_PROVIDER: 'lancedb'
          DB_PROVIDER: 'sqlite'
        run: poetry run python ./cognee/tests/test_search_db.py

  run-neo4j-lance-sqlite-search-tests:
    name: Search test for Neo4j/LanceDB/Sqlite
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'neo4j/lance/sqlite') }}
    services:
      neo4j:
        image: neo4j:5.11
        env:
          NEO4J_AUTH: neo4j/pleaseletmein
          NEO4J_PLUGINS: '["apoc","graph-data-science"]'
        ports:
          - 7474:7474
          - 7687:7687
        options: >-
          --health-cmd="cypher-shell -u neo4j -p pleaseletmein 'RETURN 1'"
          --health-interval=10s
          --health-timeout=5s
          --health-retries=5

    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 neo4j

      - name: Run Neo4j search 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: 'neo4j'
          VECTOR_DB_PROVIDER: 'lancedb'
          DB_PROVIDER: 'sqlite'
          GRAPH_DATABASE_URL: bolt://localhost:7687
          GRAPH_DATABASE_USERNAME: neo4j
          GRAPH_DATABASE_PASSWORD: pleaseletmein
        run: poetry run python ./cognee/tests/test_search_db.py

  run-kuzu-pgvector-postgres-search-tests:
      name: Search test for Kuzu/PGVector/Postgres
      runs-on: ubuntu-22.04
      if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'kuzu/pgvector/postgres') }}
      services:
        postgres:
          image: pgvector/pgvector:pg17
          env:
            POSTGRES_USER: cognee
            POSTGRES_PASSWORD: cognee
            POSTGRES_DB: cognee_db
          options: >-
            --health-cmd pg_isready
            --health-interval 10s
            --health-timeout 5s
            --health-retries 5
          ports:
            - 5432:5432
      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 dependencies
          run: poetry install -E postgres

        - name: Run Kuzu/PGVector/Postgres 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: 'kuzu'
            VECTOR_DB_PROVIDER: 'pgvector'
            DB_PROVIDER: 'postgres'
            DB_NAME: 'cognee_db'
            DB_HOST: '127.0.0.1'
            DB_PORT: 5432
            DB_USERNAME: cognee
            DB_PASSWORD: cognee
          run: poetry run python ./cognee/tests/test_search_db.py

  run-neo4j-pgvector-postgres-search-tests:
    name: Search test for Neo4j/PGVector/Postgres
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'neo4j/pgvector/postgres') }}
    services:
      neo4j:
        image: neo4j:5.11
        env:
          NEO4J_AUTH: neo4j/pleaseletmein
          NEO4J_PLUGINS: '["apoc","graph-data-science"]'
        ports:
          - 7474:7474
          - 7687:7687
        options: >-
          --health-cmd="cypher-shell -u neo4j -p pleaseletmein 'RETURN 1'"
          --health-interval=10s
          --health-timeout=5s
          --health-retries=5
      postgres:
        image: pgvector/pgvector:pg17
        env:
          POSTGRES_USER: cognee
          POSTGRES_PASSWORD: cognee
          POSTGRES_DB: cognee_db
        ports:
          - 5432:5432
        options: >-
          --health-cmd pg_isready
          --health-interval 10s
          --health-timeout 5s
          --health-retries=5
    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 dependencies
        run: |
          poetry install -E neo4j -E postgres

      - name: Run Neo4j + PGVector + Postgres search 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: 'neo4j'
          VECTOR_DB_PROVIDER: 'pgvector'
          DB_PROVIDER: 'postgres'
          GRAPH_DATABASE_URL: bolt://localhost:7687
          GRAPH_DATABASE_USERNAME: neo4j
          GRAPH_DATABASE_PASSWORD: pleaseletmein
          DB_NAME: cognee_db
          DB_HOST: 127.0.0.1
          DB_PORT: 5432
          DB_USERNAME: cognee
          DB_PASSWORD: cognee
        run: poetry run python ./cognee/tests/test_search_db.py