Skip to main content

RAG with Agents

This is an agent specifically optimized for doing retrieval when necessary and also holding a conversation.

To start, we will set up the retriever we want to use, and then turn it into a retriever tool. Next, we will use the high level constructor for this type of agent. Finally, we will walk through how to construct a conversational retrieval agent from components.

The Retriever​

To start, we need a retriever to use! The code here is mostly just example code. Feel free to use your own retriever and skip to the section on creating a retriever tool.

from langchain.document_loaders import TextLoader

loader = TextLoader("../../../../../docs/docs/modules/state_of_the_union.txt")
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS

documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
embeddings = OpenAIEmbeddings()
db = FAISS.from_documents(texts, embeddings)
retriever = db.as_retriever()

Retriever Tool​

Now we need to create a tool for our retriever. The main things we need to pass in are a name for the retriever as well as a description. These will both be used by the language model, so they should be informative.

from langchain.agents.agent_toolkits import create_retriever_tool
tool = create_retriever_tool(
    retriever,
    "search_state_of_union",
    "Searches and returns documents regarding the state-of-the-union.",
)
tools = [tool]

Agent Constructor​

Here, we will use the high level create_conversational_retrieval_agent API to construct the agent.

Notice that beside the list of tools, the only thing we need to pass in is a language model to use. Under the hood, this agent is using the OpenAIFunctionsAgent, so we need to use an ChatOpenAI model.

from langchain.agents.agent_toolkits import create_conversational_retrieval_agent
from langchain.chat_models import ChatOpenAI

llm = ChatOpenAI(temperature=0)
agent_executor = create_conversational_retrieval_agent(llm, tools, verbose=True)

We can now try it out!

result = agent_executor({"input": "hi, im bob"})


> Entering new AgentExecutor chain...
Hello Bob! How can I assist you today?

> Finished chain.
result["output"]
'Hello Bob! How can I assist you today?'

Notice that it remembers your name

result = agent_executor({"input": "whats my name?"})


> Entering new AgentExecutor chain...
Your name is Bob.

> Finished chain.
result["output"]
'Your name is Bob.'

Notice that it now does retrieval

result = agent_executor(
    {
        "input": "what did the president say about kentaji brown jackson in the most recent state of the union?"
    }
)


> Entering new AgentExecutor chain...

Invoking: `search_state_of_union` with `{'query': 'Kentaji Brown Jackson'}`


[Document(page_content='Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n\nTonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n\nOne of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n\nAnd I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.', metadata={'source': '../../../../../docs/docs/modules/state_of_the_union.txt'}), Document(page_content='One was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. \n\nWhen they came home, many of the world’s fittest and best trained warriors were never the same. \n\nHeadaches. Numbness. Dizziness. \n\nA cancer that would put them in a flag-draped coffin. \n\nI know. \n\nOne of those soldiers was my son Major Beau Biden. \n\nWe don’t know for sure if a burn pit was the cause of his brain cancer, or the diseases of so many of our troops. \n\nBut I’m committed to finding out everything we can. \n\nCommitted to military families like Danielle Robinson from Ohio. \n\nThe widow of Sergeant First Class Heath Robinson.  \n\nHe was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. \n\nStationed near Baghdad, just yards from burn pits the size of football fields. \n\nHeath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter.', metadata={'source': '../../../../../docs/docs/modules/state_of_the_union.txt'}), Document(page_content='A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling.  \n\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers.  \n\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.', metadata={'source': '../../../../../docs/docs/modules/state_of_the_union.txt'}), Document(page_content='We can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. \n\nI recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. \n\nThey were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. \n\nOfficer Mora was 27 years old. \n\nOfficer Rivera was 22. \n\nBoth Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. \n\nI spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \n\nI’ve worked on these issues a long time. \n\nI know what works: Investing in crime prevention and community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety.', metadata={'source': '../../../../../docs/docs/modules/state_of_the_union.txt'})]In the most recent state of the union, the President mentioned Kentaji Brown Jackson. The President nominated Circuit Court of Appeals Judge Ketanji Brown Jackson to serve on the United States Supreme Court. The President described Judge Ketanji Brown Jackson as one of our nation's top legal minds who will continue Justice Breyer's legacy of excellence.

> Finished chain.
result["output"]
"In the most recent state of the union, the President mentioned Kentaji Brown Jackson. The President nominated Circuit Court of Appeals Judge Ketanji Brown Jackson to serve on the United States Supreme Court. The President described Judge Ketanji Brown Jackson as one of our nation's top legal minds who will continue Justice Breyer's legacy of excellence."

Notice that the follow up question asks about information previously retrieved, so no need to do another retrieval

result = agent_executor({"input": "how long ago did he nominate her?"})


> Entering new AgentExecutor chain...
The President nominated Judge Ketanji Brown Jackson four days ago.

> Finished chain.
result["output"]
'The President nominated Judge Ketanji Brown Jackson four days ago.'

Creating from components​

What actually is going on underneath the hood? Let’s take a look so we can understand how to modify going forward.

There are a few components:

  • The memory
  • The prompt template
  • The agent
  • The agent executor
# This is needed for both the memory and the prompt
memory_key = "history"

The Memory​

In this example, we want the agent to remember not only previous conversations, but also previous intermediate steps. For that, we can use AgentTokenBufferMemory. Note that if you want to change whether the agent remembers intermediate steps, or how the long the buffer is, or anything like that you should change this part.

from langchain.agents.openai_functions_agent.agent_token_buffer_memory import (
    AgentTokenBufferMemory,
)

memory = AgentTokenBufferMemory(memory_key=memory_key, llm=llm)

The Prompt Template​

For the prompt template, we will use the OpenAIFunctionsAgent default way of creating one, but pass in a system prompt and a placeholder for memory.

from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
from langchain.prompts import MessagesPlaceholder
from langchain_core.messages import SystemMessage
system_message = SystemMessage(
    content=(
        "Do your best to answer the questions. "
        "Feel free to use any tools available to look up "
        "relevant information, only if necessary"
    )
)
prompt = OpenAIFunctionsAgent.create_prompt(
    system_message=system_message,
    extra_prompt_messages=[MessagesPlaceholder(variable_name=memory_key)],
)

The Agent​

We will use the OpenAIFunctionsAgent

agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)

The Agent Executor​

Importantly, we pass in return_intermediate_steps=True since we are recording that with our memory object

from langchain.agents import AgentExecutor
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True,
    return_intermediate_steps=True,
)
result = agent_executor({"input": "hi, im bob"})


> Entering new AgentExecutor chain...
Hello Bob! How can I assist you today?

> Finished chain.
result = agent_executor({"input": "whats my name"})


> Entering new AgentExecutor chain...
Your name is Bob.

> Finished chain.