cognee/.github/workflows/test_llms.yml

name: LLM Test Suites

permissions:
  contents: read

on:
  workflow_call:

env:
  RUNTIME__LOG_LEVEL: ERROR
  ENV: 'dev'

jobs:
  test-gemini:
    name: Run Gemini Test
    runs-on: ubuntu-22.04
    steps:
      - name: Check out repository
        uses: actions/checkout@v4

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: '3.11.x'

      - name: Run Gemini Simple Example
        env:
          LLM_PROVIDER: "gemini"
          LLM_API_KEY: ${{ secrets.GEMINI_API_KEY }}
          LLM_MODEL: "gemini/gemini-2.0-flash"
          EMBEDDING_PROVIDER: "gemini"
          EMBEDDING_API_KEY: ${{ secrets.GEMINI_API_KEY }}
          EMBEDDING_MODEL: "gemini/text-embedding-004"
          EMBEDDING_DIMENSIONS: "768"
          EMBEDDING_MAX_TOKENS: "8076"
        run: uv run python ./examples/python/simple_example.py

  test-fastembed:
    name: Run Fastembed Test
    runs-on: ubuntu-22.04
    steps:
      - name: Check out repository
        uses: actions/checkout@v4

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: '3.11.x'

      - name: Run Fastembed Simple Example
        env:
          LLM_PROVIDER: "openai"
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_PROVIDER: "fastembed"
          EMBEDDING_MODEL: "sentence-transformers/all-MiniLM-L6-v2"
          EMBEDDING_DIMENSIONS: "384"
          EMBEDDING_MAX_TOKENS: "256"
        run: uv run python ./examples/python/simple_example.py

  test-openrouter:
    name: Run OpenRouter Test
    runs-on: ubuntu-22.04
    steps:
      - name: Check out repository
        uses: actions/checkout@v4

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: '3.11.x'

      - name: Run OpenRouter Simple Example
        env:
          LLM_PROVIDER: "custom"
          LLM_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
          LLM_MODEL: "openrouter/x-ai/grok-code-fast-1"
          LLM_ENDPOINT: "https://openrouter.ai/api/v1"
          EMBEDDING_PROVIDER: "openai"
          EMBEDDING_API_KEY: ${{ secrets.OPENAI_API_KEY }}
          EMBEDDING_MODEL: "openai/text-embedding-3-large"
          EMBEDDING_DIMENSIONS: "3072"
          EMBEDDING_MAX_TOKENS: "8191"
        run: uv run python ./examples/python/simple_example.py

  test-bedrock-api-key:
    name: Run Bedrock API Key Test
    runs-on: ubuntu-22.04
    steps:
      - name: Check out repository
        uses: actions/checkout@v4

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: '3.11.x'
          extra-dependencies: "aws"

      - name: Run Bedrock API Key Simple Example
        env:
          LLM_PROVIDER: "bedrock"
          LLM_API_KEY: ${{ secrets.BEDROCK_API_KEY }}
          LLM_MODEL: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
          LLM_MAX_TOKENS: "16384"
          AWS_REGION_NAME: "eu-west-1"
          EMBEDDING_PROVIDER: "bedrock"
          EMBEDDING_API_KEY: ${{ secrets.BEDROCK_API_KEY }}
          EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0"
          EMBEDDING_DIMENSIONS: "1024"
          EMBEDDING_MAX_TOKENS: "8191"
        run: uv run python ./examples/python/simple_example.py

  test-bedrock-aws-credentials:
    name: Run Bedrock AWS Credentials Test
    runs-on: ubuntu-22.04
    steps:
      - name: Check out repository
        uses: actions/checkout@v4

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: '3.11.x'
          extra-dependencies: "aws"

      - name: Run Bedrock AWS Credentials Simple Example
        env:
          LLM_PROVIDER: "bedrock"
          LLM_MODEL: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
          LLM_MAX_TOKENS: "16384"
          AWS_REGION_NAME: "eu-west-1"
          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
          EMBEDDING_PROVIDER: "bedrock"
          EMBEDDING_API_KEY: ${{ secrets.BEDROCK_API_KEY }}
          EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0"
          EMBEDDING_DIMENSIONS: "1024"
          EMBEDDING_MAX_TOKENS: "8191"
        run: uv run python ./examples/python/simple_example.py

  test-bedrock-aws-profile:
    name: Run Bedrock AWS Profile Test
    runs-on: ubuntu-22.04
    steps:
      - name: Check out repository
        uses: actions/checkout@v4

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: '3.11.x'
          extra-dependencies: "aws"

      - name: Configure AWS Profile
        run: |
          mkdir -p ~/.aws
          cat > ~/.aws/credentials << EOF
          [bedrock-test]
          aws_access_key_id = ${{ secrets.AWS_ACCESS_KEY_ID }}
          aws_secret_access_key = ${{ secrets.AWS_SECRET_ACCESS_KEY }}
          EOF

      - name: Run Bedrock AWS Profile Simple Example
        env:
          LLM_PROVIDER: "bedrock"
          LLM_MODEL: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
          LLM_MAX_TOKENS: "16384"
          AWS_PROFILE_NAME: "bedrock-test"
          AWS_REGION_NAME: "eu-west-1"
          EMBEDDING_PROVIDER: "bedrock"
          EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0"
          EMBEDDING_DIMENSIONS: "1024"
          EMBEDDING_MAX_TOKENS: "8191"
        run: uv run python ./examples/python/simple_example.py