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')"

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: Dependencies already installed
        run: echo "Dependencies already installed in setup"

      - 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: uv 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') }}

    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: Setup Neo4j with GDS
        uses: ./.github/actions/setup_neo4j
        id: neo4j

      - name: Dependencies already installed
        run: echo "Dependencies already installed in setup"

      - 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: ${{ steps.neo4j.outputs.neo4j-url }}
          GRAPH_DATABASE_USERNAME: ${{ steps.neo4j.outputs.neo4j-username }}
          GRAPH_DATABASE_PASSWORD: ${{ steps.neo4j.outputs.neo4j-password }}
        run: uv 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 }}
            extra-dependencies: "postgres"

        - name: Dependencies already installed
          run: echo "Dependencies already installed in setup"

        - 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: uv 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:
      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 }}
          extra-dependencies: "postgres"

      - name: Setup Neo4j with GDS
        uses: ./.github/actions/setup_neo4j
        id: neo4j

      - name: Dependencies already installed
        run: echo "Dependencies already installed in setup"

      - 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: ${{ steps.neo4j.outputs.neo4j-url }}
          GRAPH_DATABASE_USERNAME: ${{ steps.neo4j.outputs.neo4j-username }}
          GRAPH_DATABASE_PASSWORD: ${{ steps.neo4j.outputs.neo4j-password }}
          DB_NAME: cognee_db
          DB_HOST: 127.0.0.1
          DB_PORT: 5432
          DB_USERNAME: cognee
          DB_PASSWORD: cognee
        run: uv run python ./cognee/tests/test_search_db.py