Add New Retriever Interface with Callbacks (#5962)

Handle the new retriever events in a way that (I think) is entirely
backwards compatible? Needs more testing for some of the chain changes
and all.

This creates an entire new run type, however. We could also just treat
this as an event within a chain run presumably (same with memory)

Adds a subclass initializer that upgrades old retriever implementations
to the new schema, along with tests to ensure they work.

First commit doesn't upgrade any of our retriever implementations (to
show that we can pass the tests along with additional ones testing the
upgrade logic).

Second commit upgrades the known universe of retrievers in langchain.

- [X] Add callback handling methods for retriever start/end/error (open
to renaming to 'retrieval' if you want that)
- [X] Update BaseRetriever schema to support callbacks
- [X] Tests for upgrading old "v1" retrievers for backwards
compatibility
- [X] Update existing retriever implementations to implement the new
interface
- [X] Update calls within chains to .{a]get_relevant_documents to pass
the child callback manager
- [X] Update the notebooks/docs to reflect the new interface
- [X] Test notebooks thoroughly


Not handled:
- Memory pass throughs: retrieval memory doesn't have a parent callback
manager passed through the method

---------

Co-authored-by: Nuno Campos <nuno@boringbits.io>
Co-authored-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
This commit is contained in:
Zander Chase
2023-06-30 14:44:03 -07:00
committed by GitHub
parent a5b206caf3
commit b0859c9b18
60 changed files with 2401 additions and 232 deletions

View File

@@ -71,11 +71,13 @@
"import numpy as np\n",
"\n",
"from langchain.schema import BaseRetriever\n",
"from langchain.callbacks.manager import AsyncCallbackManagerForRetrieverRun, CallbackManagerForRetrieverRun\n",
"from langchain.utilities import GoogleSerperAPIWrapper\n",
"from langchain.embeddings import OpenAIEmbeddings\n",
"from langchain.chat_models import ChatOpenAI\n",
"from langchain.llms import OpenAI\n",
"from langchain.schema import Document"
"from langchain.schema import Document\n",
"from typing import Any"
]
},
{
@@ -97,11 +99,16 @@
" def __init__(self, search):\n",
" self.search = search\n",
"\n",
" def get_relevant_documents(self, query: str):\n",
" def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any) -> List[Document]:\n",
" return [Document(page_content=self.search.run(query))]\n",
"\n",
" async def aget_relevant_documents(self, query: str):\n",
" raise NotImplemented\n",
" async def _aget_relevant_documents(self,\n",
" query: str,\n",
" *,\n",
" run_manager: AsyncCallbackManagerForRetrieverRun,\n",
" **kwargs: Any,\n",
" ) -> List[Document]:\n",
" raise NotImplementedError()\n",
"\n",
"\n",
"retriever = SerperSearchRetriever(GoogleSerperAPIWrapper())"

View File

@@ -43,7 +43,6 @@
"import openai\n",
"from dotenv import load_dotenv\n",
"from langchain.embeddings.openai import OpenAIEmbeddings\n",
"from langchain.schema import BaseRetriever\n",
"from langchain.vectorstores.azuresearch import AzureSearch"
]
},