Repository: topoteretes/cognee
File: cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py (Python)

Code before:
    retry_if_not_exception_type,
    before_sleep_log,
)

from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
    LLMInterface,
)
from cognee.infrastructure.llm.config import get_llm_config

logger = get_logger()


class AnthropicAdapter(LLMInterface):
    """
    Adapter for interfacing with the Anthropic API, enabling structured output generation
    and prompt display.
    """

    name = "Anthropic"
    model: str
    default_instructor_mode = "anthropic_tools"

    def __init__(self, max_completion_tokens: int, model: str = None):
        import anthropic

        config_instructor_mode = get_llm_config().llm_instructor_mode
        instructor_mode = (
            config_instructor_mode if config_instructor_mode else self.default_instructor_mode
        )

        self.aclient = instructor.patch(
            create=anthropic.AsyncAnthropic(api_key=get_llm_config().llm_api_key).messages.create,
            mode=instructor.Mode(instructor_mode),
        )

        self.model = model
        self.max_completion_tokens = max_completion_tokens

    @retry(
        stop=stop_after_delay(128),
        wait=wait_exponential_jitter(2, 128),
        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
        before_sleep=before_sleep_log(logger, logging.DEBUG),
        reraise=True,
    )
    async def acreate_structured_output(
        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
    ) -> BaseModel:
        """
        Generate a response from a user query.


Review comment (suggestion):
We should have instructor mode be defined on a class level as input and if the input is None use the default_instructor_mode.

Check the get_llm_client function and forward instructor_mode there with the rest of the params

Code after:
    retry_if_not_exception_type,
    before_sleep_log,
)

from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
    LLMInterface,
)
from cognee.infrastructure.llm.config import get_llm_config

logger = get_logger()


class AnthropicAdapter(LLMInterface):
    """
    Adapter for interfacing with the Anthropic API, enabling structured output generation
    and prompt display.
    """

    name = "Anthropic"
    model: str
    default_instructor_mode = "anthropic_tools"

    def __init__(self, max_completion_tokens: int, model: str = None, instructor_mode: str = None):
        import anthropic

        self.instructor_mode = instructor_mode if instructor_mode else self.default_instructor_mode

        self.aclient = instructor.patch(
            create=anthropic.AsyncAnthropic(api_key=get_llm_config().llm_api_key).messages.create,
            mode=instructor.Mode(self.instructor_mode),
        )

        self.model = model
        self.max_completion_tokens = max_completion_tokens

    @retry(
        stop=stop_after_delay(128),
        wait=wait_exponential_jitter(2, 128),
        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
        before_sleep=before_sleep_log(logger, logging.DEBUG),
        reraise=True,
    )
    async def acreate_structured_output(
        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
    ) -> BaseModel:
        """
        Generate a response from a user query.

        Parameters:
        -----------