A basic LangChain.js chain with prompt template, structured JSON output and OpenAI / Ollama LLMs

This is the fourth part of my AI-Boosted Development series:

The introduction summarized recent changes enabling easy AI app prototyping in TypeScript.
The second article covered installing Jupyter Lab IDE and tools for rapid prototyping.
The third part explained how to use JupyterLab and demonstrated basic coding workflow.

This article focuses on building the first part of the "Text Reviewer" app. As I explained in the introduction, the app works in two steps:

The user enters text, which an LLM model reviews and improves.
The tool compares the LLM's result to the original text and shows the changes.

We'll cover the first step here, showing a basic LangChain chain that reviews and improves text.

The second step doesn't use a language model. If we ask a language model to list changes compared to the original text, it often returns incorrect or incomplete results. These models can generate textual outputs or improve a text's quality, but they're not good at algorithmic tasks. They process prompts and try to predict the expected answer. They do not "think" or possess human-like logical capabilities. However, we can use a language model to generate the source code of a text comparison function. We'll use an LLM to generate code for text comparison in future articles.

Language models can respond in different formats, such as Markdown, JSON, or XML. In the examples below, we ask the models to provide JSON responses in a predefined schema. JSON responses work well if the schema is simple and the response doesn't contain many special characters.

Prompt template + OpenAI model + JSON output

In the example below, we implement the reviewTextOpenAI function with the following signature:

/**
 * Reviews and corrects the input text using OpenAI's GPT-4o model.
 * 
 * @param instruction - Instructions to be given to the model.
 * @param inputText - The text to be reviewed and corrected.
 * @returns - The reviewed and corrected text, or undefined if the response is invalid.
 */
reviewTextOpenAI(instruction: string, inputText: string): Promise<string | undefined>

Example function call and output:

// Define the instruction and input text for the prompt
const instruction = "Fix the grammar issues in the following text.";
const inputText = "How to stays relevant as the developer in the world of ai?";

// Log the result of the review function (OpenAI version)
console.log(await reviewTextOpenAI(instruction, inputText));
// CONSOLE: How to stay relevant as a developer in the world of AI?

The reviewTextOpenAI function does the following:

Creates a prompt template.
Defines a JSON schema using Zod.
Creates an language model (GPT-4o) wrapper, that returns the response in the format we defined with our JSON schema. We use the .withStructuredOutput method to get JSON output from the model.
Connects the prompt template with the language model to create a chain.
Calls the chain with the provided inputText and instruction

Here is the complete code for the function (you can download it in a Jupyter Notebook):

/**
 * Reviews and corrects the input text using OpenAI's GPT-4o model.
 * 
 * @param instruction - Instructions to be given to the model.
 * @param inputText - The text to be reviewed and corrected.
 * @returns - The reviewed and corrected text, or undefined if the response is invalid.
 */
async function reviewTextOpenAI(instruction: string, inputText: string): Promise<string | undefined> {
  // Create a prompt template using the provided instruction and input text
  const prompt = PromptTemplate.fromTemplate(
`{instruction}
---
{inputText}`);

  // Initialize the OpenAI chat model with specified options
  const llm = new ChatOpenAI({
    modelName: "gpt-4o", // Use the GPT-4 model
    verbose: false, // Disable verbose logging
  });

  // Define the schema for the model's output, it contains the reviewed text  
  const reviewedTextSchema = z.object({
    reviewedText: z.string().describe("The reviewed text.") // The reviewed text must be a string
  });
  type ReviewedTextSchema = z.infer<typeof reviewedTextSchema>; // Infer the TypeScript type from the Zod schema

  // We expect structured JSON output, we achieve this using OpenAI's function calling feature
  const llmWithStructuredOutput = llm.withStructuredOutput(reviewedTextSchema, {
    method: "functionCalling",
    name: "withStructuredOutput"
  });  

  // Create a processing chain combining the prompt and the LLM
  const chain = prompt.pipe(llmWithStructuredOutput);

  // Invoke the chain with the instruction and input text, and wait for the response
  const response: ReviewedTextSchema = await chain.invoke({ instruction, inputText });

  // Return the reviewed text if present in the response, otherwise undefined  
  return response?.reviewedText;
}

// Define the instruction and input text for the prompt
const instruction = "Fix the grammar issues in the following text.";
const inputText = "How to stays relevant as the developer in the world of ai?";

// show the reviewed text returned by the LLM
console.log(await reviewTextOpenAI(instruction, inputText));

Let's break down each part of the code. The function reviewTextOpenAI takes two parameters: an instruction and the input text. It returns a reviewed and corrected text or undefined if the response is invalid:

async function reviewTextOpenAI(instruction: string, inputText: string): Promise<string | undefined> {

We use the declarative LangChain Expression Language (LCEL) to compose chains. The first element of our chain is the prompt template that has two parameters: instruction and inputText. We assign values to these parameters when we execute the chain.

const prompt = PromptTemplate.fromTemplate(
`{instruction}
---
{inputText}`);

We initialize the OpenAI chat model wrapper. We use the gpt-4o model and disable verbose logging. Learn more about the ChatOpenAI model wrapper here.

const llm = new ChatOpenAI({
 modelName: "gpt-4o",
 verbose: false,
});

We define the schema for the model's output using the Zod library. The schema specifies that the output should have a property reviewedText, which must be a string. Then, we use z.infer to create a TypeScript type from this schema. This ensures our code knows exactly what type of data to expect.

const reviewedTextSchema = z.object({
 reviewedText: z.string().describe("The reviewed text.")
});
type ReviewedTextSchema = z.infer<typeof reviewedTextSchema>; // Infer the TypeScript type from the Zod schema

We configure the model to produce structured JSON output using Langchain's .withStructuredOutput() method:

const llmWithStructuredOutput = llm.withStructuredOutput(reviewedTextSchema, {
 method: "functionCalling",
 name: "withStructuredOutput"
});

We create a processing chain that combines the prompt and the model configured for structured output. Learn more about Langchain's LCEL chains and the pipe() method here.

const chain = prompt.pipe(llmWithStructuredOutput);

Finally, we invoke the processing chain with the instruction and input text, then wait for the response. We return the reviewed text if it is present in the response. Otherwise, we return undefined.

const response: ReviewedTextSchema = await chain.invoke({ instruction, inputText });   
return response?.reviewedText;

The OpenAI API requires an API key. In JupyterLab, we provide this key with the OPENAI_API_KEY environment variable, so we have to create a .env file with an OpenAI API key:

OPENAI_API_KEY=[your OpenAI API key]

Then we use the import "https://deno.land/std@0.215.0/dotenv/load.ts"; statement for reading the .env file and set the environment variables.

Prompt template + Ollama model + JSON output

Ollama-based models need a different approach for JSON output. LangChain's .withStructuredOutput doesn't support Ollama yet, so we use the OllamaFunctions wrapper's function calling feature.

Example function call and output:

// Define the instruction and input text for the prompt
const instruction = "Fix the grammar issues in the following text.";
const inputText = "How to stays relevant as the developer in the world of ai?";

// Log the result of the review function (Ollama version)
console.log(await reviewTextOllama(instruction, inputText));
// CONSOLE: How to stay relevant as a developer in the world of AI?

The steps are like those of the OpenAI implementations, with one difference: the method used to get JSON the output. The reviewTextOllama function does the following:

Creates a prompt template.
Defines a JSON schema using Zod.
Creates an LLM (Ollama / Codellama) wrapper that returns the response in the format defined by our JSON schema. We use function calling to get JSON output from the model.
Connects the prompt template with the language model to create a chain.
Calls the chain with the given inputText and instruction

Here is the complete code for the function (you can download it in a Jupyter Notebook):

/**
 * Processes a given text using a language model to review and correct it based on the provided instruction.
 * 
 * @param instruction - Instruction for the language model on how to process the text.
 * @param inputText - The text that needs to be reviewed and corrected.
 * @returns The reviewed text if successful, otherwise undefined.
 */
async function reviewTextOllama(instruction: string, inputText: string): Promise<string | undefined> {
  // Create a prompt template by combining the instruction and input text
  const prompt = PromptTemplate.fromTemplate(
`{instruction}
---
{inputText}`);

  // Define a schema for the expected output using zod
  const reviewedTextSchema = z.object({
    reviewedText: z.string().describe("The reviewed text.") // Define the structure and description of the reviewed text
  });
  type ReviewedTextSchema = z.infer<typeof reviewedTextSchema>; // Infer the TypeScript type from the zod schema

  // Initialize the language model with specific configuration
  const llm = new OllamaFunctions({
    baseUrl: "http://localhost:11434", // Base URL for the language model server
    model: "codellama:7b-code", // Specify the model to use
    verbose: false, // Disable verbose logging
  }).bind({
    functions: [
      {
        name: "storeResultTool", // Function name used in the language model
        description: "Gets the reviewed text", // Description of the function
        parameters: {
          type: "object", // Define the type of parameters expected by the function
          properties: zodToJsonSchema(reviewedTextSchema), // Convert zod schema to JSON schema
        },
      },
    ],
    function_call: {
      name: "storeResultTool", // Specify the function to be called
    },
  });

  // Create a processing chain: prompt -> language model -> JSON output parser
  const chain = prompt.pipe(llm).pipe(new JsonOutputFunctionsParser());

  // Invoke the chain with the instruction and input text
  const response = await chain.invoke({ instruction, inputText });

  // Return the reviewed text if available
  return response?.reviewedText;
}

// Define the instruction and input text for the prompt
const instruction = "Fix the grammar issues in the following text.";
const inputText = "How to stays relevant as the developer in the world of ai?";

// Log the result of the review function
console.log(await reviewTextOllama(instruction, inputText));

Let's see how we set up the Ollama wrapper to use the codellama model with JSON response in our code.

  // Initialize the language model with specific configuration
  const llm = new OllamaFunctions({
    baseUrl: "http://localhost:11434", // Base URL for the language model server
    model: "codellama:7b-code", // Specify the model to use
    verbose: false, // Enable verbose logging
  }).bind({
    functions: [
      {
        name: "storeResultTool", // Function name used in the language model
        description: "Gets the reviewed text", // Description of the function
        parameters: {
          type: "object", // Define the type of parameters expected by the function
          properties: zodToJsonSchema(reviewedTextSchema), // Convert zod schema to JSON schema
        },
      },
    ],
    function_call: {
      name: "storeResultTool", // Specify the function to be called
    },
  });

When we create the Ollama wrapper (OllamaFunctions) , we pass a configuration object to it with the model's name and the baseUrl for the Ollama server.
We use the .bind function on the created OllamaFunctions instance to define the storeResultTool function. This function's parameter has the reviewedTextSchema schema, the schema for our expected response. The function_call.name = 'storeResultTool' configuration option forces the model send the response to the storeResultTool function.

We create a processing chain that combines the prompt and the model configured for structured output. This chain has an extra step compared to the OpenAI implementation: using the JsonOutputFunctionsParser to convert the JSON output string to an object:

  // Create a processing chain: prompt -> language model -> JSON output parser
  const chain = prompt.pipe(llm).pipe(new JsonOutputFunctionsParser());

Summary

Congrats on completing this tutorial! In the fourth part of the AI-Boosted Development series, I showed how to create a basic LLM chain using LangChain.js. We used prompt templates, got structured JSON output, and integrated with OpenAI and Ollama LLMs. In the next article, I will show how to generate a function that compares two strings character by character and returns the differences in an HTML string. It is going to be exciting, so make sure to subscribe!

👨‍💻About the author

My name is Gergely Szerovay, I worked as a data scientist and full-stack developer for many years, and I have been working as frontend tech lead, focusing on Angular-based frontend development. As part of my role, I'm constantly following how Angular and the frontend development scene in general are evolving.

Angular has been advancing very rapidly over the past few years, and in the past year, with the rise of generative AI, our software development workflows have also evolved rapidly. In order to closely follow the evolution of AI-assisted software development, I decided to start building AI tools in public, and publish my progress on AIBoosted.dev , Subscribe here 🚀

Follow me on Substack (Angular Addicts), Substack (AIBoosted.dev), Medium, Dev.to, X or LinkedIn to learn more about Angular, and how to build AI apps with AI, Typescript, React and Angular!

Blog

A basic LangChain.js chain with prompt template, structured JSON output and OpenAI / Ollama LLMs

Gergely Szerovay

Prompt template + OpenAI model + JSON output

Prompt template + Ollama model + JSON output

Summary

👨‍💻About the author

Join Our Newsletter. No Spam, Only the good stuff.

Related