Skip to main content


This example goes over how to use LangChain to interact with AI21 models.


!pip install -qU langchain-ai21

Environment Setup

We'll need to get a AI21 API key and set the AI21_API_KEY environment variable:

import os
from getpass import getpass

os.environ["AI21_API_KEY"] = getpass()


from langchain_ai21 import AI21LLM
from langchain_core.prompts import PromptTemplate

template = """Question: {question}

Answer: Let's think step by step."""

prompt = PromptTemplate.from_template(template)

model = AI21LLM(model="j2-ultra")

chain = prompt | model

chain.invoke({"question": "What is LangChain?"})
API Reference:AI21LLM | PromptTemplate
'\nLangChain is a (database)\nLangChain is a database for storing and processing documents'

AI21 Contextual Answer

You can use AI21's contextual answers model to receives text or document, serving as a context, and a question and returns an answer based entirely on this context.

This means that if the answer to your question is not in the document, the model will indicate it (instead of providing a false answer)

from langchain_ai21 import AI21ContextualAnswers

tsm = AI21ContextualAnswers()

response = tsm.invoke(input={"context": "Your context", "question": "Your question"})
API Reference:AI21ContextualAnswers

You can also use it with chains and output parsers and vector DBs

from langchain_ai21 import AI21ContextualAnswers
from langchain_core.output_parsers import StrOutputParser

tsm = AI21ContextualAnswers()
chain = tsm | StrOutputParser()

response = chain.invoke(
{"context": "Your context", "question": "Your question"},

Was this page helpful?

You can leave detailed feedback on GitHub.