fix: Add upper to BAML structured output name

This commit is contained in:
Igor Ilic 2025-08-05 21:24:58 +02:00
parent 9ca050a70d
commit 0941c64ba7

View file

@ -107,7 +107,7 @@ class LLMGateway:
@staticmethod @staticmethod
def extract_code_summary(content: str) -> Coroutine: def extract_code_summary(content: str) -> Coroutine:
llm_config = get_llm_config() llm_config = get_llm_config()
if llm_config.structured_output_framework == "BAML": if llm_config.structured_output_framework.upper() == "BAML":
from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import ( from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
extract_code_summary, extract_code_summary,
) )
@ -123,7 +123,7 @@ class LLMGateway:
@staticmethod @staticmethod
def extract_summary(content: str, response_model: Type[BaseModel]) -> Coroutine: def extract_summary(content: str, response_model: Type[BaseModel]) -> Coroutine:
llm_config = get_llm_config() llm_config = get_llm_config()
if llm_config.structured_output_framework == "BAML": if llm_config.structured_output_framework.upper() == "BAML":
from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import ( from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
extract_summary, extract_summary,
) )