cognee/.github/workflows/db_examples_tests.yml

name: Reusable DB Examples Tests

on:
  workflow_call:
    inputs:
      databases:
        required: false
        type: string
        default: "all"
        description: "Which databases to run (comma-separated or 'all')"
      python-version:
        required: false
        type: string
        default: "3.11.x"
    secrets:
      LLM_MODEL:
        required: true
      LLM_ENDPOINT:
        required: true
      LLM_API_KEY:
        required: true
      LLM_API_VERSION:
        required: true
      EMBEDDING_MODEL:
        required: true
      EMBEDDING_ENDPOINT:
        required: true
      EMBEDDING_API_KEY:
        required: true
      EMBEDDING_API_VERSION:
        required: true
      QDRANT_API_URL:
        required: false
      QDRANT_API_KEY:
        required: false
      WEAVIATE_API_URL:
        required: false
      WEAVIATE_API_KEY:
        required: false
      POSTGRES_PASSWORD:
        required: false
      NEO4J_API_URL:
        required: false
      NEO4J_API_KEY:
        required: false




jobs:
  run-db-example-neo4j:
    name: "Neo4j DB Example Test"
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'neo4j') }}
    steps:
      - name: Check out
        uses: actions/checkout@master

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

      - name: Install Neo4j extra
        run: |
          poetry install -E neo4j

      - name: Run Neo4j Example
        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"
          GRAPH_DATABASE_URL: ${{ secrets.NEO4J_API_URL }}
          GRAPH_DATABASE_USERNAME: "neo4j"
          GRAPH_DATABASE_PASSWORD: ${{ secrets.NEO4J_API_KEY }}
        run: |
          poetry run python examples/database_examples/neo4j_example.py

  run-db-example-kuzu:
    name: "Kuzu DB Example Test"
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'kuzu') }}

    steps:
      - name: Check out
        uses: actions/checkout@v4

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

      - name: Install Kuzu extra
        run: |
          poetry install -E kuzu

      - name: Run Kuzu Example
        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"
        run: |
          poetry run python examples/database_examples/kuzu_example.py

  run-db-example-milvus:
    name: "Milvus DB Example Test"
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'milvus') }}

    steps:
      - name: Check out
        uses: actions/checkout@v4

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

      - name: Install Milvus extra
        run: |
          poetry install -E milvus

      - name: Run Milvus Example
        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 }}
        run: |
          poetry run python examples/database_examples/milvus_example.py

  run-db-example-weaviate:
    name: "Weaviate DB Example Test"
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'weaviate') }}
    steps:
      - name: Check out
        uses: actions/checkout@v4

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

      - name: Install Weaviate extra
        run: |
          poetry install -E weaviate

      - name: Run Weaviate Example
        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.WEAVIATE_API_URL }}
          VECTOR_DB_KEY: ${{ secrets.WEAVIATE_API_KEY }}
        run: |
          poetry run python examples/database_examples/weaviate_example.py

  run-db-example-qdrant:
    name: "Qdrant DB Example Test"
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'qdrant') }}
    defaults:
        run:
          shell: bash
  
    steps:
      - name: Check out
        uses: actions/checkout@master

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

      - name: Install Qdrant extra
        run: |
          poetry install -E qdrant

      - name: Run Qdrant Example
        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.QDRANT_API_URL }}
          VECTOR_DB_KEY: ${{ secrets.QDRANT_API_KEY }}
        run: |
          poetry run python examples/database_examples/qdrant_example.py

  run-db-example-pgvector:
    name: "PostgreSQL PGVector DB Example Test"
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'postgres') }}
    services:
      postgres:
        image: pgvector/pgvector:pg17
        env:
          POSTGRES_USER: cognee
          POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
          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 PGVector extra
        run: |
          poetry install -E postgres

      - name: Run PGVector Example
        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 }}
        run: |
          poetry run python examples/database_examples/pgvector_example.py