Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import configparser
 
-from typing import List, Union, Optional, Any, Dict
+from typing import List, Union, Optional, Any, Dict, cast
 import re
 import sys
 import time
@@ -14,9 +14,11 @@ import threading
 import pandas as pd
 from langchain import SerpAPIWrapper, LLMChain
 from langchain.agents import Tool, AgentType, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
+from langchain.callbacks.manager import Callbacks
 from langchain.callbacks.streaming_stdout_final_only import FinalStreamingStdOutCallbackHandler
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import LLMChain, SimpleSequentialChain
+from langchain.chains.query_constructor.ir import StructuredQuery
 from langchain.chains.query_constructor.base import AttributeInfo
 from langchain.document_loaders import DataFrameLoader, SeleniumURLLoader
 from langchain.embeddings import OpenAIEmbeddings
@@ -24,7 +26,7 @@ from langchain.indexes import VectorstoreIndexCreator
 from langchain.prompts import PromptTemplate, StringPromptTemplate, load_prompt, BaseChatPromptTemplate
 from langchain.llms import OpenAI
 from langchain.retrievers.self_query.base import SelfQueryRetriever
-from langchain.schema import AgentAction, AgentFinish, HumanMessage
+from langchain.schema import AgentAction, AgentFinish, HumanMessage, Document
 from langchain.vectorstores import Chroma
 
 import gradio as gr
@@ -112,7 +114,30 @@ metadata_field_info = [
 wine_vectorstore = Chroma.from_documents(docs, embeddings)
 document_content_description = "A database of wines. 'name' and 'pairing' must be included in the query, and 'Body', 'Tannin', 'Sweetness', 'Alcohol', 'Price', 'Rating', 'Wine_Type', and 'Country' can be included in the filter. query and filter must be form of 'key: value'. For example, query: 'name: 돔페리뇽, pairing:육류'."
 llm = OpenAI(temperature=0)
-wine_retriever = SelfQueryRetriever.from_llm(
+
+class CustomSelfQueryRetriever(SelfQueryRetriever):
+    async def aget_relevant_documents(self, query: str, callbacks: Callbacks = None) -> List[Document]:
+        inputs = self.llm_chain.prep_inputs({"query": query})
+        structured_query = cast(
+            StructuredQuery,
+            self.llm_chain.predict_and_parse(callbacks=callbacks, **inputs),
+        )
+        if self.verbose:
+            print(structured_query)
+        new_query, new_kwargs = self.structured_query_translator.visit_structured_query(
+            structured_query
+        )
+        if structured_query.limit is not None:
+            new_kwargs["k"] = structured_query.limit
+
+        if self.use_original_query:
+            new_query = query
+
+        search_kwargs = {**self.search_kwargs, **new_kwargs}
+        docs = self.vectorstore.search(new_query, self.search_type, **search_kwargs)
+        return docs
+
+wine_retriever = CustomSelfQueryRetriever.from_llm(
     llm, wine_vectorstore, document_content_description, metadata_field_info, verbose=True
 )  # Added missing closing parenthesis
 
@@ -150,7 +175,7 @@ metadata_field_info = [
 
 document_content_description = "Database of a winebar"
 llm = OpenAI(temperature=0)
-wine_bar_retriever = SelfQueryRetriever.from_llm(
+wine_bar_retriever = CustomSelfQueryRetriever.from_llm(
     llm, wine_bar_vectorstore, document_content_description, metadata_field_info=metadata_field_info, verbose=True
 )  # Added missing closing parenthesis
 #### Tool3: Search in Google