Llm

This module contains functionality related to the the llm module for augmentation.components.llms.lite_llm.

Llm

ConfigurableLiteLLM

Bases: LiteLLM

Extended LiteLLM that allows context_window and num_output override.

This subclass overrides the metadata property to allow configuration-based context_window and num_output values, rather than relying solely on LiteLLM's model registry defaults.

Source code in src/augmentation/components/llms/lite_llm/llm.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class ConfigurableLiteLLM(LiteLLM):
    """
    Extended LiteLLM that allows context_window and num_output override.

    This subclass overrides the metadata property to allow configuration-based
    context_window and num_output values, rather than relying solely on LiteLLM's
    model registry defaults.
    """

    def __init__(
        self,
        *args,
        context_window: int | None = None,
        num_output: int | None = None,
        **kwargs,
    ):
        """Initialize with optional context_window and num_output overrides."""
        super().__init__(*args, **kwargs)
        self._context_window_override = context_window
        self._num_output_override = num_output

    @property
    def metadata(self) -> LLMMetadata:
        """
        Get LLM metadata with configuration overrides applied.

        Returns:
            LLMMetadata with context_window and num_output from configuration
            if specified, otherwise uses parent class defaults.
        """
        # Get default metadata from parent
        default_metadata = super().metadata

        # Apply overrides if configured
        context_window = (
            self._context_window_override
            if self._context_window_override is not None
            else default_metadata.context_window
        )

        num_output = (
            self._num_output_override
            if self._num_output_override is not None
            else default_metadata.num_output
        )

        return LLMMetadata(
            context_window=context_window,
            num_output=num_output,
            is_chat_model=default_metadata.is_chat_model,
            is_function_calling_model=default_metadata.is_function_calling_model,
            model_name=default_metadata.model_name,
            system_role=default_metadata.system_role,
        )

metadata property

Get LLM metadata with configuration overrides applied.

Returns:
  • LLMMetadata –

    LLMMetadata with context_window and num_output from configuration

  • LLMMetadata –

    if specified, otherwise uses parent class defaults.

__init__(*args, context_window=None, num_output=None, **kwargs)

Initialize with optional context_window and num_output overrides.

Source code in src/augmentation/components/llms/lite_llm/llm.py 
21
22
23
24
25
26
27
28
29
30
31
def __init__(
    self,
    *args,
    context_window: int | None = None,
    num_output: int | None = None,
    **kwargs,
):
    """Initialize with optional context_window and num_output overrides."""
    super().__init__(*args, **kwargs)
    self._context_window_override = context_window
    self._num_output_override = num_output

LiteLLMFactory

Bases: SingletonFactory

Factory class for creating LiteLLM language model instances.

This class implements the Singleton pattern to ensure only one instance of an LiteLLM model with a specific configuration exists in the application. It uses LiteLLMConfiguration to configure the model parameters.

Attributes:
  • _configuration_class (Type) –

    Type of the configuration class used for creating LiteLLM model instances.
Source code in src/augmentation/components/llms/lite_llm/llm.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
class LiteLLMFactory(SingletonFactory):
    """
    Factory class for creating LiteLLM language model instances.

    This class implements the Singleton pattern to ensure only one instance
    of an LiteLLM model with a specific configuration exists in the application.
    It uses LiteLLMConfiguration to configure the model parameters.

    Attributes:
        _configuration_class: Type of the configuration class used for
            creating LiteLLM model instances.
    """

    _configuration_class: Type = LiteLLMConfiguration

    @classmethod
    def _create_instance(cls, configuration: LiteLLMConfiguration) -> LiteLLM:
        """
        Creates a new instance of the LiteLLM language model.

        Args:
            configuration (LiteLLMConfiguration): Configuration object containing
                settings for the LiteLLM model, including API key, model name,
                maximum tokens, retry settings, context_window, and num output.

        Returns:
            ConfigurableLiteLLM: An instance of the LiteLLM language model configured
                with the provided settings, including optional context_window and
                num_output overrides.
        """
        return ConfigurableLiteLLM(
            api_key=configuration.secrets.api_key.get_secret_value(),
            api_base=configuration.api_base,
            model=configuration.name,
            max_tokens=configuration.max_tokens,
            max_retries=configuration.max_retries,
            context_window=configuration.context_window,
            num_output=configuration.num_output,
        )