mirror of
https://github.com/Mintplex-Labs/langchain-python.git
synced 2026-07-19 13:26:32 -04:00
Zander Chase
b0859c9b18
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>
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
"""Retriever that wraps a base retriever and filters the results."""
|
|
|
|
from typing import Any, List
|
|
|
|
from pydantic import BaseModel, Extra
|
|
|
|
from langchain.callbacks.manager import (
|
|
AsyncCallbackManagerForRetrieverRun,
|
|
CallbackManagerForRetrieverRun,
|
|
)
|
|
from langchain.retrievers.document_compressors.base import (
|
|
BaseDocumentCompressor,
|
|
)
|
|
from langchain.schema import BaseRetriever, Document
|
|
|
|
|
|
class ContextualCompressionRetriever(BaseRetriever, BaseModel):
|
|
"""Retriever that wraps a base retriever and compresses the results."""
|
|
|
|
base_compressor: BaseDocumentCompressor
|
|
"""Compressor for compressing retrieved documents."""
|
|
|
|
base_retriever: BaseRetriever
|
|
"""Base Retriever to use for getting relevant documents."""
|
|
|
|
class Config:
|
|
"""Configuration for this pydantic object."""
|
|
|
|
extra = Extra.forbid
|
|
arbitrary_types_allowed = True
|
|
|
|
def _get_relevant_documents(
|
|
self,
|
|
query: str,
|
|
*,
|
|
run_manager: CallbackManagerForRetrieverRun,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Get documents relevant for a query.
|
|
|
|
Args:
|
|
query: string to find relevant documents for
|
|
|
|
Returns:
|
|
Sequence of relevant documents
|
|
"""
|
|
docs = self.base_retriever.get_relevant_documents(
|
|
query, callbacks=run_manager.get_child(), **kwargs
|
|
)
|
|
if docs:
|
|
compressed_docs = self.base_compressor.compress_documents(
|
|
docs, query, callbacks=run_manager.get_child()
|
|
)
|
|
return list(compressed_docs)
|
|
else:
|
|
return []
|
|
|
|
async def _aget_relevant_documents(
|
|
self,
|
|
query: str,
|
|
*,
|
|
run_manager: AsyncCallbackManagerForRetrieverRun,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Get documents relevant for a query.
|
|
|
|
Args:
|
|
query: string to find relevant documents for
|
|
|
|
Returns:
|
|
List of relevant documents
|
|
"""
|
|
docs = await self.base_retriever.aget_relevant_documents(
|
|
query, callbacks=run_manager.get_child(), **kwargs
|
|
)
|
|
if docs:
|
|
compressed_docs = await self.base_compressor.acompress_documents(
|
|
docs, query, callbacks=run_manager.get_child()
|
|
)
|
|
return list(compressed_docs)
|
|
else:
|
|
return []
|