Structured Output & LLM Routers in Langchain

Title Image

Introduction

In the realm of building applications using large language models (LLMs), particularly in Retrieval-Augmented Generation (RAG) applications, an emerging concept known as LLM routers is gaining traction. LLM routers serve as a pivotal component that directs the flow of operations based on intermediate actions validated by the LLM. This routing mechanism ensures that applications execute the correct functions, enhancing efficiency and accuracy.

Understanding LLM Routers

LLM routers are designed to manage the workflow within an application by leveraging the capabilities of LLMs to analyze and validate intermediate actions. For instance, in a RAG system, the LLM router can direct user queries to the appropriate collections based on the nature of the query. Furthermore, once a response is generated, the LLM router can act as a response grader, validating the accuracy and relevance of the output before presenting it to the user.

To achieve this functionality, it is crucial to obtain structured outputs from the LLM. Structured outputs can be attained through various methods, including:

1. Prompting for Structured Output: By designing prompts that request specific structured data, the LLM can be guided to generate responses in a desired format.
2. Using Structured Output Parsers: LangChain offers structured output parsers that can be utilized to format the LLM’s output appropriately.
3. Pydantic Models: A highly favored approach involves using Pydantic models to ensure the LLM’s output adheres to a predefined structure. This method provides robustness and clarity in handling the data generated by the LLM.

Consider the following scenarios where we can use LLM routers:

• Query Routing: When a user poses a query, the LLM router evaluates the query’s content and directs it to the most relevant collection within the vector store database. This ensures that the response generated is both accurate and contextually appropriate.
• Field Extraction for API Integration: Based on the user query, the LLM router can extract specific content or fields from the query and pass them as payload to an API endpoint. This allows the application to hit the endpoint and present the data fetched from the API.
• Response Validation: After generating a response, the LLM router acts as a response grader, verifying the validity and relevance of the output. If the response does not meet the required standards, the LLM router can trigger further actions to refine or regenerate the response.

By implementing these mechanisms, applications can harness the full potential of LLMs, ensuring optimal performance and user satisfaction.

Utilizing LangChain for Structured Output and LLM Routing

LangChain provides a robust framework for working with language models, enabling structured output through prompt templates and output parsers. This guide will walk through three examples: a basic string prompt template, a JSON output parser, and a Pydantic output parser. Each example showcases how to achieve structured outputs, which are essential for implementing LLM routers that route the outputs to different applications based on the content.

Example 1: String Prompt Template

The first example demonstrates a simple prompt template for customer support. The task is to provide sentiment and summarization of user remarks and return the results in JSON format.

from langchain_core.prompts import PromptTemplate
import os
from langchain_openai import ChatOpenAI

os.environ["OPENAI_API_KEY"] = "your_openai_api_key"

model = ChatOpenAI(temperature=0)
prompt_template = PromptTemplate.from_template(
    "You are a customer support assistant and your task is to provide the sentiment and summarization of the user remarks. Return summary and sentiment in JSON format. Below is the remarks provided by customer: {remarks}"
)
chain = prompt_template | model
response = chain.invoke({"remarks": "Hi, I have ordered a peter england t white t shirt a brand new one but I received a brown t shirt with some stinks on it and the packaging is also not proper. I just wanted a complete refund for my order and I am really disappointed about the service you gave."})

print(response.content)

In this example, we use a string prompt template to format a single string. The output is returned as a string in JSON format, which can be converted to valid JSON later.

Example 2: JSON Output Parser

To directly obtain a JSON response, we can use the JsonOutputParser. This parser simplifies the process and ensures the output is in the desired JSON format.

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
import os
from langchain_openai import ChatOpenAI

os.environ["OPENAI_API_KEY"] = "your_openai_api_key"

model = ChatOpenAI(temperature=0)
prompt_template = ChatPromptTemplate.from_messages([
    ("system", "You are a customer support assistant and your task is to provide the sentiment and summarization of the user remarks. Return summary and sentiment in JSON format."),
    ("user", " {query}")
])

chain = prompt_template | model | JsonOutputParser()
response = chain.invoke({"query": "Hi, I have ordered a peter england t white t shirt a brand new one but I received a brown t shirt with some stinks on it and the packaging is also not proper. I just wanted a complete refund for my order and I am really disappointed about the service you gave."})

print(response)

With the JsonOutputParser, the response is directly obtained in JSON format, facilitating further processing or integration with other applications.

Example 3: Pydantic Output Parser

For more complex use cases, such as validating and extracting specific fields from a query, we can use the PydanticOutputParser. This example demonstrates extracting journey details from a travel query and ensuring the validity of the data using Pydantic models.

from langchain.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field, validator
from langchain_openai import ChatOpenAI
import os

os.environ["OPENAI_API_KEY"] = "your_openai_api_key"

# Define your desired data structure
class JourneyDetails(BaseModel):
    journey_date: str = Field(description="Date of the journey in yyyy-mm-dd format")
    destination_name: str = Field(description="Destination of the journey")
    source_name: str = Field(description="Source or origin of the journey")
    
    @validator("destination_name")
    def validate_destination_name(cls, field):
        if field == 'unknown':
            raise ValueError("Please mention the destination!")
        return field

# Query to prompt a language model to populate the data structure
journey_query = "give me mumbai to chennai trains for aug 12."

# Set up a parser and inject instructions into the prompt template
parser = PydanticOutputParser(pydantic_object=JourneyDetails)
prompt = PromptTemplate(
    template="Given the query, extract the fields such as journey_date, destination_name, source_name.\n{format_instructions}\n{query}\n",
    input_variables=["query"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

model = ChatOpenAI(temperature=0)
chain = prompt | model | parser
result = chain.invoke({"query": journey_query})

print(result)

In this example, we create a Pydantic model with fields journey_date, destination_name, and source_name. The destination_name field is validated to ensure it's not set to 'unknown'. After extracting the data, it can be used to construct a payload for an API request or other processing tasks.

Beyond just the structure of the Pydantic class, the name of the Pydantic class, the docstring, and the names and provided descriptions of parameters are very important

The .with_structured_output() method

The `with_structured_output()` method is a simple and reliable way to get structured outputs from models that support native APIs for structuring data, like tool/function calling or JSON mode. This method takes a schema specifying the desired output attributes’ names, types, and descriptions. It returns a model-like Runnable that outputs objects matching the given schema instead of just strings or messages. If a JSON Schema is used, the method returns a dictionary; if a Pydantic class is used, it returns Pydantic objects. This approach ensures clear, structured data that is easy to work with in your applications.

from typing import Optional
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
import os

os.environ["OPENAI_API_KEY"] = "your_openai_api_key"

llm = ChatOpenAI(temperature=0)

class Joke(BaseModel):
    """Joke to tell user."""

    setup: str = Field(description="The setup of the joke")
    punchline: str = Field(description="The punchline to the joke")
    rating: Optional[int] = Field(description="How funny the joke is, from 1 to 10")


structured_llm = llm.with_structured_output(Joke)

structured_llm.invoke("Tell me a joke about cats")

output: Joke(setup='Why was the cat sitting on the computer?', punchline='To keep an eye on the mouse!', rating=8)

We can also pass in a JSON Schema dict if you prefer not to use Pydantic.

json_schema = {
    "title": "joke",
    "description": "Joke to tell user.",
    "type": "object",
    "properties": {
        "setup": {
            "type": "string",
            "description": "The setup of the joke",
        },
        "punchline": {
            "type": "string",
            "description": "The punchline to the joke",
        },
        "rating": {
            "type": "integer",
            "description": "How funny the joke is, from 1 to 10",
        },
    },
    "required": ["setup", "punchline"],
}
structured_llm = llm.with_structured_output(json_schema)

structured_llm.invoke("Tell me a joke about cats")

Output:{'setup': 'Why was the cat sitting on the computer?',
 'punchline': 'To keep an eye on the mouse!',
 'rating': 8}

Refer more about this structured output here : How to return structured data from a model | 🦜️🔗 LangChain

These examples demonstrate how LangChain can be utilized to achieve structured outputs, which are crucial for routing the results to different applications based on the content. By leveraging prompt templates and output parsers, you can create a flexible and powerful system for handling various language model outputs.

Conclusion:

Implementing LLM routers in LangChain streamlines workflows in RAG systems by using structured outputs through prompt templates and parsers like JSON and Pydantic. This ensures data from LLMs is accurate and actionable, facilitating tasks like query routing, API field extraction, and response validation.

Thank you for exploring LangChain’s capabilities for structured output and LLM routing. I hope this guide has been helpful.

References:

Langchain Docs: Introduction | 🦜️🔗 LangChain