Modifying LLM responses on the fly

Abenezer Belachew

Abenezer Belachew · July 20, 2023

10 min read

Intro

  • Recently, I’ve been playing around with Vercel’s AI SDK, a great tool that allows you to easily integrate language models into your projects. In one of my projects, I needed to customize the response stream from the LLM to meet specific requirements. However, as of writing this article, the SDK doesn’t offer a built-in method for such modifications. Therefore, I’ve decided to write this article to demonstrate how you can alter the responses to achieve a desired format.

  • While the example I am about to use is basic, the technique can be effectively adapted for more intricate tasks.

What are we making?

  • In this article, we'll build a lightweight markdown component that accepts plain text from an LLM stream as a prop. The component will then convert the URLs in the text into clickable links that open in new tabs. We'll even add emojis next to the link based on the industry sector the website covers. Let's get started! 🦘
Demo

Prerequisites

  • I am going to use NextJs13 and Vercel AI SDK but you can use the component in other react frameworks too.

Let’s start

> npx create-next-app@latest
  • I’ve named my project my-markdown but you can call it whatever you want. I’ve also picked the default option for all the questions.
Default options

Alright, cd into the new project you created and run it.

> cd my-markdown
> npm run dev
  • Your server will most likely run on http://localhost:3000. Head there and check if the default Next.JS template page is there. If it’s not, make sure you’re on the right port.

The AI SDK

  • We’ll be using Vercel’s AI SDK to communicate with an LLM. For our case, I’ve chosen OpenAI since, at the moment, the most popular and widely used one.
  • Quit your server and install the package.
npm install ai openai-edge

Add the OpenAI API key to .env.local

OPENAI_API_KEY=xxxxxxxxx
  • You can get your OpenAI API keys here.

Chat Route

import { Configuration, OpenAIApi } from 'openai-edge'
import { OpenAIStream, StreamingTextResponse } from 'ai'
 
// Create an OpenAI API client (that's edge friendly!)
const config = new Configuration({
  apiKey: process.env.OPENAI_API_KEY
})
const openai = new OpenAIApi(config)
 
// IMPORTANT! Set the runtime to edge
export const runtime = 'edge'
 
export async function POST(req: Request) {
  // Extract the `messages` from the body of the request
  const { messages } = await req.json()
 
  // Ask OpenAI for a streaming chat completion given the prompt
  const response = await openai.createChatCompletion({
    model: 'gpt-3.5-turbo',
    stream: true,
    messages
  })
  // Convert the response into a friendly text-stream
  const stream = OpenAIStream(response)
  // Respond with the stream
  return new StreamingTextResponse(stream)
}
  • Basically, this route extracts the messages from the request and sends them to the openai.CreateChatCompletion function from openai-edge. It includes the settings you pass in (model and stream) to convert and return the stream of responses in a friendly manner. We’ve also set stream to true so that the model can return partial responses before the entire conversation is complete.
  • Now, let’s now create a client component with a form that we will use to send and display responses.
  • I’ll be using the main page for this because it’s easier to create the UI: Head to src/app/page.tsx and add the following lines of code. There is a slight styling modification from the one on the docs.
"use client";
import { useChat } from "ai/react";

export default function Home() {
  const { messages, input, handleInputChange, handleSubmit } = useChat();

  return (
    <main className="container mt-12">
      <div className="flex flex-col items-center justify-center">
        <div className="max-w-3xl max-h-[600px] overflow-auto items-center">
          {messages.map((m) => (
            <div key={m.id}>
              <div className="font-semibold">
                {m.role === "user" ? "You: " : "AI: "}
              </div>
              <div className="p-2 rounded-md mb-2">
                {m.content}
              </div>
            </div>
          ))}
        </div>

        <form
          onSubmit={handleSubmit}
          className="fixed flex flex-row w-full max-w-md bottom-0 bg-gray-50 border border-gray-700 rounded mb-8 shadow-xl p-2"
        >
          <input
            className="border border-gray-300 rounded p-2 w-full"
            value={input}
            onChange={handleInputChange}
            placeholder="Type a message..."
          />
          <button type="submit" className="p-2 ml-4 bg-gray-300 rounded-md mt-4 align-middle">
            Send
          </button>
        </form>
      </div>
    </main>
  );
}

Now run your server, and check it out. If you see some weird lines like the pic below, head to globals.css and remove everything but the first three lines

@tailwind base;
@tailwind components;
@tailwind utilities;
Weird Lines

You should now see this:

No weird lines

Type in a message and hit send or enter.

Weird Lines

Nice! It seems to be working.

The problem

One problem here is that the links are treated as plain text. If we had trained the LLM on our own data and wanted to direct it to specific pages or open all links in a different tab, we wouldn't be able to achieve that with the current setup.

The solution

To address this, we can encapsulate the message element in a custom component that allows us to manipulate incoming streams as per our requirements. In this example, I'll detect all the links in a response and generate <a> tags with a _blank attribute, enabling them to open in a new tab.

  • Let's create a new file called my-markdown.tsx in the src/components/ directory for our component. This component will receive a text prop, identify the links within it, and generate <a> tags with a _blank attribute.
import React from 'react'

const MyMarkdown = ({ text }: { text: string }) => {
  // Get the whole regex from github 
  // (https://github.com/abenezerBelachew/modifying-llm-responses/blob/main/src/app/components/my-markdown.tsx#L27)
  const linkRegex = /((?:(http|https|Http|Https|rtsp|Rtsp): ...
  const parts = []

  let lastIndex = 0
  let match

  while ((match = linkRegex.exec(text)) !== null) {
    const [fullMatch, linkText] = match
    const matchStart = match.index
    const matchEnd = matchStart + fullMatch.length
    const linkUrl = linkText.startsWith('http') || linkText.startsWith('https://') ? linkText : `http://${linkText}`

    if (lastIndex < matchStart) {
      parts.push(text.slice(lastIndex, matchStart))
    }

    parts.push(
      <a
        target='_blank'
        rel='noopener noreferrer'
        className='break-words underline underline-offset-2 text-blue-600'
        href={linkUrl}>
        {linkText}
      </a>
    )

    lastIndex = matchEnd
  }

  if (lastIndex < text.length) {
    parts.push(text.slice(lastIndex))
  }

  return (
    <>
      {parts.map((part, i) => (
        <React.Fragment key={i}>{part}</React.Fragment>
      ))}
    </>
  )
}

export default MyMarkdown
  • I would be lying if I said I knew how to implement or fully explain the linkRegex but it seems to catch most type of urls. I found it on stackoverflow .
  • The purpose of this component is to take a text prop (a string containing some text content) and identify URLs within that text. For each URL found, it generates an anchor tag <a> with the attribute target='_blank' (to open the link in a new tab) and rel='noopener noreferrer' (best practice for security when opening links in a new tab).
  • The component uses a regular expression (linkRegex) to find URLs within the text prop. It then iterates through the text and separates it into multiple parts. Each part is either a text segment or an anchor tag generated for a URL. After identifying the URLs and creating anchor tags, it returns the parts as a React fragment.
  • When using this component with some text containing URLs, it will return the text with the URLs wrapped in anchor tags, allowing the URLs to be clickable and open in new tabs.
  • Keep in mind that the provided regular expression (linkRegex) is quite complex and can match a variety of URL formats. It is used to handle different cases like URLs with or without protocols (e.g., http://, https://), IP addresses, domain names, etc.

Cool, try it out.

  • Head to src/app/page.tsx and import our new component. After that, pass in m.content as a prop.
"use client";
import { useChat } from "ai/react";
import MyMarkdown from "./components/my-markdown";           // New

export default function Home() {
  const { messages, input, handleInputChange, handleSubmit } = useChat();

  return (
    <main className="container mt-12">
      <div className="flex flex-col items-center justify-center">
        <div className="max-w-3xl max-h-[600px] overflow-auto items-center">
          {messages.map((m) => (
            <div key={m.id}>
              <div className="font-semibold">
                {m.role === "user" ? "You: " : "AI: "}
              </div>
              <div className="p-2 rounded-md mb-2">
                <MyMarkdown text={m.content} />              // New
              </div>
            </div>
          ))}
        </div>
Demo

Doing interesting stuff with it

  • Say we have an emoji AI that takes in a URL and generates an emoji to represent the industry sector covered by that website. We can now incorporate this feature into our component.
  • First, let’s add our very fancyAI above our MyMarkdown object.
const veryFancyAI = (url: string) => {
    const education: string[] = ['www.google.com', 'www.wikipedia.org']
    const social: string[] = ['www.facebook.com', 'www.instagram.com', 'www.reddit.com', ]
    const movies: string[] = ['www.netflix.com']
    const shopping: string[] = ['www.amazon.com', 'www.ebay.com']
    const renting: string[] = ['www.airbnb.com']

    if (education.includes(url)) {
        return '🏫'
    } else if (social.includes(url)) {
        return '🫂'
    } else if (movies.includes(url)) {
        return '📽️'
    } else if (shopping.includes(url)) {
        return '🛒'
    } else if (renting.includes(url)) {
        return '🏠'
    } else {
        return '🌐'
    }
}
  • And then add it to below the linkUrl
const matchEnd = matchStart + fullMatch.length
const linkUrl = linkText.startsWith('http') || linkText.startsWith('https://') ? linkText : `https://${linkText}`
const emoji = veryFancyAI(linkText)           // New

if (lastIndex < matchStart) {
      parts.push(text.slice(lastIndex, matchStart))
 }

parts.push(
      <a
        target='_blank'
        rel='noopener noreferrer'
        className='break-words underline underline-offset-2 text-blue-600'
        href={linkUrl}>
        {linkText} 
        <span className='text-2xl border-2 p-2'>{emoji}</span>  // New
      </a>
  )

Here’s what the final code for MyMarkdown looks like:

import React from 'react'

const veryFancyAI = (url: string) => {
    const education: string[] = ['www.google.com', 'www.wikipedia.org']
    const social: string[] = ['www.facebook.com', 'www.instagram.com', 'www.reddit.com', ]
    const movies: string[] = ['www.netflix.com']
    const shopping: string[] = ['www.amazon.com', 'www.ebay.com']
    const renting: string[] = ['www.airbnb.com']

    if (education.includes(url)) {
        return '🏫'
    } else if (social.includes(url)) {
        return '🫂'
    } else if (movies.includes(url)) {
        return '📽️'
    } else if (shopping.includes(url)) {
        return '🛒'
    } else if (renting.includes(url)) {
        return '🏠'
    } else {
        return '🌐'
    }
}

const MyMarkdown = ({ text }: { text: string }) => {
  // Get the whole regex from github 
  // (https://github.com/abenezerBelachew/modifying-llm-responses/blob/main/src/app/components/my-markdown.tsx#L27)
  const linkRegex = /((?:(http|https|Http|Https|rtsp|Rtsp): ...
  const parts = []

  let lastIndex = 0
  let match

  while ((match = linkRegex.exec(text)) !== null) {
    const [fullMatch, linkText] = match
    const matchStart = match.index
    const matchEnd = matchStart + fullMatch.length
    const linkUrl = linkText.startsWith('http') || linkText.startsWith('https://') ? linkText : `https://${linkText}`
    const emoji = veryFancyAI(linkText)

    if (lastIndex < matchStart) {
      parts.push(text.slice(lastIndex, matchStart))
    }

    parts.push(
      <a
        target='_blank'
        rel='noopener noreferrer'
        className='break-words underline underline-offset-2 text-blue-600'
        href={linkUrl}>
        {linkText} 
        <span className='text-2xl border-2 p-2'>{emoji}</span>
      </a>
    )

    lastIndex = matchEnd
  }

  if (lastIndex < text.length) {
    parts.push(text.slice(lastIndex))
  }

  return (
    <>
      {parts.map((part, i) => (
        <React.Fragment key={i}>{part}</React.Fragment>
      ))}
    </>
  )
}

export default MyMarkdown
Demo
  • This is obviously a very simple example but you can let your imagination run wild. Now that you know how to modify the responses to look a certain way, you can craft great experiences for yourself and your users.
  • I wrote this article to highlight how even with minimal modifications, you can tailor the responses of your LLM to suit specific applications. If your LLM is trained on custom data and follows predefined rules for formatting, creating your own wrappers like this can empower you to achieve your desired outcomes. Embrace this flexibility to mold the LLM's responses and leverage its power in unique ways for your projects. The possibilities are endless, and you have the opportunity to create truly customized and impactful applications.
  • If you want to go into more details on how to render rich responses from LLMs, this article by Spencer Miskoviak is a great read.
  • The Github repo for this project can be found here.