Generating Structured Output From OpenAI, Anthropic, and Azure AI Using LangChain following Pydantic Approach

Image generated by Dalle

Introduction

LangChain is an amazing framework providing easy access to various large language models (LLMs) such as OpenAI, Anthropic, and Azure AI. Often, obtaining structured output from these LLMs is crucial for further processing or storage. I present a simple technique utilizing Pydantic models to generate structured output from these three LLMs.

Approach

Our goal is to standardize the output received from LLMs so that we can easily manipulate the answers programmatically. Our strategy consists of four primary components:

1. Select the preferred LLM (OpenAI, Anthropic, or Azure AI)
2. Define the question asked to the LLM
3. Create a common Pydantic schema for the output
4. Parse the raw output and convert it into the Pydantic structure

Code Implementation

import os
import json
import re
from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import PromptTemplate
from langchain_anthropic import ChatAnthropic

os.environ["OPENAI_API_KEY"] = "<your_openai_api_key>"
os.environ["ANTHROPIC_API_KEY"] = "<your_anthropic_api_key>"

template = """Provide answer in valid JSON format only. Answer should be consistent, clear and concise. User question is: ```{human_input}``` Follow this Pydantic specifications for the output: 
            {{ class format(BaseModel): 
               answer: str 
               follow_up_questions: List[str] }} 
            Dont include class name in the beginining."""

prompt = PromptTemplate(input_variables=["human_input"], template=template)

def chat_with_llm():
    model_type = input("Please choose 'openai', 'azure', or 'anthropic': ").lower()
    human_input = input("Please enter your question: ")

    if model_type == "openai":
        model_version = input("Please choose 'gpt-3.5-turbo' or 'gpt-4': ").lower()

        if model_version == "gpt-3.5-turbo":
            model = ChatOpenAI(temperature=0, model="gpt-3.5-turbo")
        elif model_version == "gpt-4":
            model = ChatOpenAI(temperature=0, model="gpt-4")
        else:
            raise ValueError("Invalid model version. Please choose 'gpt-3.5-turbo' or 'gpt-4'.")

        chatgpt_chain = LLMChain(
            llm=model,
            prompt=prompt,
            verbose=True,
            memory=ConversationBufferWindowMemory(k=5),
            llm_kwargs={"max_length": 4096},
        )
    elif model_type == "azure":
        os.environ["AZURE_OPENAI_API_KEY"] = "..."
        os.environ["AZURE_OPENAI_ENDPOINT"] = "..."
        os.environ["OPENAI_API_VERSION"] = "...

        model = AzureChatOpenAI(
            openai_api_version="...,
            azure_deployment="...",
        )

        chatgpt_chain = LLMChain(
            llm=model,
            prompt=prompt,
            verbose=True,
            memory=ConversationBufferWindowMemory(k=5),
        )
    elif model_type == "anthropic":
        model = ChatAnthropic(temperature=0, anthropic_api_key="<your_anthropic_api_key>", model_name="...")

        chatgpt_chain = LLMChain(
            llm=model,
            prompt=prompt,
            verbose=True,
            memory=ConversationBufferWindowMemory(k=5),
        )
    else:
        raise ValueError("Invalid model type. Please choose 'openai', 'azure', or 'anthropic'.")

    input_data = {"human_input": human_input}
    output = chatgpt_chain.invoke(input=input_data)

    json_output = {}
    if model_type != "azure":
        response_text = output['text']
        clean_response = re.sub(r'^\{"json\_meta.*"\}\n|^\{\n|\}$', '', response_text, flags=re.DOTALL)
        try:
            json_output = json.loads(clean_response)
            if isinstance(json_output, dict):
                if "answer" in json_output and "follow_up_questions" in json_output:
                    return json_output
                else:
                    return {"answer": clean_response, "follow_up_questions": []}
            else:
                return {"answer": clean_response, "follow_up_questions": []}
        except json.JSONDecodeError:
            return {"answer": clean_response, "follow_up_questions": []}
    else:
        output_dict = json.loads(output['text'])
        json_output = {
            'answer': output_dict['answer'].replace('\n', ' '),
            'follow_up_questions': [q.replace('\n', ' ') for q in output_dict['follow_up_questions']]
        }

    return json_output

output = chat_with_llm()
print(json.dumps(output, indent=4))

• The code starts by importing the necessary modules: os, json, re, langchain_openai, langchain.chains, langchain.memory, langchain.prompts, and langchain_anthropic.
• Then, it defines environment variables, OPENAI_API_KEY, AZUREAI_API_KEY and ANTHROPIC_API_KEY, assigning placeholders for actual API keys.
• Following this, a prompt template called template is created, which includes instructions for the LLM to reply in valid JSON format with a strict requirement for the output to follow a specific Pydantic model.
• The chat_with_llm() function definition initiates next:
• Retrieves user input for the desired model type (‘openai’, ‘azure’, or ‘anthropic’).
• Collects user input for the question.
• Determines the correct LLM based on the user input.
• Constructs the LLMChain by passing the relevant LLM, prompt, and additional configurations.
• Within the LLMChain instantiation, ConversationBufferWindowMemory ensures that the conversational context is taken into account for producing the most informed replies.
• Upon constructing the LLMChain, the user’s question is passed to the invoke() function, retrieving the generated response.
• Subsequently, the code cleans the response, stripping away unwanted metadata tags, and attempts to decode the cleaned response as JSON.
• Should the JSON parsing fail, the JSONDecodeError exception handler returns the response as plain text alongside empty follow-up questions. Otherwise, the response gets parsed and checked against the Pydantic model.
• Finally, the full response is printed out in a neat, indented fashion using the json.dumps() function.

By following this approach, you can easily obtain structured output from OpenAI, Anthropic, and Azure AI large language models using LangChain and the Pydantic library. The structured output can then be seamlessly integrated into your applications or used for further processing and analysis.