How to Use and Update Complex Pydantic Models as Input in LangChain with Non-Linear Chains?

[JamesLarkinWhite]

I'm working with LangChain and have a complex pydantic model that encapsulates all the necessary information for my process. I need to pass only parts of this object to different stages of a chain and update the object with the outputs of these stages so that the data is persisted throughout the process. Additionally, my chain is non-linear, involving routers and forks.

Here's a simplified version of my pydantic model and the chain setup that obviously does not work but should help you understand my reaoning:

from pydantic import BaseModel, Field
from typing import List, Optional
from langchain.chains import LLMChain, RunnableLambda, RunnablePassthrough
from langchain.prompts import PromptTemplate
from langchain.tools import ClassificationResponse

class Message(BaseModel):
    role: str
    content: str

class QuestionContext(BaseModel):
    current_description: str
    current__comments: Optional[str]
    current_question: str
    current_context: str
    category: Optional[str]

class InputData(BaseModel):
    messages: List[Message]
    question_context: QuestionContext

# Define your chains with specific prompts
main_chain_prompt = PromptTemplate(
    template="Category: {category}\nQuestion: {question}\nContext: {context}\n\nPlease provide a response based on the given category and context.",
    input_variables=["category", "question", "context"],
)

tagging_prompt_new = PromptTemplate(
    template="Question: {question}\nContext: {context}\n\nPlease determine the category for this question.",
    input_variables=["question", "context"],
)

main_chain = LLMChain(llm=llm, prompt=main_chain_prompt)
# this is just a mockup. here i would use a pydantic class to classify
from pydantic import BaseModel, Field, conint
from typing import Literal
from enum import Enum
from langchain.prompts import ChatPromptTemplate
from langchain.chains import create_tagging_chain, create_tagging_chain_pydantic


class ClassificationResponse(BaseModel):
    """
    A few-shot example of text classification.

    """

    label: Literal[
        "Its red",
        "its blue",
....
    ] = Field(
        ...,
        description="The predicted class labels for the question",
    )
tagging_prompt= PromptTemplate(
    template="Question: {question}\nContext: {context}\n\nPlease determine the category of this text.",
    input_variables=["question", "context"],
)

tagging_chain = tagging_prompt | llm.with_structured_output(ClassificationResponse)

sequential_chain = RunnablePassthrough.assign(tagging=tagging_chain) | main_chain

def router(input_data: InputData):
    if input_data.question_context.category:
        return main_chain
    else:
        return sequential_chain

full_chain = RunnableLambda(router) | RunnableLambda(
    lambda chain, input_data: chain.run(
        {
# i am unsure what i need to write here
        }
    )
)

# Example usage
input_data = InputData(
    messages=[Message(role="user", content="Hello")],
    question_context=QuestionContext(
        current_ticket_description="Issue with blue ballon",
        current_ticket_comments=None,
        current_question="What could be the problem?",
        current_context="my blue ballon is red",
        automotive_category_of_request="Engine"
    )
)

# Run the full chain with the input data
result = full_chain.run(input_data)

In reality the result would be used to updated Input_data such that i have all information in one place.
I am not sure how i could that use a nother sequantial chain that does RAG for example and persist the InputData object throuut the whole chain.

Is this even a valid approach?

[dosubot[bot]]

Hey there, @JamesLarkinWhite! I'm here to help you with any bugs, questions, or contributions you might have. Let's work together to solve this!

To handle complex Pydantic models in LangChain with non-linear chains, you can follow these steps to selectively pass parts of the model to different stages and update the model with outputs:

1. Extract Relevant Parts of the Pydantic Model: Use a function to extract the necessary fields from your Pydantic model and pass them to the chain stages. This can be done using a RunnableLambda to transform the input data.

2. Update the Pydantic Model: After each stage, update the Pydantic model with the results. You can use a RunnableLambda or a custom function to merge the outputs back into the model.

3. Use Routers for Non-Linear Logic: Implement routing logic using a custom function or RunnableLambda to decide which chain to execute based on the current state of the Pydantic model.

Here's a simplified example of how you might structure this:

from pydantic import BaseModel, Field
from typing import List, Optional
from langchain.chains import LLMChain, RunnableLambda
from langchain.prompts import PromptTemplate

class Message(BaseModel):
    role: str
    content: str

class QuestionContext(BaseModel):
    current_description: str
    current_comments: Optional[str]
    current_question: str
    current_context: str
    category: Optional[str]

class InputData(BaseModel):
    messages: List[Message]
    question_context: QuestionContext

# Define your chains with specific prompts
main_chain_prompt = PromptTemplate(
    template="Category: {category}\nQuestion: {question}\nContext: {context}\n\nPlease provide a response based on the given category and context.",
    input_variables=["category", "question", "context"],
)

main_chain = LLMChain(llm=llm, prompt=main_chain_prompt)

def extract_input_data(input_data: InputData):
    return {
        "category": input_data.question_context.category,
        "question": input_data.question_context.current_question,
        "context": input_data.question_context.current_context,
    }

def update_input_data(input_data: InputData, result):
    # Update the input_data with the result
    input_data.question_context.category = result.get("category", input_data.question_context.category)
    # Add more updates as needed
    return input_data

def router(input_data: InputData):
    if input_data.question_context.category:
        return main_chain
    else:
        return tagging_chain

full_chain = RunnableLambda(router) | RunnableLambda(
    lambda chain, input_data: chain.run(extract_input_data(input_data))
) | RunnableLambda(update_input_data)

# Example usage
input_data = InputData(
    messages=[Message(role="user", content="Hello")],
    question_context=QuestionContext(
        current_description="Issue with blue balloon",
        current_comments=None,
        current_question="What could be the problem?",
        current_context="my blue balloon is red",
        category=None
    )
)

# Run the full chain with the input data
result = full_chain.run(input_data)

In this setup:

• extract_input_data: Extracts the necessary fields from the InputData model to pass to the chain.
• update_input_data: Updates the InputData model with the results from the chain.
• router: Determines which chain to execute based on the current state of the InputData.

This approach allows you to maintain and update the state of your Pydantic model throughout the chain execution.

To continue talking to Dosu, mention @dosu.

[JamesLarkinWhite]

@dosu

I needed to adapt your code a little to get it to run:

from pydantic import BaseModel, Field
from typing import List, Optional, Literal
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain.llms import OpenAI  # Assuming you're using OpenAI


# Define your models
class ClassificationResponse(BaseModel):
    label: Literal["Red Balloon", "Blue Balloon", "Green Balloon", "Other Balloon color"] = Field(
        ..., description="The predicted class labels for the question"
    )


class Message(BaseModel):
    role: str
    content: str


class QuestionContext(BaseModel):
    current_description: str
    current_comments: Optional[str]
    current_question: str
    current_context: str
    category: Optional[str]


class InputData(BaseModel):
    messages: List[Message]
    question_context: QuestionContext


# Define chains
main_chain_prompt = PromptTemplate(
    template="Category: {category}\nQuestion: {question}\nContext: {context}\n\nPlease provide a response based on the given category and context.",
    input_variables=["category", "question", "context"],
)


main_chain = LLMChain(llm=llm, prompt=main_chain_prompt)

tagging_prompt = PromptTemplate(
    template="Question: {question}\nContext: {context}\n\nPlease determine the category of this text.",
    input_variables=["question", "context"],
)
tagging_chain = tagging_prompt | llm.with_structured_output(ClassificationResponse)


# Helper functions
def extract_for_tagging(input_data: InputData):
    return {
        "question": input_data.question_context.current_question,
        "context": input_data.question_context.current_context,
    }


def extract_for_main_chain(input_data: InputData):
    return {
        "category": input_data.question_context.category,
        "question": input_data.question_context.current_question,
        "context": input_data.question_context.current_context,
    }


def update_input_data(input_data: InputData, result):
    if isinstance(result, ClassificationResponse):
        input_data.question_context.category = result.label
    elif isinstance(result, dict):
        input_data.question_context.category = result.get("output", {}).get(
            "label", input_data.question_context.category
        )
    return input_data


def router(input_data: InputData):
    if input_data.question_context.category:
        return main_chain, extract_for_main_chain(input_data), input_data
    else:
        print("Using tagging chain")
        return tagging_chain, extract_for_tagging(input_data), input_data


# Define the full chain
def process_chain_result(args):
    chain, chain_input, original_input_data = args
    result = chain.invoke(chain_input)
    return result, original_input_data


full_chain = (
    RunnableLambda(router)
    | RunnableLambda(process_chain_result)
    | RunnableLambda(lambda args: update_input_data(args[1], args[0]))
)

# Example usage
input_data = InputData(
    messages=[Message(role="user", content="Hello")],
    question_context=QuestionContext(
        current_description="Issue with blue balloon",
        current_comments=None,
        current_question="My Balloon is pink and I do not want this.",
        current_context="It might be that there was a mistake in the order.",
        category=None,
    ),
)
print("Input data before:")
pprint(input_data.dict())
# Run the full chain
result = full_chain.invoke(input_data)
print("Final result:")
print(result)
print(f"Updated category: {result.question_context.category}")

print("Input data after:")
pprint(input_data.dict())
´´´

The llm is defined outside of this code. This seems rather cumbersome. Is there a better way to do this? Is this even the correct approach?

[dosubot[bot]]

To streamline your approach and make it less cumbersome, consider the following improvements:

1. Simplify Chain Execution: Instead of manually handling the chain execution and updating the input data, you can encapsulate these operations within a single function or class. This reduces redundancy and improves readability.

2. Use a State Management Pattern: Implement a state management pattern where the InputData object is updated in a more structured way. This can be done by creating a method within the InputData class to handle updates, which keeps the logic encapsulated.

3. Leverage LangChain's Built-in Features: Ensure you are using LangChain's built-in features effectively, such as structured outputs and routing, to minimize custom logic.

Here's a refined version of your setup:

from pydantic import BaseModel, Field
from typing import List, Optional, Literal
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain.llms import OpenAI  # Assuming you're using OpenAI

# Define your models
class ClassificationResponse(BaseModel):
    label: Literal["Red Balloon", "Blue Balloon", "Green Balloon", "Other Balloon color"] = Field(
        ..., description="The predicted class labels for the question"
    )

class Message(BaseModel):
    role: str
    content: str

class QuestionContext(BaseModel):
    current_description: str
    current_comments: Optional[str]
    current_question: str
    current_context: str
    category: Optional[str]

class InputData(BaseModel):
    messages: List[Message]
    question_context: QuestionContext

    def update_category(self, category: str):
        self.question_context.category = category

# Define chains
main_chain_prompt = PromptTemplate(
    template="Category: {category}\nQuestion: {question}\nContext: {context}\n\nPlease provide a response based on the given category and context.",
    input_variables=["category", "question", "context"],
)

main_chain = LLMChain(llm=llm, prompt=main_chain_prompt)

tagging_prompt = PromptTemplate(
    template="Question: {question}\nContext: {context}\n\nPlease determine the category of this text.",
    input_variables=["question", "context"],
)
tagging_chain = tagging_prompt | llm.with_structured_output(ClassificationResponse)

# Helper functions
def extract_for_tagging(input_data: InputData):
    return {
        "question": input_data.question_context.current_question,
        "context": input_data.question_context.current_context,
    }

def extract_for_main_chain(input_data: InputData):
    return {
        "category": input_data.question_context.category,
        "question": input_data.question_context.current_question,
        "context": input_data.question_context.current_context,
    }

def router(input_data: InputData):
    if input_data.question_context.category:
        return main_chain, extract_for_main_chain(input_data)
    else:
        return tagging_chain, extract_for_tagging(input_data)

# Define the full chain
def process_chain(input_data: InputData):
    chain, chain_input = router(input_data)
    result = chain.invoke(chain_input)
    if isinstance(result, ClassificationResponse):
        input_data.update_category(result.label)
    elif isinstance(result, dict):
        input_data.update_category(result.get("output", {}).get("label", input_data.question_context.category))
    return input_data

# Example usage
input_data = InputData(
    messages=[Message(role="user", content="Hello")],
    question_context=QuestionContext(
        current_description="Issue with blue balloon",
        current_comments=None,
        current_question="My Balloon is pink and I do not want this.",
        current_context="It might be that there was a mistake in the order.",
        category=None,
    ),
)

# Run the full chain
result = process_chain(input_data)
print(f"Updated category: {result.question_context.category}")

This version encapsulates the update logic within the InputData class and simplifies the chain execution process. It also uses a single function to handle the routing and processing, making the code cleaner and easier to maintain.

To continue talking to Dosu, mention @dosu.

---

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