diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/config.py index 2e300dc0c..7c0221993 100644 --- a/cognee/infrastructure/llm/config.py +++ b/cognee/infrastructure/llm/config.py @@ -74,6 +74,41 @@ class LLMConfig(BaseSettings): model_config = SettingsConfigDict(env_file=".env", extra="allow") + @model_validator(mode="after") + def strip_quotes_from_strings(self) -> "LLMConfig": + """ + Strip surrounding quotes from specific string fields that often come from + environment variables with extra quotes (e.g., via Docker's --env-file). + + Only applies to known config keys where quotes are invalid or cause issues. + """ + string_fields_to_strip = [ + "llm_api_key", + "llm_endpoint", + "llm_api_version", + "baml_llm_api_key", + "baml_llm_endpoint", + "baml_llm_api_version", + "fallback_api_key", + "fallback_endpoint", + "fallback_model", + "llm_provider", + "llm_model", + "baml_llm_provider", + "baml_llm_model", + ] + + cls = self.__class__ + for field_name in string_fields_to_strip: + if field_name not in cls.model_fields: + continue + value = getattr(self, field_name, None) + if isinstance(value, str) and len(value) >= 2: + if value[0] == value[-1] and value[0] in ("'", '"'): + setattr(self, field_name, value[1:-1]) + + return self + def model_post_init(self, __context) -> None: """Initialize the BAML registry after the model is created.""" # Check if BAML is selected as structured output framework but not available diff --git a/cognee/tests/unit/infrastructure/llm/test_llm_config.py b/cognee/tests/unit/infrastructure/llm/test_llm_config.py new file mode 100644 index 000000000..0c024db2c --- /dev/null +++ b/cognee/tests/unit/infrastructure/llm/test_llm_config.py @@ -0,0 +1,46 @@ +import pytest + +from cognee.infrastructure.llm.config import LLMConfig + + +def test_strip_quotes_from_strings(): + """ + Test if the LLMConfig.strip_quotes_from_strings model validator behaves as expected. + """ + config = LLMConfig( + # Strings with surrounding double quotes ("value" → value) + llm_api_key='"double_value"', + # Strings with surrounding single quotes ('value' → value) + llm_endpoint="'single_value'", + # Strings without quotes (value → value) + llm_api_version="no_quotes_value", + # Empty quoted strings ("" → empty string) + fallback_model='""', + # None values (should remain None) + baml_llm_api_key=None, + # Mixed quotes ("value' → unchanged) + fallback_endpoint="\"mixed_quote'", + # Strings with internal quotes ("internal\"quotes" → internal"quotes") + baml_llm_model='"internal"quotes"', + ) + + # Strings with surrounding double quotes ("value" → value) + assert config.llm_api_key == "double_value" + + # Strings with surrounding single quotes ('value' → value) + assert config.llm_endpoint == "single_value" + + # Strings without quotes (value → value) + assert config.llm_api_version == "no_quotes_value" + + # Empty quoted strings ("" → empty string) + assert config.fallback_model == "" + + # None values (should remain None) + assert config.baml_llm_api_key is None + + # Mixed quotes ("value' → unchanged) + assert config.fallback_endpoint == "\"mixed_quote'" + + # Strings with internal quotes ("internal\"quotes" → internal"quotes") + assert config.baml_llm_model == 'internal"quotes'