Using Snowflake LLM via LangChain API in spacy-llm

[andrebittar-ko]

[APOLOGIES - realise I should maybe have posted this in spacy-llm discussions]

Taking inspiration from this blog post, I would like to use LLMs hosted by Snowflake with spacy (including annotation tasks in Prodigy).

I have slightly adapted the code in the blog to create a wrapper class that inherits from the langchain LLM class and implements the specified mehtods (_call and _identifying_params property):

from typing import Any, Dict, Iterator, List, Mapping, Optional
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM

class SnowflakeCortexLLM(LLM,extra='allow'):
    def __init__(self):
        super().__init__()
        self.connection_params = CONNECTION_PARAMTERS = {
            "user": "MY_EMAIL",
            "account": "MY_SNOWFLAKE_ACCOUNT",
            "authenticator": "externalbrowser",
            "role": "MY_SNOWFLAKE_USER_ROLE",
            "database": "MY_DATABASE",
            "schema": "MY_SCHEMA"
        }
        self.sp_session = Session.builder.configs(self.connection_params).create()
        self.model = 'llama3.1-70b' # the model I want to use for now, but this should be passed as an argument eventually
        self.cortex_function: str = 'complete'

    @property
    def _llm_type(self) -> str:
        return "snowflake_cortex"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        
        prompt_text = prompt
        sql_stmt = f'''
            select snowflake.cortex.{self.cortex_function}(
                '{self.model}'
                ,'{prompt_text}') as llm_reponse;'''
        
        l_rows = self.sp_session.sql(sql_stmt).collect()

        llm_response = l_rows[0]['LLM_REPONSE']

        return llm_response

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            "model": self.model
            ,"cortex_function" : self.cortex_function
            ,"snowpark_session": self.sp_session.session_id
        }

I instantiate the class: llm = SnowflakeCortexLLM(), which connects to the Snowflake LLM service and allows me to make calls: llm.invoke('Hello!')

Which provides a response: Hello! How can I assist you today?'

So, this part is working. However, I am a little stuck as to how to integrate this with the langchain API in spacy-llm and confgi file.

I have registered the above class as follows:

from spacy_llm.registry import registry

@registry.llm_models("langchain.SnowflakeCortexLLM.v1")
def create_snowflake_llm():
    return SnowflakeCortexLLM()

I have defined the following config.cfg, for a span categorisation task:

[paths]
examples = "path_to_example_file.yml"

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

[components]

[components.llm]
factory = "llm"

[components.llm.model]
@llm_models = "langchain.SnowflakeCortexLLM.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"

I load the config as follows:

from spacy_llm.util import assemble

nlp = assemble('config.cfg')

But this results in an error:

...
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.

I have also tried registering the model without the langchain namespace (e.g. SnowflakeCortexLLM.v1 and using the corresponding reference in the config file), but get the same error.

Which I guess means there is an extra piece of code to fit this into spaCy's langchain API...???

Is this possible? Is there a better way to go about this?

[edit]: I have now also tried using the langchain Snowflake API. to do this. I get the same error: ValueError: The 'model' Callable should have one input argument and one return value. I have posted my full code for that on the spacy-llm forum.