Snowflake LLMs via LangChain

[andrebittar-ko]

I am (still) trying to set up a config that uses Snowflake's Cortex LLM services. I posted this on the spacy forum, but figure it might be worth posting a link here.

Further to the code provided in the above post, I have also tried adapting the ChatSnoflakeCortex class in the langchain API (to allow connection via SSO which is used in my organisation), but I get the same error.

From what I have read all LangChain functionalities should be available. Any help would be welcome!

The code to reproduce the behaviour is below:

import os

from langchain_community.chat_models import ChatSnowflakeCortex
from langchain_core.messages import HumanMessage, SystemMessage

import json
from typing import Any, Dict, List, Optional

from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    ChatMessage,
    HumanMessage,
    SystemMessage,
)
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.utils import (
    convert_to_secret_str,
    get_from_dict_or_env,
    get_pydantic_field_names,
    pre_init,
)
from langchain_core.utils.utils import _build_model_kwargs
from pydantic import Field, SecretStr, model_validator

import warnings
warnings.filterwarnings("ignore")

SUPPORTED_ROLES: List[str] = [
    "system",
    "user",
    "assistant",
]


class ChatSnowflakeCortexError(Exception):
    """Error with Snowpark client."""


def _convert_message_to_dict(message: BaseMessage) -> dict:
    """Convert a LangChain message to a dictionary.

    Args:
        message: The LangChain message.

    Returns:
        The dictionary.
    """
    message_dict: Dict[str, Any] = {
        "content": message.content,
    }

    # populate role and additional message data
    if isinstance(message, ChatMessage) and message.role in SUPPORTED_ROLES:
        message_dict["role"] = message.role
    elif isinstance(message, SystemMessage):
        message_dict["role"] = "system"
    elif isinstance(message, HumanMessage):
        message_dict["role"] = "user"
    elif isinstance(message, AIMessage):
        message_dict["role"] = "assistant"
    else:
        raise TypeError(f"Got unknown type {message}")
    return message_dict


def _truncate_at_stop_tokens(
    text: str,
    stop: Optional[List[str]],
) -> str:
    """Truncates text at the earliest stop token found."""
    if stop is None:
        return text

    for stop_token in stop:
        stop_token_idx = text.find(stop_token)
        if stop_token_idx != -1:
            text = text[:stop_token_idx]
    return text

class MyChatSnowflakeCortex(BaseChatModel):
    """Snowflake Cortex based Chat model

    To use you must have the ``snowflake-snowpark-python`` Python package installed and
    either:

        1. environment variables set with your snowflake credentials or
        2. directly passed in as kwargs to the ChatSnowflakeCortex constructor.

    Example:
        .. code-block:: python

            from langchain_community.chat_models import ChatSnowflakeCortex
            chat = ChatSnowflakeCortex()
    """

    _sp_session: Any = None
    """Snowpark session object."""

    model: str = "llama3.1-70b"
    """Snowflake cortex hosted LLM model name, defaulted to `snowflake-arctic`.
        Refer to docs for more options."""

    cortex_function: str = "complete"
    """Cortex function to use, defaulted to `complete`.
        Refer to docs for more options."""

    temperature: float = 0.7
    """Model temperature. Value should be >= 0 and <= 1.0"""

    max_tokens: Optional[int] = None
    """The maximum number of output tokens in the response."""

    top_p: Optional[float] = None
    """top_p adjusts the number of choices for each predicted tokens based on
        cumulative probabilities. Value should be ranging between 0.0 and 1.0. 
    """

    @model_validator(mode="before")
    @classmethod
    def build_extra(cls, values: Dict[str, Any]) -> Any:
        """Build extra kwargs from additional params that were passed in."""
        all_required_field_names = get_pydantic_field_names(cls)
        values = _build_model_kwargs(values, all_required_field_names)
        return values

    @pre_init
    def validate_environment(cls, values: Dict) -> Dict:
        try:
            setattr(cls, '_sp_session', values['model_kwargs']['session'])
        except Exception as e:
            raise ChatSnowflakeCortexError(f"Failed to create session: {e}")

        return values


    def __del__(self) -> None:
        if getattr(self, "_sp_session", None) is not None:
            self._sp_session.close()

    @property
    def _llm_type(self) -> str:
        """Get the type of language model used by this chat model."""
        return f"snowflake-cortex-{self.model}"

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        message_dicts = [_convert_message_to_dict(m) for m in messages]
        message_str = str(message_dicts)
        options = {"temperature": self.temperature}
        if self.top_p is not None:
            options["top_p"] = self.top_p
        if self.max_tokens is not None:
            options["max_tokens"] = self.max_tokens
        options_str = str(options)
        sql_stmt = f"""
            select snowflake.cortex.{self.cortex_function}(
                '{self.model}'
                ,{message_str},{options_str}) as llm_response;"""

        try:
            l_rows = self._sp_session.sql(sql_stmt).collect()
        except Exception as e:
            raise ChatSnowflakeCortexError(
                f"Error while making request to Snowflake Cortex via Snowpark: {e}"
            )

        response = json.loads(l_rows[0]["LLM_RESPONSE"])
        ai_message_content = response["choices"][0]["messages"]

        content = _truncate_at_stop_tokens(ai_message_content, stop)
        message = AIMessage(
            content=content,
            response_metadata=response["usage"],
        )
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

CONNCETION_PARAMETERS = {
... # all params here
}

# create a session using SSO
session = Session.builder.configs(CONNECTION_PARAMETERS).create()

# instantiate the LLM class
llm = MyChatSnowflakeCortex(session=session)

# call the LLM with a message
llm.invoke('Hello! What is your name?').content

# output: 'Nice to meet you! I don\'t have a personal name, but I\'m an AI designed to assist and communicate with users in a helpful and friendly way. I\'m often referred to as a "language model" or a "chatbot." You can think of me as a conversational AI companion. How can I assist you today?'

# now register the class with spacy
from spacy_llm.registry import registry

@registry.llm_models("langchain.MyChatSnowflakeCortex.v1")
def create_snowflake_llm():
    return MyChatSnowflakeCortex(session=session)

# try to load the LLM with my spacy config
from spacy_llm.util import assemble

config_path = 'config.cfg'
nlp = assemble(config_path)

The spacy config:

[paths]
examples = "examples.yml"

[nlp]
lang = "en"
pipeline = ["llm"]
batch_size = 128

[components]

[components.llm]
factory = "llm"

[components.llm.model]
@llm_models = "langchain.MyChatSnowflakeCortex.v1"

[components.llm.task]
@llm_tasks = "spacy.SpanCat.v3"
labels = LABEL1,LABEL2,LABEL3

[components.llm.task.examples]
@misc = "spacy.FewShotReader.v1"
path = ${paths.examples}

[components.llm.task.normalizer]
@misc = "spacy.LowercaseNormalizer.v1"

This causes the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[8], line 4
      1 from spacy_llm.util import assemble
      3 config_path = 'config.cfg'
----> 4 nlp = assemble(config_path)

File ~\.conda\envs\snowflake\lib\site-packages\spacy_llm\util.py:48, in assemble(config_path, overrides)
     46 config_path = Path(config_path)
     47 config = load_config(config_path, overrides=overrides, interpolate=False)
---> 48 return assemble_from_config(config)

File ~\.conda\envs\snowflake\lib\site-packages\spacy_llm\util.py:28, in assemble_from_config(config)
     22 def assemble_from_config(config: Config) -> Language:
     23     """Assemble a spaCy pipeline from a confection Config object.
     24 
     25     config (Config): Config to load spaCy pipeline from.
     26     RETURNS (Language): An initialized spaCy pipeline.
     27     """
---> 28     nlp = load_model_from_config(config, auto_fill=True)
     29     config = config.interpolate()
     30     sourced = get_sourced_components(config)

File ~\.conda\envs\snowflake\lib\site-packages\spacy\util.py:587, in load_model_from_config(config, meta, vocab, disable, enable, exclude, auto_fill, validate)
    584 # This will automatically handle all codes registered via the languages
    585 # registry, including custom subclasses provided via entry points
    586 lang_cls = get_lang_class(nlp_config["lang"])
--> 587 nlp = lang_cls.from_config(
    588     config,
    589     vocab=vocab,
    590     disable=disable,
    591     enable=enable,
    592     exclude=exclude,
    593     auto_fill=auto_fill,
    594     validate=validate,
    595     meta=meta,
    596 )
    597 return nlp

File ~\.conda\envs\snowflake\lib\site-packages\spacy\language.py:1847, in Language.from_config(cls, config, vocab, disable, enable, exclude, meta, auto_fill, validate)
   1844     factory = pipe_cfg.pop("factory")
   1845     # The pipe name (key in the config) here is the unique name
   1846     # of the component, not necessarily the factory
-> 1847     nlp.add_pipe(
   1848         factory,
   1849         name=pipe_name,
   1850         config=pipe_cfg,
   1851         validate=validate,
   1852         raw_config=raw_config,
   1853     )
   1854 else:
   1855     assert "source" in pipe_cfg

File ~\.conda\envs\snowflake\lib\site-packages\spacy\language.py:814, in Language.add_pipe(self, factory_name, name, before, after, first, last, source, config, raw_config, validate)
    810     pipe_component, factory_name = self.create_pipe_from_source(
    811         factory_name, source, name=name
    812     )
    813 else:
--> 814     pipe_component = self.create_pipe(
    815         factory_name,
    816         name=name,
    817         config=config,
    818         raw_config=raw_config,
    819         validate=validate,
    820     )
    821 pipe_index = self._get_pipe_index(before, after, first, last)
    822 self._pipe_meta[name] = self.get_factory_meta(factory_name)

File ~\.conda\envs\snowflake\lib\site-packages\spacy\language.py:702, in Language.create_pipe(self, factory_name, name, config, raw_config, validate)
    699 cfg = {factory_name: config}
    700 # We're calling the internal _fill here to avoid constructing the
    701 # registered functions twice
--> 702 resolved = registry.resolve(cfg, validate=validate)
    703 filled = registry.fill({"cfg": cfg[factory_name]}, validate=validate)["cfg"]
    704 filled = Config(filled)

File ~\.conda\envs\snowflake\lib\site-packages\confection\__init__.py:759, in registry.resolve(cls, config, schema, overrides, validate)
    750 @classmethod
    751 def resolve(
    752     cls,
   (...)
    757     validate: bool = True,
    758 ) -> Dict[str, Any]:
--> 759     resolved, _ = cls._make(
    760         config, schema=schema, overrides=overrides, validate=validate, resolve=True
    761     )
    762     return resolved

File ~\.conda\envs\snowflake\lib\site-packages\confection\__init__.py:808, in registry._make(cls, config, schema, overrides, resolve, validate)
    806 if not is_interpolated:
    807     config = Config(orig_config).interpolate()
--> 808 filled, _, resolved = cls._fill(
    809     config, schema, validate=validate, overrides=overrides, resolve=resolve
    810 )
    811 filled = Config(filled, section_order=section_order)
    812 # Check that overrides didn't include invalid properties not in config

File ~\.conda\envs\snowflake\lib\site-packages\confection\__init__.py:880, in registry._fill(cls, config, schema, validate, resolve, parent, overrides)
    877     getter = cls.get(reg_name, func_name)
    878     # We don't want to try/except this and raise our own error
    879     # here, because we want the traceback if the function fails.
--> 880     getter_result = getter(*args, **kwargs)
    881 else:
    882     # We're not resolving and calling the function, so replace
    883     # the getter_result with a Promise class
    884     getter_result = Promise(
    885         registry=reg_name, name=func_name, args=args, kwargs=kwargs
    886     )

File ~\.conda\envs\snowflake\lib\site-packages\spacy_llm\pipeline\llm.py:86, in make_llm(nlp, name, task, model, cache, save_io, validate_types)
     81     raise ValueError(
     82         "Argument `task` has not been specified, but is required (e. g. {'@llm_tasks': "
     83         "'spacy.NER.v3'})."
     84     )
     85 if validate_types:
---> 86     validate_type_consistency(task, model)
     88 return LLMWrapper(
     89     name=name,
     90     task=task,
   (...)
     94     vocab=nlp.vocab,
     95 )

File ~\.conda\envs\snowflake\lib\site-packages\spacy_llm\ty.py:380, in validate_type_consistency(task, model)
    378 # Validate the 'model' object
    379 if not (len(type_hints["model"]) == 2 and "return" in type_hints["model"]):
--> 380     raise ValueError(
    381         "The 'model' Callable should have one input argument and one return value."
    382     )
    383 for k in type_hints["model"]:
    384     if k == "return":

ValueError: The 'model' Callable should have one input argument and one return value.