Language model kernels

def create_embedding_model_from_settings(
    settings: EmbeddingSettings,
) -> Embeddings:
    """
    Create langchain embedding model from an EmbeddingSettings
    object. Raises a ValueError if the source argument
    is not supported.

    Args:
        settings: an EmbeddingSettings object containing model
            configuration.

    Returns:
        a Langchain embeddings model object.

    Raises ValueError, TypeError, ValidationError

    Example:
        ```python
        settings = EmbeddingSettings(
            dense_model="OpenAI/text-embedding-3-small"
        )
        model = create_embedding_model_from_settings(settings)
        ```
    """
    return langchain_embeddings[settings]

def create_embedding_model_from_spec(
    dense_model: str, *, sparse_model: SparseModel = "Qdrant/bm25"
) -> Embeddings:
    """
    Create langchain embedding model from source_name and
    model_name. Raises a ValueError if the source_name argument
    is not supported.

    Args:
        dense_model: the model specification in the form source/model,
            such as 'OpenAI/text-embedding-3-small'

    Returns:
        a Langchain embeddings model object.

    Raises ValueError, TypeError, ValidationError

    Example:
        ```python
        spec = {'dense_model': "OpenAI/text-embedding-3-small"}
        model = create_embedding_model_from_spec(**spec)
        ```
    """
    spec = EmbeddingSettings(
        dense_model=dense_model,
        sparse_model=sparse_model,
    )
    return langchain_embeddings[spec]

def create_model_from_settings(
    settings: LanguageModelSettings,
) -> BaseChatModel:
    """
    Create langchain model from a LanguageModelSettings object.
    Raises a ValueError if the source argument is not supported.

    Args:
        settings: a LanguageModelSettings object containing model
            configuration.

    Returns:
        a Langchain model object.

    Raises ValueError, TypeError, ValidationError

    Example:
        ```python
        # Create settings explicitly.
        config = LanguageModelSettings(
            model="OpenAI/gpt-4o-mini",
            temperature=0.7,
            max_tokens=1000,
            max_retries=3,
        )
        model = create_model_from_settings(config)
        response = model.invoke("Why is the sky blue?")

        # Load settings from config.toml.
        settings = Settings()
        system_prompt = "You are a helpful assistant."
        model = create_model_from_settings(
            settings.minor,
        )
        response = model.invoke("Why is the grass green?")
        ```
    """
    return langchain_models[settings]

def create_model_from_spec(
    model: str,
    *,
    temperature: float = 0.1,
    max_tokens: int | None = None,
    max_retries: int = 2,
    timeout: float | None = None,
    provider_params: dict[str, MetadataPrimitiveWithList] | None = None,
) -> BaseChatModel:
    """
    Create langchain model from specifications.

    Args:
        model: the model in the form source/model, such as
            'OpenAI/gpt-4o'

    Returns:
        a Langchain model object.

    Raises ValueError, TypeError, ValidationError

    Example:
        ```python
        spec = {'model': "OpenAI/gpt-4o-mini"}
        model = create_model_from_spec(**spec)
        ```
    """
    if provider_params is None:
        provider_params = {}

    spec = LanguageModelSettings(
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        max_retries=max_retries,
        timeout=timeout,
        provider_params=provider_params,
    )
    return langchain_models[spec]

def create_prompt(
    prompt: str,
    name: str,
    *,
    system_prompt: str | None = None,
    replace: bool = False,
) -> None:
    """
    Adds a custom prompt template to the prompt dictionary.

    Args:
        prompt: the prompt text.
        name: the name of the tool. This will also define a kernel
            with the same name (i.e. a Langchain runnable)
        system_prompt: an optional system prompt text.
    """

    if name in prompt_library.keys():
        if replace:
            prompt_library.pop(name)  # type: ignore
        else:
            raise ValueError(
                f"'{name}' is already a registered prompt. "
                "Use another name to register a custom prompt."
            )

    # We abuse the lack of run-time checks for Literals here. We do
    # this because we want the availability of the preformed prompts
    # given by Literal but also the flexibility to add new prompts.
    definition = PromptDefinition(
        name=name,
        prompt=prompt,
        system_prompt=system_prompt,
    )
    prompt_library[name] = definition  # type: ignore

def create_embeddings(
    settings: (
        dict[str, str] | EmbeddingSettings | Settings | None
    ) = None,
) -> Embeddings:
    """
    Creates a Langchain embeddings runnable from a configuration
    object.

    Args:
        settings: an EmbeddingSettings object with the following
            fields:

            - dense_model: a specification in the form provider/
            model, for example 'OpenAI/text-embedding-3-small'
            - sparse_model: a sparse model specification.

            Alternatively, a dictionary with the same fields and
            text. If None (default), the settings will be read
            from the configuration file. If no configuration file
            exists, a settings object will be created with default
            parameters.

    Returns:
        a Langchain object that embeds text by calling embed_documents
            or embed_query.

    Raises:
        ValidationError, TypeError: for invalid spec
        ImportError: for missing libraries
        requests.exceptions.ConnectionError: if not online

    Example:
    ```python
    from lmm.models.langchain.runnables import (
        create_embeddings,
    )

    try:
        encoder: Embeddings = create_embeddings()
        vector = encoder.embed_query("Why is the sky blue?")
        documents = ["The sky is blue due to its oxygen content"]
        vectors = encoder.embed_documents(documents)
    except Exception ...
    ```
    """
    if not bool(settings):  # includes empty dict
        sets = Settings()
        settings = sets.embeddings
    elif isinstance(settings, Settings):
        settings = settings.embeddings
    elif isinstance(settings, dict):
        # checked by pydantic model
        settings = EmbeddingSettings(**settings)  # type: ignore

    return embeddings_library[settings]

def create_kernel_from_objects(
    human_prompt: str,
    *,
    system_prompt: str | None = None,
    language_model: (
        BaseChatModel | LanguageModelSettings | Settings | None
    ) = None,
) -> RunnableType:
    """
    Creates a Langchain runnable from a prompt template and
    a language settings object. This name is not registered in the
    prompts library; it is available directly.

    Args:
        human_prompt: prompt text
        system_prompt: system prompt text
        language_model: either a Langchain BaseChatModel, or
            a LanguageModelSettings object, or None (default). In
            this latter case the language.minor from the config
            file is used to create the model.

    Returns:
        a Langchain runnable, a type aliased as `RunnableType`.

    Example:
    ```python
    human_prompt = '''
    Provide the questions to which the text answers.

    TEXT:
    {text}
    '''
    settings = Settings()
    try:
        model = create_kernel_from_objects(
            human_prompt=human_prompt,
            system_prompt="You are a helpful assistant",
            language_model=settings.aux,
        )
    except Exception ...

    # model use:
    try:
        response = model.invoke({'text': "Logistic regression is used"
            + " when the outcome variable is binary"})
    except Exception ...
    ```
    """
    if language_model is None:
        settings = Settings()
        language_model = create_model_from_settings(settings.minor)
        name = f"Custom:{settings.minor}"
    elif isinstance(language_model, Settings):
        name = f"Custom:{language_model.minor}"
        language_model = create_model_from_settings(
            language_model.minor
        )
    elif isinstance(language_model, LanguageModelSettings):
        name = f"Custom:{language_model}"
        language_model = create_model_from_settings(language_model)
    else:  # it's a BaseChatModel
        name = "Custom"

    # Langchain prompt
    prompt: ChatPromptTemplate
    if system_prompt is not None:
        prompt = ChatPromptTemplate.from_messages(  # type: ignore
            [
                SystemMessagePromptTemplate.from_template(
                    system_prompt
                ),
                HumanMessagePromptTemplate.from_template(
                    human_prompt
                ),
            ]
        )
    else:
        prompt = ChatPromptTemplate.from_template(human_prompt)

    # combine into a runnable
    runnable: RunnableType = prompt | language_model | StrOutputParser()  # type: ignore
    # .name is a member function of RunnableSerializable
    # inited to None, which we re-initialize here
    runnable.name = name
    return runnable

def create_runnable(
    runnable_name: PromptNames | str,
    user_settings: (
        dict[str, str] | LanguageModelSettings | Settings | None
    ) = None,
    system_prompt: str | None = None,
    **kwargs: MetadataPrimitiveWithList,
) -> RunnableType:  # RunnableSerializable[dict[str, str], str]
    """
    Creates a Langchain chain (a 'runnable') by combining tools/prompts
    created under the runnable_name parameters and configurations from
    config.toml with optional override settings.

    The function maps different chains to their appropriate
    language model settings categories. For example,
    - 'query', 'query_with_context' -> major model settings
    - 'question_generator', 'summarizer' -> minor model settings
    - 'allowed_content_validator', 'context_validator' ->
                                            aux model settings

    Settings Hierarchy (highest to lowest priority):
    1. user_settings parameter (if provided)
    2. config.toml file settings
    3. Default settings from Settings class

    Args:
        runnable_name: The name of the runnable to create. If one of
            the supported names defined in the PromptNames
            literal type, returns a cached runnable object. Otherwise,
            looks up in the prompt_library dictionary if there is
            a prompt with that runnable_name, and returns a runnable
            object for a chat with that prompt.
        user_settings: Optional settings to override the default
            configuration. Can be either:
            - dict[str, str]: Dictionary with 'model' key
            - LanguageModelSettings: Pydantic model instance
            - None: Use settings from config.toml or defaults
        system_prompt: System prompt used in messages with the
            language model.

    Returns:
        A KernelType object RunnableSerializable[dict[str, str], str],
            A Langchain runnable chain that combines a prompt template,
            language model, and string output parser. The chain accepts a
            dictionary of template variables and returns a string
            response.

    Raises:
        ValueError: If runnable_name is not supported or if user_settings
            contains invalid model source names. No check is made at this
            stage that the model names are correct (as they frequently
            change); instead, failure occurs when the .invoke member function
            is called.
        ValidationError, TypeError: alternative errors raised in the same
            circumstances as above.
        ImportError: for not installed libraries.

    Examples:
        Create runnable with default settings read from configuration
        file:
        ```python
        try:
            runnable = create_runnable("query")
        except Exception ...
        ```

        Override with dictionary:
        ```python
        try:
            runnable = create_runnable("summarizer",
                {"model": "OpenAI/gpt-4o"})
        except Exception ...
        ```

        Override with settings object:
        ```python
        from lmm.config.config import LanguageModelSettings
        settings = LanguageModelSettings(
            model="Mistral/mistral-small-latest"
        )
        try:
            runnable = create_runnable("question_generator", settings)
        except Exception ...
        ```

        The runnable object may be used with Langchain `.invoke` syntax:

        ```python
        try:
            response = runnable.invoke(
                {'text': "Logistic regression is used when the outcome"
                    + " variable is binary."}
            )
        except Exception ...
        ```
    """

    def _create_or_get(
        settings: LanguageModelSettings,
        runnable_name: PromptNames,
        system_prompt: str | None,
        **kwargs: MetadataPrimitiveWithList,
    ) -> RunnableType:
        model: RunnableDefinition = RunnableDefinition(
            runnable_name=runnable_name,
            settings=settings,
            system_prompt_override=system_prompt,
            params=_dict_to_runnable_par(kwargs),
        )
        return runnable_library[model]

    settings: Settings
    match user_settings:
        case dict() if bool(user_settings):
            try:
                settings = Settings(
                    major=user_settings,  # type: ignore
                    minor=user_settings,  # type: ignore
                    aux=user_settings,  # type: ignore
                )
            except Exception as e:
                raise ValueError(f"Invalid model definition:\n{e}")
        case LanguageModelSettings():
            settings = Settings(
                major=user_settings,
                minor=user_settings,
                aux=user_settings,
            )
        case Settings():
            settings = user_settings
        case None | {}:
            settings = Settings()
        case _:
            raise ValueError(
                f"Invalid model definition: {user_settings}"
            )

    # Logic to retrieve or create prompt definition to check model tier
    prompt_definition: PromptDefinition
    params_dict = dict(kwargs)

    # We use create_prompt_definition if we have params (to handle validation)
    # or prompt_library if we don't (to handle custom prompts).
    if params_dict:
        try:
            prompt_definition = create_prompt_definition(runnable_name, **params_dict) # type: ignore
        except ValueError:
            # Fallback: if create_prompt_definition fails (e.g. custom prompt), 
            # check library. But strictly custom prompts shouldn't have params 
            # in current architecture unless handled by create_prompt_definition.
            # raising the error is appropriate if params were intended for a 
            # prompt that doesn't support them.
            raise
    else:
        try:
            prompt_definition = prompt_library[runnable_name] # type: ignore
        except Exception as e:
            raise ValueError(f"{runnable_name} is not a valid runnable name ({e})")

    match prompt_definition.model_tier:
        case 'major':
            settings_to_use = settings.major
        case 'minor':
            settings_to_use = settings.minor
        case 'aux':
            settings_to_use = settings.aux
        case _:
            settings_to_use = settings.minor

    return _create_or_get(
        settings_to_use, 
        runnable_name, # type: ignore
        system_prompt,
        **kwargs
    )

def yield_constant_message(
    message: str = "Message",
) -> ConstantMessageIterator:
    """
    Create and return a ConstantMessageIterator instance.

    This function creates an iterator that generates the same
    message repeatedly. The iterator is infinite and will continue
    generating the message indefinitely.

    Args:
        message: The message to generate. Defaults to "Message".

    Returns:
        A ConstantMessageIterator instance.

    Example:
        >>> iterator = constant_message()
        >>> next(iterator)
        'Message'
        >>> next(iterator)
        'Message'

        >>> custom_iterator = constant_message("Alert")
        >>> next(custom_iterator)
        'Alert'
        >>> next(custom_iterator)
        'Alert'
    """
    return ConstantMessageIterator(message)

def yield_message(prefix: str = "Message") -> MessageIterator:
    """
    Create and return a MessageIterator instance.

    This function creates an iterator that generates sequential messages
    with the specified prefix. The iterator is infinite and will continue
    generating messages indefinitely.

    Args:
        prefix: The prefix to use for generated messages. Defaults to "Message"

    Returns:
        A MessageIterator instance that generates messages like:
        "Message 1", "Message 2", "Message 3", etc.

    Example:
        >>> iterator = yield_message()
        >>> next(iterator)
        'Message 1'
        >>> next(iterator)
        'Message 2'

        >>> custom_iterator = yield_message("Alert")
        >>> next(custom_iterator)
        'Alert 1'
        >>> next(custom_iterator)
        'Alert 2'
    """
    return MessageIterator(prefix)