gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: Resolve issue with BAML rate limit handling

Browse source
This commit is contained in:
Igor Ilic 2025-11-19 23:04:19 +01:00
parent 487635b71b
commit 7360729db1
1 changed files with 16 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/acreate_structured_output.py
View file

@ -1,7 +1,15 @@
import asyncio import asyncio
from typing import Type from typing import Type
from cognee.shared.logging_utils import get_logger from pydantic import BaseModel
from tenacity import (
retry,
stop_after_delay,
wait_exponential_jitter,
retry_if_not_exception_type,
before_sleep_log,
)
from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.llm.config import get_llm_config from cognee.infrastructure.llm.config import get_llm_config
from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction.create_dynamic_baml_type import ( from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction.create_dynamic_baml_type import (
create_dynamic_baml_type, create_dynamic_baml_type,
@ -10,12 +18,17 @@ from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type
TypeBuilder, TypeBuilder,
) )
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client import b from cognee.infrastructure.llm.structured_output_framework.baml.baml_client import b
from pydantic import BaseModel import logging
logger = get_logger() logger = get_logger()
@retry(
stop=stop_after_delay(128),
wait=wait_exponential_jitter(2, 128),
before_sleep=before_sleep_log(logger, logging.DEBUG),
reraise=True,
)
async def acreate_structured_output( async def acreate_structured_output(
text_input: str, system_prompt: str, response_model: Type[BaseModel] text_input: str, system_prompt: str, response_model: Type[BaseModel]
): ):