How to eliminate global variable usage in LangChain Runnable chain

[quartermaine]

I'm working with LangChain's Runnable interface to build a processing chain. Currently, I'm using a global variable llm_az to hold the language model instance, which I want to eliminate for better modularity and testability.

Here's a simplified version of my code:

 from langchain_openai import AzureChatOpenAI
 from langchain_core.prompts import ChatPromptTemplate

 # Global variable I want to remove
 llm_az = RunnableLambda(lambda x: None)

 def get_llm_from_state(state):
     global llm_az
     llm_az = AzureChatOpenAI(
         model=state.get('model_name'),
         deployment_name=state.get('deployment_name'),
         azure_endpoint=state.get('azure_openai_endpoint'),
         openai_api_key=state.get('openai_api_key'),
         openai_api_type=state.get('openai_api_type'),
         openai_api_version=state.get('openai_api_version'),
         max_tokens=int(state.get('max_tokens', 1000))
     )

 def get_final_prompt(x):
     prompt = ChatPromptTemplate.from_messages([
         ("system", x.get("answer_prompt", "Default system prompt")),
         ("human", "{question}")
     ])
     return prompt.format_messages(question=x['question'])

 def build_chain():
     return (
         RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: get_llm_from_state(x)) 
         )
         | RunnableLambda(lambda x: get_final_prompt(x))
         | RunnableLambda(lambda x: llm_az.invoke(x))
     )

How can I refactor this code to avoid using the global llm_az variable, perhaps by passing the language model instance through the chain in a more functional and modular way?

[dosubot[bot]]

To refactor your code and eliminate the use of the global variable llm_az, you can pass the language model instance through the chain as a parameter. This approach enhances modularity and testability. Here's how you can refactor your code:

1. Modify the function signatures to accept the language model as a parameter.
2. Update the chain to pass the language model instance appropriately.

Here's a refactored version of your code:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

def get_llm_from_state(state):
    return AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: {'llm': get_llm_from_state(x)})
        )
        | RunnableLambda(lambda x: {'prompt': get_final_prompt(x), 'llm': x['llm']})
        | RunnableLambda(lambda x: x['llm'].invoke(x['prompt']))
    )

Key Changes:

• get_llm_from_state: Returns the language model instance instead of setting a global variable.
• Chain Construction: Each RunnableLambda passes the language model instance along with other data through the chain using dictionaries. This ensures that the language model is available where needed without relying on global state.

This refactoring makes your code more functional and modular, allowing for easier testing and maintenance.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

When I use the suggested fix I am getting the following error :

raise ValueError(msg) ValueError: Invalid input type <class 'dict'>. Must be a PromptValue, str, or list of BaseMessages.

[dosubot[bot]]

To resolve the "Invalid input type <class 'dict'>. Must be a PromptValue, str, or list of BaseMessages" error, ensure that the input you provide to the invoke method is of the correct type. The invoke method expects a PromptValue, str, or a list of BaseMessages.

In your refactored code, you're passing a dictionary to the invoke method. You need to extract the appropriate value from the dictionary and convert it to a valid type. Here's how you can adjust the final RunnableLambda to pass the correct input:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

def get_llm_from_state(state):
    return AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: {'llm': get_llm_from_state(x)})
        )
        | RunnableLambda(lambda x: {'prompt': get_final_prompt(x), 'llm': x['llm']})
        | RunnableLambda(lambda x: x['llm'].invoke(x['prompt'][0]))  # Adjusted to pass the first message
    )

In this adjustment, x['prompt'][0] is used to extract the first message from the list of BaseMessages returned by get_final_prompt. Ensure that this message is of a type that invoke can accept. Adjust the extraction logic if your prompt structure differs.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

Still getting the same error, after the initialization of the llm in the get_llm_from_state function

[dosubot[bot]]

Ensure that the input to the invoke method is a valid type, such as a str or a list of BaseMessages. If get_final_prompt returns a list of messages, you should pass this list directly to invoke. Here's how you can adjust the code:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

def get_llm_from_state(state):
    return AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: {'llm': get_llm_from_state(x)})
        )
        | RunnableLambda(lambda x: {'prompt': get_final_prompt(x), 'llm': x['llm']})
        | RunnableLambda(lambda x: x['llm'].invoke(x['prompt']))  # Pass the entire list of messages
    )

Make sure that get_final_prompt returns a list of BaseMessages that invoke can process. If invoke requires a different format, adjust the input accordingly.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

The error happens after the initialization of the llm and in the RunnablePassthrough.assign

RunnablePassthrough.assign( llm=RunnableLambda(lambda x: get_llm_from_state(x)) )

[dosubot[bot]]

The error you're encountering might be due to the way the RunnablePassthrough.assign method is handling the input. The assign method expects the input to be a dictionary with specific properties, and it seems like the input type might not be matching the expected schema.

To resolve this, ensure that the input to RunnablePassthrough.assign is structured correctly. Here's how you can adjust the code:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

def get_llm_from_state(state):
    return AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: {'llm': get_llm_from_state(x)})
        )
        | RunnableLambda(lambda x: {'prompt': get_final_prompt(x), 'llm': x['llm']})
        | RunnableLambda(lambda x: x['llm'].invoke(x['prompt']))
    )

Ensure that the input to RunnablePassthrough.assign is a dictionary with the required properties, such as question, as specified in the input schema ^[1][2]. Adjust the input structure to match the expected schema, and verify that the RunnableLambda is returning the correct type.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

I have managed to solve the llm initialization problem by converting the output to dictionary with:

    return {'llm': AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )}

but now if I do the following :

    return (
        RunnablePassthrough.assign(
            llm=RunnableLambda(lambda x: {'llm': get_llm_from_state(x)})
        )
        | RunnableLambda( lambda x: x['llm']['llm'].invoke(lambda x: get_final_prompt(x)))
    )

I am getting the error:

    | RunnableLambda( lambda x: x['llm']['llm'].invoke(lambda x: get_final_prompt(x)))
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    raise ValueError(msg)
ValueError: Invalid input type <class 'function'>. Must be a PromptValue, str, or list of BaseMessages.

[dosubot[bot]]

The error you're encountering is due to passing a function to the invoke method, which expects a PromptValue, str, or a list of BaseMessages. To fix this, ensure that you pass the correct type to invoke. Here's how you can adjust your code:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

def get_llm_from_state(state):
    return {'llm': AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )}

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: get_llm_from_state(x))
        )
        | RunnableLambda(lambda x: {'prompt': get_final_prompt(x), 'llm': x['llm']})
        | RunnableLambda(lambda x: x['llm'].invoke(x['prompt']))  # Pass the formatted prompt directly
    )

In this adjustment, get_final_prompt(x) is called to generate the prompt, and the result is passed directly to invoke. Ensure that get_final_prompt returns a type that invoke can accept, such as a str or a list of BaseMessages ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

If I use :

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: get_llm_from_state(x))
        )
        |  RunnablePassthrough.assign( (
           final_prompt = RunnableLambda(lambda x: : get_final_prompt(x))
       )
        | RunnableLambda(lambda x: x['llm']['llm'].invoke(x['final_prompt']))  # Pass the formatted prompt directly
    )

I am getting a completely wrong answer and the sql code used is correct

[dosubot[bot]]

To address the issue of getting incorrect answers despite using correct SQL code, ensure that the SQL query is being processed and validated correctly. You can enhance the chain by adding a validation step to check for common SQL mistakes and ensure the query is accurate before execution. Here's how you can implement this:

1. Add a validation step to check and correct common SQL mistakes.
2. Ensure that the SQL query is being executed against the correct tables.

Here's an example of how you can implement this:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

def get_llm_from_state(state):
    return {'llm': AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )}

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def validate_sql_query(x):
    # Define the system prompt for validation
    system = """Double check the user's SQL query for common mistakes, including:
    - Using NOT IN with NULL values
    - Using UNION when UNION ALL should have been used
    - Using BETWEEN for exclusive ranges
    - Data type mismatch in predicates
    - Properly quoting identifiers
    - Using the correct number of arguments for functions
    - Casting to the correct data type
    - Using the proper columns for joins

    If there are any of the above mistakes, rewrite the query.
    If there are no mistakes, just reproduce the original query with no further commentary.

    Output the final SQL query only."""
    prompt = ChatPromptTemplate.from_messages(
        [("system", system), ("human", "{query}")]
    ).partial(dialect=x['llm'].dialect)

    # Create the validation chain
    validation_chain = prompt | x['llm'] | StrOutputParser()

    # Validate the query
    return validation_chain.invoke({"query": x['final_prompt']})

def build_chain():
    return (
        RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: get_llm_from_state(x))
        )
        | RunnablePassthrough.assign(
            final_prompt=RunnableLambda(lambda x: get_final_prompt(x))
        )
        | RunnableLambda(lambda x: {'validated_query': validate_sql_query(x), 'llm': x['llm']})
        | RunnableLambda(lambda x: x['llm'].invoke(x['validated_query']))  # Pass the validated query
    )

This approach adds a validation step to ensure the SQL query is correct before execution. Adjust the validate_sql_query function to match your specific validation needs ^[1][2].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

This solution is still not working and invoking the final_prompt doesn't contain all the values in the runnable x in order for the llm to get the result

[dosubot[bot]]

To ensure that the final_prompt contains all necessary values from the runnable x, you need to manage the context and data flow correctly. This involves using context setters and getters, as well as appropriate prompt templates. Here's how you can adjust your code:

1. Use Context.setter and Context.getter to manage the flow of data.
2. Ensure that all necessary values are included in the final_prompt.

Here's an example of how you can implement this:

from langchain_core.runnables import RunnableLambda, RunnablePassthrough, Context
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

def get_llm_from_state(state):
    return {'llm': AzureChatOpenAI(
        model=state.get('model_name'),
        deployment_name=state.get('deployment_name'),
        azure_endpoint=state.get('azure_openai_endpoint'),
        openai_api_key=state.get('openai_api_key'),
        openai_api_type=state.get('openai_api_type'),
        openai_api_version=state.get('openai_api_version'),
        max_tokens=int(state.get('max_tokens', 1000))
    )}

def get_final_prompt(x):
    prompt = ChatPromptTemplate.from_messages([
        ("system", x.get("answer_prompt", "Default system prompt")),
        ("human", "{question}")
    ])
    return prompt.format_messages(question=x['question'])

def build_chain():
    return (
        Context.setter("input")
        | RunnablePassthrough.assign(
            init=RunnableLambda(lambda x: get_llm_from_state(x))
        )
        | {
            "final_prompt": RunnableLambda(lambda x: get_final_prompt(x)) | Context.setter("final_prompt"),
            "llm": Context.getter("input")["llm"]
        }
        | RunnableLambda(lambda x: x['llm'].invoke(Context.getter("final_prompt")))
    )

In this setup:

• Context.setter("input") is used to set the initial context.
• RunnableLambda retrieves and sets the final_prompt in the context.
• The invoke method uses Context.getter("final_prompt") to ensure it has all necessary values.

This approach ensures that the final_prompt is constructed with all required data, allowing the language model to produce the correct result ^[1][2].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

@dosu

there's no Context in the langchain_core.runnables

[dosubot[bot]]

To ensure that all necessary values are included in the final prompt when using LangChain's Runnable interface, you can use the ensure_config function. This function ensures that a configuration is a dictionary with all keys present, including default values for any missing keys. Here's an example of how you can use it:

from langchain_core.runnables.config import ensure_config

# Example usage
config = {
    "tags": ["example"],
    "metadata": {"key": "value"},
    # other config values...
}

# Ensure all necessary values are included
final_config = ensure_config(config)

The ensure_config function will fill in any missing keys with default values, ensuring that your configuration is complete and ready for use with the Runnable interface ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[quartermaine]

Can you give me an example for my previous chain ?