From 6af613e2b986a319e0b12e030d7e36f58b4de2d3 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Mon, 21 Aug 2023 22:05:11 -0700 Subject: [PATCH] retrieval agents --- .../integrations/retrievers/agent.ipynb | 205 ++++++++++++++++++ .../agents/retrieval_agent/__init__.py | 0 .../langchain/agents/retrieval_agent/base.py | 162 ++++++++++++++ .../agents/retrieval_agent/prompt.py | 41 ++++ 4 files changed, 408 insertions(+) create mode 100644 docs/extras/integrations/retrievers/agent.ipynb create mode 100644 libs/langchain/langchain/agents/retrieval_agent/__init__.py create mode 100644 libs/langchain/langchain/agents/retrieval_agent/base.py create mode 100644 libs/langchain/langchain/agents/retrieval_agent/prompt.py diff --git a/docs/extras/integrations/retrievers/agent.ipynb b/docs/extras/integrations/retrievers/agent.ipynb new file mode 100644 index 00000000000..c24eb9c8696 --- /dev/null +++ b/docs/extras/integrations/retrievers/agent.ipynb @@ -0,0 +1,205 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "45d73d1f", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.document_loaders import TextLoader\n", + "loader = TextLoader('../../../../docs/extras/modules/state_of_the_union.txt')" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f5a61aff", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.text_splitter import CharacterTextSplitter\n", + "from langchain.vectorstores import FAISS\n", + "from langchain.embeddings import OpenAIEmbeddings\n", + "\n", + "documents = loader.load()\n", + "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n", + "texts = text_splitter.split_documents(documents)\n", + "embeddings = OpenAIEmbeddings()\n", + "db = FAISS.from_documents(texts, embeddings)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "daabf286", + "metadata": {}, + "outputs": [], + "source": [ + "retriever = db.as_retriever()" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "cab036df", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.agents.agent_toolkits import create_retriever_tool" + ] + }, + { + "cell_type": "code", + "execution_count": 78, + "id": "a4576106", + "metadata": {}, + "outputs": [], + "source": [ + "tool = create_retriever_tool(\n", + " retriever, \n", + " \"search_state_of_union\",\n", + " \"Searches and returns documents regarding the state-of-the-union. The input should be a search query\"\n", + ")\n", + "tools = [tool]" + ] + }, + { + "cell_type": "code", + "execution_count": 79, + "id": "e5df294c", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 87, + "id": "3a3ef33a", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 91, + "id": "f70a3fce", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.chat_models import ChatAnthropic, ChatOpenAI\n", + "\n", + "model = ChatAnthropic(model='claude-2', temperature=0)\n", + "model = ChatOpenAI(model=\"gpt-4\", temperature=0)" + ] + }, + { + "cell_type": "code", + "execution_count": 92, + "id": "fc100599", + "metadata": {}, + "outputs": [], + "source": [ + "chain = LLMChain(\n", + " llm=model,\n", + " prompt=XMLRetrievalAgent.get_default_prompt(),\n", + " output_parser=XMLRetrievalAgent.get_default_output_parser()\n", + ")\n", + "agent = XMLRetrievalAgent(tools=tools, llm_chain=chain)" + ] + }, + { + "cell_type": "code", + "execution_count": 93, + "id": "9bc1473e", + "metadata": {}, + "outputs": [], + "source": [ + "ae = AgentRetriever(agent=agent, tools=tools, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 94, + "id": "2e17c298", + "metadata": { + "scrolled": false + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\n", + "\u001b[1m> Entering new AgentRetriever chain...\u001b[0m\n", + "\u001b[32;1m\u001b[1;3msearch_state_of_union\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "justice breyer state of the union\u001b[0m\u001b[36;1m\u001b[1;3m[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/extras/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/extras/modules/state_of_the_union.txt'}), Document(page_content='And built the strongest, freest, and most prosperous nation the world has ever known. \\n\\nNow is the hour. \\n\\nOur moment of responsibility. \\n\\nOur test of resolve and conscience, of history itself. \\n\\nIt is in this moment that our character is formed. Our purpose is found. Our future is forged. \\n\\nWell I know this nation. \\n\\nWe will meet the test. \\n\\nTo protect freedom and liberty, to expand fairness and opportunity. \\n\\nWe will save democracy. \\n\\nAs hard as these times have been, I am more optimistic about America today than I have been my whole life. \\n\\nBecause I see the future that is within our grasp. \\n\\nBecause I know there is simply nothing beyond our capacity. \\n\\nWe are the only nation on Earth that has always turned every crisis we have faced into an opportunity. \\n\\nThe only nation that can be defined by a single word: possibilities. \\n\\nSo on this night, in our 245th year as a nation, I have come to report on the State of the Union.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='We can do all this while keeping lit the torch of liberty that has led generations of immigrants to this land—my forefathers and so many of yours. \\n\\nProvide a pathway to citizenship for Dreamers, those on temporary status, farm workers, and essential workers. \\n\\nRevise our laws so businesses have the workers they need and families don’t wait decades to reunite. \\n\\nIt’s not only the right thing to do—it’s the economically smart thing to do. \\n\\nThat’s why immigration reform is supported by everyone from labor unions to religious leaders to the U.S. Chamber of Commerce. \\n\\nLet’s get it done once and for all. \\n\\nAdvancing liberty and justice also requires protecting the rights of women. \\n\\nThe constitutional right affirmed in Roe v. Wade—standing precedent for half a century—is under attack as never before. \\n\\nIf we want to go forward—not backward—we must protect access to health care. Preserve a woman’s right to choose. And let’s continue to advance maternal health care in America.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'})]\u001b[0m\u001b[32;1m\u001b[1;3msearch_state_of_unionukraine state of the union\u001b[0m\u001b[36;1m\u001b[1;3m[Document(page_content='Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. \\n\\nIn this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight. \\n\\nLet each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world. \\n\\nPlease rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people. \\n\\nThroughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos. \\n\\nThey keep moving. \\n\\nAnd the costs and the threats to America and the world keep rising. \\n\\nThat’s why the NATO Alliance was created to secure peace and stability in Europe after World War 2. \\n\\nThe United States is a member along with 29 other nations. \\n\\nIt matters. American diplomacy matters. American resolve matters.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='We see the unity among leaders of nations and a more unified Europe a more unified West. And we see unity among the people who are gathering in cities in large crowds around the world even in Russia to demonstrate their support for Ukraine. \\n\\nIn the battle between democracy and autocracy, democracies are rising to the moment, and the world is clearly choosing the side of peace and security. \\n\\nThis is a real test. It’s going to take time. So let us continue to draw inspiration from the iron will of the Ukrainian people. \\n\\nTo our fellow Ukrainian Americans who forge a deep bond that connects our two nations we stand with you. \\n\\nPutin may circle Kyiv with tanks, but he will never gain the hearts and souls of the Ukrainian people. \\n\\nHe will never extinguish their love of freedom. He will never weaken the resolve of the free world. \\n\\nWe meet tonight in an America that has lived through two of the hardest years this nation has ever faced. \\n\\nThe pandemic has been punishing.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans. \\n\\nLast year COVID-19 kept us apart. This year we are finally together again. \\n\\nTonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \\n\\nWith a duty to one another to the American people to the Constitution. \\n\\nAnd with an unwavering resolve that freedom will always triumph over tyranny. \\n\\nSix days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. \\n\\nHe thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. \\n\\nHe met the Ukrainian people. \\n\\nFrom President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}), Document(page_content='For that purpose we’ve mobilized American ground forces, air squadrons, and ship deployments to protect NATO countries including Poland, Romania, Latvia, Lithuania, and Estonia. \\n\\nAs I have made crystal clear the United States and our Allies will defend every inch of territory of NATO countries with the full force of our collective power. \\n\\nAnd we remain clear-eyed. The Ukrainians are fighting back with pure courage. But the next few days weeks, months, will be hard on them. \\n\\nPutin has unleashed violence and chaos. But while he may make gains on the battlefield – he will pay a continuing high price over the long run. \\n\\nAnd a proud Ukrainian people, who have known 30 years of independence, have repeatedly shown that they will not tolerate anyone who tries to take their country backwards. \\n\\nTo all Americans, I will be honest with you, as I’ve always promised. A Russian dictator, invading a foreign country, has costs around the world.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'})]\u001b[0m\u001b[32;1m\u001b[1;3m04567\u001b[0m\n", + "\n", + "\u001b[1m> Finished chain.\u001b[0m\n" + ] + }, + { + "data": { + "text/plain": [ + "[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/extras/modules/state_of_the_union.txt'}),\n", + " Document(page_content='Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. \\n\\nIn this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight. \\n\\nLet each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world. \\n\\nPlease rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people. \\n\\nThroughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos. \\n\\nThey keep moving. \\n\\nAnd the costs and the threats to America and the world keep rising. \\n\\nThat’s why the NATO Alliance was created to secure peace and stability in Europe after World War 2. \\n\\nThe United States is a member along with 29 other nations. \\n\\nIt matters. American diplomacy matters. American resolve matters.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}),\n", + " Document(page_content='We see the unity among leaders of nations and a more unified Europe a more unified West. And we see unity among the people who are gathering in cities in large crowds around the world even in Russia to demonstrate their support for Ukraine. \\n\\nIn the battle between democracy and autocracy, democracies are rising to the moment, and the world is clearly choosing the side of peace and security. \\n\\nThis is a real test. It’s going to take time. So let us continue to draw inspiration from the iron will of the Ukrainian people. \\n\\nTo our fellow Ukrainian Americans who forge a deep bond that connects our two nations we stand with you. \\n\\nPutin may circle Kyiv with tanks, but he will never gain the hearts and souls of the Ukrainian people. \\n\\nHe will never extinguish their love of freedom. He will never weaken the resolve of the free world. \\n\\nWe meet tonight in an America that has lived through two of the hardest years this nation has ever faced. \\n\\nThe pandemic has been punishing.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}),\n", + " Document(page_content='Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans. \\n\\nLast year COVID-19 kept us apart. This year we are finally together again. \\n\\nTonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \\n\\nWith a duty to one another to the American people to the Constitution. \\n\\nAnd with an unwavering resolve that freedom will always triumph over tyranny. \\n\\nSix days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. \\n\\nHe thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. \\n\\nHe met the Ukrainian people. \\n\\nFrom President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'}),\n", + " Document(page_content='For that purpose we’ve mobilized American ground forces, air squadrons, and ship deployments to protect NATO countries including Poland, Romania, Latvia, Lithuania, and Estonia. \\n\\nAs I have made crystal clear the United States and our Allies will defend every inch of territory of NATO countries with the full force of our collective power. \\n\\nAnd we remain clear-eyed. The Ukrainians are fighting back with pure courage. But the next few days weeks, months, will be hard on them. \\n\\nPutin has unleashed violence and chaos. But while he may make gains on the battlefield – he will pay a continuing high price over the long run. \\n\\nAnd a proud Ukrainian people, who have known 30 years of independence, have repeatedly shown that they will not tolerate anyone who tries to take their country backwards. \\n\\nTo all Americans, I will be honest with you, as I’ve always promised. A Russian dictator, invading a foreign country, has costs around the world.', metadata={'source': '../../../../docs/extras/modules/state_of_the_union.txt'})]" + ] + }, + "execution_count": 94, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ae.get_relevant_documents(\"what did he say about justice breyer and ukraine?\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4639d501", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3663162a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/libs/langchain/langchain/agents/retrieval_agent/__init__.py b/libs/langchain/langchain/agents/retrieval_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/libs/langchain/langchain/agents/retrieval_agent/base.py b/libs/langchain/langchain/agents/retrieval_agent/base.py new file mode 100644 index 00000000000..a144af6a5b9 --- /dev/null +++ b/libs/langchain/langchain/agents/retrieval_agent/base.py @@ -0,0 +1,162 @@ +import xml.etree.ElementTree as ET +from typing import Any, List, Tuple, Union + +from langchain.agents.agent import ( + AgentExecutor, + AgentOutputParser, + BaseSingleActionAgent, +) +from langchain.agents.xml.prompt import agent_instructions +from langchain.callbacks.base import Callbacks +from langchain.callbacks.manager import CallbackManagerForRetrieverRun +from langchain.chains.llm import LLMChain +from langchain.prompts.chat import AIMessagePromptTemplate, ChatPromptTemplate +from langchain.schema import AgentAction, AgentFinish, BaseRetriever, Document +from langchain.tools.base import BaseTool + + +class XMLAgentOutputParser(AgentOutputParser): + """Output parser for XMLAgent.""" + + def parse(self, text: str) -> Union[AgentAction, AgentFinish]: + if "" in text: + tool, tool_input = text.split("") + _tool = tool.split("")[1] + _tool_input = tool_input.split("")[1] + return AgentAction(tool=_tool, tool_input=_tool_input, log=text) + elif "" in text: + _, answer = text.split("") + return AgentFinish(return_values={"output": answer}, log=text) + else: + raise ValueError + + def get_format_instructions(self) -> str: + raise NotImplementedError + + @property + def _type(self) -> str: + return "xml-retrieval-agent" + + +class XMLRetrievalAgent(BaseSingleActionAgent): + """Agent that uses XML tags to do retrieval. + + Args: + tools: list of tools the agent can choose from + llm_chain: The LLMChain to call to predict the next action + + Examples: + + .. code-block:: python + + from langchain.agents import XMLAgent + from langchain + + tools = ... + model = + + + """ + + tools: List[BaseTool] + """List of tools this agent has access to.""" + llm_chain: LLMChain + """Chain to use to predict action.""" + + @property + def input_keys(self) -> List[str]: + return ["input"] + + @staticmethod + def get_default_prompt() -> ChatPromptTemplate: + return ChatPromptTemplate.from_template( + agent_instructions + ) + AIMessagePromptTemplate.from_template("{intermediate_steps}") + + @staticmethod + def get_default_output_parser() -> XMLAgentOutputParser: + return XMLAgentOutputParser() + + def plan( + self, + intermediate_steps: List[Tuple[AgentAction, str]], + callbacks: Callbacks = None, + **kwargs: Any, + ) -> Union[AgentAction, AgentFinish]: + log = "" + _id = 0 + doc_mapping = {} + for action, observation in intermediate_steps: + doc_string = "" + for doc in observation: + doc_mapping[_id] = doc + doc_string += f"{_id}{doc.page_content}" + _id += 1 + log += ( + f"{action.tool}{action.tool_input}" + f"{doc_string}" + ) + tools = "" + for tool in self.tools: + tools += f"{tool.name}: {tool.description}\n" + inputs = { + "intermediate_steps": log, + "tools": tools, + "question": kwargs["input"], + "stop": ["", ""], + } + response = self.llm_chain(inputs, callbacks=callbacks) + result = response[self.llm_chain.output_key] + if isinstance(result, AgentAction): + return result + else: + root = ET.fromstring("" + result.return_values["output"] + "") + ids = [elem.text for elem in root.findall("id")] + docs = [doc_mapping[int(i)] for i in ids] + return AgentFinish(return_values={"output": docs}, log=result.log) + + async def aplan( + self, + intermediate_steps: List[Tuple[AgentAction, str]], + callbacks: Callbacks = None, + **kwargs: Any, + ) -> Union[AgentAction, AgentFinish]: + log = "" + _id = 0 + doc_mapping = {} + for action, observation in intermediate_steps: + doc_string = "" + for doc in observation: + doc_mapping[_id] = doc + doc_string += f"{_id}{doc.page_content}" + _id += 1 + log += ( + f"{action.tool}{action.tool_input}" + f"{doc_string}" + ) + tools = "" + for tool in self.tools: + tools += f"{tool.name}: {tool.description}\n" + inputs = { + "intermediate_steps": log, + "tools": tools, + "question": kwargs["input"], + "stop": ["", ""], + } + response = await self.llm_chain.acall(inputs, callbacks=callbacks) + result = response[self.llm_chain.output_key] + if isinstance(result, AgentAction): + return result + else: + root = ET.fromstring("" + result.return_values["output"] + "") + ids = [elem.text for elem in root.findall("id")] + docs = [doc_mapping[int(i)] for i in ids] + return AgentFinish(return_values={"output": docs}, log=result.log) + + +class AgentRetriever(BaseRetriever, AgentExecutor): + def _get_relevant_documents( + self, query: str, *, run_manager: CallbackManagerForRetrieverRun + ) -> List[Document]: + result = self({"input": query}, callbacks=run_manager.get_child()) + return result["output"] diff --git a/libs/langchain/langchain/agents/retrieval_agent/prompt.py b/libs/langchain/langchain/agents/retrieval_agent/prompt.py new file mode 100644 index 00000000000..a59e345c85b --- /dev/null +++ b/libs/langchain/langchain/agents/retrieval_agent/prompt.py @@ -0,0 +1,41 @@ +# flake8: noqa +agent_instructions = """You are a helpful research assistant who helps find relevant documents for the user's questions. + +You have access to the following tools: + +{tools} + +In order to use a tool, you MUST use AND tags. + +Each tool will return a list of documents. This will be returned in the format of: ......... +For example, if you have a tool called 'search' that could run a google search, in order to search for the weather in SF you would respond: + +search +weather in SF +164 degrees + +When you are done, respond with a list of the document ids that are relevant. The documents corresponding to these ids will be returned. For example: + +14... + +Only respond with the ids of the documents that are actually relevant to the question at hand. \ +You can make as many queries as are necessary in order to get the correct documents. + +Some example of how you should act: + +Scenario 1: +- The user asks for topic X +- You run a query for Y but don't get any good results +- You run another for Z and get better results, and so you return some from those + +Scenario 2: +- The user asks for topic X and Y +- You run a query for X +- You run a query for Y +- You return the relevant documents from each + +Ready? + +Begin! + +Question: {question}"""