🦜🔗 LangChain Integration

LMQL can also be used together with the 🦜🔗 LangChain python library. Both, using langchain libraries from LMQL code and using LMQL queries as part of chains are supported.

Using LangChain from LMQL

We first consider the case, where one may want to use LangChain modules as part of an LMQL query.

In this example, we want to leverage the LangChain Chroma retrieval model, to enable question answering about a text document (the LMQL paper in this case).

First, we need to import the required libraries.

import lmql
import asyncio
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma

Next, we load and embed the text of the relevant document (lmql.txt in our case).

# load text of LMQL paper
with open("lmql.txt") as f:
    contents = f.read()
texts = []
for i in range(0, len(contents), 120):
    texts.append(contents[i:i+120])

embeddings = OpenAIEmbeddings()
docsearch = Chroma.from_texts(texts, embeddings,
    metadatas=[{"text": t} for t in texts], persist_directory="lmql-index")

We then construct chatbot using a simple LMQL query, that first prompts the user for a question via await input(...), retrieves relevant text paragraphs using LangChain and then produces an answer using openai/gpt-3.5-turbo (ChatGPT).

import termcolor

@lmql.query
async def chatbot():
        '''
sample(temperature=0.2, max_len=2048, openai_chunksize=2048)
    """You are a chatbot that helps users answer questions.
    You are first provided with the question and relevant information."""
    while True:
        q = await input("\\nQuestion: ")
        "Question: {q}\\n"
        print(termcolor.colored("Reading relevant pages...", "green"))
        results = set([d.page_content for d in docsearch.similarity_search(q, 4)])
        information = "\\n\\n".join(["..." + r + "..." for r in list(results)])
        "\\nRelevant Information: {information}\\n"
        "Your response based on relevant information:[RESPONSE]"
from
    "openai/gpt-3.5-turbo"
        '''

await chatbot(output_writer=lmql.stream(variable="RESPONSE"))

Question: What is LMQL?
Reading relevant pages...
LMQL stands for Language Model Query Language. It is a high-level language with declarative SQL-like elements and an imperative syntax for querying language models. LMQL allows for great expressiveness and supports scripted prompting and just a set of high-level constraints. It has been evaluated comprehensively and shown to be able to express a wide range of techniques.

As shown in the query, inline LMQL code appearing in a Python script can access the outer scope containing e.g. the docsearch variable, and access any relevant utility functions and object defined in Python.

Using LMQL from LangChain

In addition to using langchain utilities in LMQL query code, LMQL queries can also seamlessly be integrated as a langchain Chain component.

For this consider, the sequential prompting example from the langchain documentation, where we first prompt the language model to propose a company name for a given product, and then ask it for a catchphrase.

To get started, we first import the relevant langchain components, as well as LMQL.

from langchain import LLMChain, PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (ChatPromptTemplate,HumanMessagePromptTemplate)
from langchain.llms import OpenAI

import lmql

Our chain has two stages: (1) Asking the model for a company name, and (2) asking the model for a catchphrase. For the sake of this example, we will implement (1) in with a langchain prompt and (2) with an LMQL query.

First, we define the langchain prompt for the company name and instantiate the resulting LLMChain:

# setup the LM to be used by langchain
llm = OpenAI(temperature=0.9)

human_message_prompt = HumanMessagePromptTemplate(
        prompt=PromptTemplate(
            template="What is a good name for a company that makes {product}?",
            input_variables=["product"],
        )
    )
chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])
chat = ChatOpenAI(temperature=0.9)
chain = LLMChain(llm=chat, prompt=chat_prompt_template)

This can already be executed to produce a company name:

chain.run("colorful socks")

'Rainbow Sock Co.'

Next, we define prompt (2) in LMQL, i.e. the LMQL query generating the catchphrase:

@lmql.query
async def write_catch_phrase(company_name: str):
    '''
    argmax "Write a catchphrase for the following company: {company_name}. [CATCHPHRASE]" from "chatgpt"
    '''

Again, we can run this part in isolation, like so:

(await write_catch_phrase("Socks Inc"))[0].variables["CATCHPHRASE"]

' "Step up your style with Socks Inc."'

To chain the two prompts together, we can use a SimpleSequentialChain from langchain:

from langchain.chains import SimpleSequentialChain
overall_chain = SimpleSequentialChain(chains=[chain, write_catch_phrase], verbose=True)

Now, we can run the overall chain, relying both on LMQL and langchain components:

# Run the chain specifying only the input variable for the first chain.
catchphrase = overall_chain.run("colorful socks")
print(catchphrase)

> Entering new SimpleSequentialChain chain...
RainbowSocks Co.
 "Step into a world of color with RainbowSocks Co.!"

> Finished chain.
 "Step into a world of color with RainbowSocks Co.!"

Note that the full chain currently only works when run outside a Jupyter environment (in a script), due to the way langchain and LMQL rely on async vs. sequential calls. With further improvements to the langchain async API, this limitation can be expected to be removed.