From e0d48c043a1594c567135a1e73cd209c1d07eba1 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Tue, 25 Nov 2025 12:58:07 +0100 Subject: [PATCH] fix: fixes to adapter and tests --- .github/workflows/test_bedrock_api_key.yml | 28 ------ .../test_bedrock_aws_credentials.yml | 29 ------ .../workflows/test_bedrock_aws_profile.yml | 37 -------- .github/workflows/test_llms.yml | 88 +++++++++++++++++++ .github/workflows/test_suites.yml | 24 ----- .../llm/bedrock/__init__.py | 1 - .../litellm_instructor/llm/bedrock/adapter.py | 35 ++++---- .../litellm_instructor/llm/get_llm_client.py | 7 +- cognee/modules/settings/get_settings.py | 12 +-- 9 files changed, 117 insertions(+), 144 deletions(-) delete mode 100644 .github/workflows/test_bedrock_api_key.yml delete mode 100644 .github/workflows/test_bedrock_aws_credentials.yml delete mode 100644 .github/workflows/test_bedrock_aws_profile.yml diff --git a/.github/workflows/test_bedrock_api_key.yml b/.github/workflows/test_bedrock_api_key.yml deleted file mode 100644 index 3f5ea94b3..000000000 --- a/.github/workflows/test_bedrock_api_key.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: test | bedrock | api key - -on: - workflow_call: - -jobs: - 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' - - - name: Run Bedrock API Key Test - env: - LLM_PROVIDER: "bedrock" - LLM_API_KEY: ${{ secrets.BEDROCK_API_KEY }} - LLM_MODEL: "us.anthropic.claude-3-5-sonnet-20241022-v2:0" - AWS_REGION_NAME: "us-east-1" - EMBEDDING_PROVIDER: "bedrock" - EMBEDDING_MODEL: "amazon.titan-embed-text-v1" - EMBEDDING_DIMENSIONS: "1536" - run: poetry run python ./examples/python/simple_example.py diff --git a/.github/workflows/test_bedrock_aws_credentials.yml b/.github/workflows/test_bedrock_aws_credentials.yml deleted file mode 100644 index c086dceb3..000000000 --- a/.github/workflows/test_bedrock_aws_credentials.yml +++ /dev/null @@ -1,29 +0,0 @@ -name: test | bedrock | aws credentials - -on: - workflow_call: - -jobs: - 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' - - - name: Run Bedrock AWS Credentials Test - env: - LLM_PROVIDER: "bedrock" - LLM_MODEL: "us.anthropic.claude-3-5-sonnet-20240620-v1:0" - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - AWS_REGION_NAME: "us-east-1" - EMBEDDING_PROVIDER: "cohere" - EMBEDDING_MODEL: "cohere.embed-english-v3" - EMBEDDING_DIMENSIONS: "1024" - run: poetry run python ./examples/python/simple_example.py diff --git a/.github/workflows/test_bedrock_aws_profile.yml b/.github/workflows/test_bedrock_aws_profile.yml deleted file mode 100644 index aa15074e1..000000000 --- a/.github/workflows/test_bedrock_aws_profile.yml +++ /dev/null @@ -1,37 +0,0 @@ -name: test | bedrock | aws profile - -on: - workflow_call: - -jobs: - 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' - - - 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 Test - env: - LLM_PROVIDER: "bedrock" - LLM_MODEL: "us.anthropic.claude-3-5-haiku-20241022-v1:0" - AWS_PROFILE_NAME: "bedrock-test" - AWS_REGION_NAME: "us-east-1" - EMBEDDING_PROVIDER: "bedrock" - EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0" - EMBEDDING_DIMENSIONS: "1024" - run: poetry run python ./examples/python/simple_example.py diff --git a/.github/workflows/test_llms.yml b/.github/workflows/test_llms.yml index 6b0221309..0cbbc7b3a 100644 --- a/.github/workflows/test_llms.yml +++ b/.github/workflows/test_llms.yml @@ -84,3 +84,91 @@ jobs: 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 Test + env: + LLM_PROVIDER: "bedrock" + LLM_API_KEY: ${{ secrets.BEDROCK_API_KEY }} + LLM_MODEL: "eu.amazon.nova-lite-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: poetry 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 API Key Test + env: + LLM_PROVIDER: "bedrock" + LLM_MODEL: "eu.amazon.nova-lite-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: poetry 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' + + - 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 Test + env: + LLM_PROVIDER: "bedrock" + LLM_MODEL: "eu.amazon.nova-lite-v1:0" + 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: poetry run python ./examples/python/simple_example.py \ No newline at end of file diff --git a/.github/workflows/test_suites.yml b/.github/workflows/test_suites.yml index 9f2767faf..be1e354fc 100644 --- a/.github/workflows/test_suites.yml +++ b/.github/workflows/test_suites.yml @@ -139,24 +139,6 @@ jobs: uses: ./.github/workflows/test_llms.yml secrets: inherit - bedrock-tests: - name: Bedrock Tests - needs: [basic-tests, e2e-tests] - uses: ./.github/workflows/test_bedrock_api_key.yml - secrets: inherit - - bedrock-aws-credentials-tests: - name: Bedrock AWS Credentials Tests - needs: [basic-tests, e2e-tests] - uses: ./.github/workflows/test_bedrock_aws_credentials.yml - secrets: inherit - - bedrock-aws-profile-tests: - name: Bedrock AWS Profile Tests - needs: [basic-tests, e2e-tests] - uses: ./.github/workflows/test_bedrock_aws_profile.yml - secrets: inherit - # Ollama tests moved to the end ollama-tests: name: Ollama Tests @@ -193,9 +175,6 @@ jobs: db-examples-tests, mcp-test, llm-tests, - bedrock-tests, - bedrock-aws-credentials-tests, - bedrock-aws-profile-tests, ollama-tests, relational-db-migration-tests, docker-compose-test, @@ -218,9 +197,6 @@ jobs: "${{ needs.db-examples-tests.result }}" == "success" && "${{ needs.relational-db-migration-tests.result }}" == "success" && "${{ needs.llm-tests.result }}" == "success" && - "${{ needs.bedrock-tests.result }}" == "success" && - "${{ needs.bedrock-aws-credentials-tests.result }}" == "success" && - "${{ needs.bedrock-aws-profile-tests.result }}" == "success" && "${{ needs.docker-compose-test.result }}" == "success" && "${{ needs.docker-ci-test.result }}" == "success" && "${{ needs.ollama-tests.result }}" == "success" ]]; then diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/__init__.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/__init__.py index 6fb964a82..ad7cdf994 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/__init__.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/__init__.py @@ -3,4 +3,3 @@ from .adapter import BedrockAdapter __all__ = ["BedrockAdapter"] - diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/adapter.py index 66f484164..c461a0886 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/adapter.py @@ -1,18 +1,19 @@ import litellm import instructor -from typing import Type, Optional +from typing import Type from pydantic import BaseModel from litellm.exceptions import ContentPolicyViolationError from instructor.exceptions import InstructorRetryException -from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import ( LLMInterface, ) -from cognee.infrastructure.llm.exceptions import ContentPolicyFilterError +from cognee.infrastructure.llm.exceptions import ( + ContentPolicyFilterError, + MissingSystemPromptPathError, +) from cognee.infrastructure.files.storage.s3_config import get_s3_config -from cognee.infrastructure.files.utils.open_data_file import open_data_file from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import ( rate_limit_async, rate_limit_sync, @@ -35,6 +36,7 @@ class BedrockAdapter(LLMInterface): name = "Bedrock" model: str api_key: str + default_instructor_mode = "json_schema_mode" MAX_RETRIES = 5 @@ -42,23 +44,23 @@ class BedrockAdapter(LLMInterface): self, model: str, api_key: str = None, - max_tokens: int = 16384, + max_completion_tokens: int = 16384, streaming: bool = False, + instructor_mode: str = None, ): - self.aclient = instructor.from_litellm(litellm.acompletion) + self.instructor_mode = instructor_mode if instructor_mode else self.default_instructor_mode + + self.aclient = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode(self.instructor_mode)) self.client = instructor.from_litellm(litellm.completion) self.model = model self.api_key = api_key - self.max_tokens = max_tokens + self.max_completion_tokens = max_completion_tokens self.streaming = streaming def _create_bedrock_request( self, text_input: str, system_prompt: str, response_model: Type[BaseModel] ) -> dict: - """Create Bedrock request with authentication and enhanced JSON formatting.""" - enhanced_system_prompt = f"""{system_prompt} - -IMPORTANT: You must respond with valid JSON only. Do not include any text before or after the JSON. The response must be a valid JSON object that can be parsed directly.""" + """Create Bedrock request with authentication.""" request_params = { "model": self.model, @@ -66,11 +68,11 @@ IMPORTANT: You must respond with valid JSON only. Do not include any text before "drop_params": True, "messages": [ {"role": "user", "content": text_input}, - {"role": "system", "content": enhanced_system_prompt}, + {"role": "system", "content": system_prompt}, ], "response_model": response_model, "max_retries": self.MAX_RETRIES, - "max_tokens": self.max_tokens, + "max_completion_tokens": self.max_completion_tokens, "stream": self.streaming, } @@ -87,9 +89,10 @@ IMPORTANT: You must respond with valid JSON only. Do not include any text before elif s3_config.aws_profile_name: request_params["aws_profile_name"] = s3_config.aws_profile_name + if s3_config.aws_region: + request_params["aws_region_name"] = s3_config.aws_region + # Add optional parameters - if s3_config.aws_region_name: - request_params["aws_region_name"] = s3_config.aws_region_name if s3_config.aws_bedrock_runtime_endpoint: request_params["aws_bedrock_runtime_endpoint"] = s3_config.aws_bedrock_runtime_endpoint @@ -137,7 +140,7 @@ IMPORTANT: You must respond with valid JSON only. Do not include any text before if not text_input: text_input = "No user input provided." if not system_prompt: - raise InvalidValueError(message="No system prompt path provided.") + raise MissingSystemPromptPathError() system_prompt = LLMGateway.read_query_prompt(system_prompt) formatted_prompt = ( diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py index 086fd84de..954d85c1d 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py @@ -172,8 +172,8 @@ def get_llm_client(raise_api_key_error: bool = True): ) elif provider == LLMProvider.BEDROCK: - if llm_config.llm_api_key is None and raise_api_key_error: - raise LLMAPIKeyNotSetError() + # if llm_config.llm_api_key is None and raise_api_key_error: + # raise LLMAPIKeyNotSetError() from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.bedrock.adapter import ( BedrockAdapter, @@ -182,8 +182,9 @@ def get_llm_client(raise_api_key_error: bool = True): return BedrockAdapter( model=llm_config.llm_model, api_key=llm_config.llm_api_key, - max_tokens=max_completion_tokens, + max_completion_tokens=max_completion_tokens, streaming=llm_config.llm_streaming, + instructor_mode=llm_config.llm_instructor_mode.lower(), ) else: diff --git a/cognee/modules/settings/get_settings.py b/cognee/modules/settings/get_settings.py index 7e58e981f..37093bb35 100644 --- a/cognee/modules/settings/get_settings.py +++ b/cognee/modules/settings/get_settings.py @@ -164,16 +164,16 @@ def get_settings() -> SettingsDict: ], "bedrock": [ { - "value": "us.anthropic.claude-3-5-sonnet-20241022-v2:0", - "label": "Claude 3.5 Sonnet", + "value": "eu.anthropic.claude-haiku-4-5-20251001-v1:0", + "label": "Claude 4.5 Sonnet", }, { - "value": "us.anthropic.claude-3-5-haiku-20241022-v1:0", - "label": "Claude 3.5 Haiku", + "value": "eu.anthropic.claude-haiku-4-5-20251001-v1:0", + "label": "Claude 4.5 Haiku", }, { - "value": "us.anthropic.claude-3-5-sonnet-20240620-v1:0", - "label": "Claude 3.5 Sonnet (June)", + "value": "eu.amazon.nova-lite-v1:0", + "label": "Amazon Nova Lite", }, ], },