From 23b2bf2358883efe26e00dba995259dd82fe710f Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Mon, 10 Jun 2024 12:59:06 -0300 Subject: [PATCH] feat: Add SelfQueryRetrieverComponent to langflow retrievers --- .../retrievers/SelfQueryRetriever.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 src/backend/base/langflow/components/retrievers/SelfQueryRetriever.py diff --git a/src/backend/base/langflow/components/retrievers/SelfQueryRetriever.py b/src/backend/base/langflow/components/retrievers/SelfQueryRetriever.py new file mode 100644 index 000000000..a46bf7bd8 --- /dev/null +++ b/src/backend/base/langflow/components/retrievers/SelfQueryRetriever.py @@ -0,0 +1,39 @@ +# from langflow.field_typing import Data +from langchain.chains.query_constructor.base import AttributeInfo +from langchain.retrievers.self_query.base import SelfQueryRetriever +from langchain_core.vectorstores import VectorStore + +from langflow.custom import CustomComponent +from langflow.field_typing import BaseLanguageModel +from langflow.schema import Record +from langflow.schema.message import Message + + +class SelfQueryRetrieverComponent(CustomComponent): + display_name: str = "Self Query Retriever" + description: str = "Retriever that uses a vector store and an LLM to generate the vector store queries." + icon = "LangChain" + + def build( + self, + query: Message, + vectorstore: VectorStore, + metadata_field_info: list[AttributeInfo], + document_content_description: str, + llm: BaseLanguageModel, + ) -> Record: + metadata_field_info = [i[0] for i in metadata_field_info] + + self_query_retriever = SelfQueryRetriever.from_llm( + llm, + vectorstore, + document_content_description, + metadata_field_info, + enable_limit=True, + ) + + input_text = query.text + documents = self_query_retriever.invoke(input=input_text) + records = [Record.from_document(document) for document in documents] + self.status = records + return records