Use similarity retriever to provide context for '.respond'

app.py

@@ -18,8 +18,8 @@ with gr.Blocks() as demo:
         with gr.Column():
             url = gr.Textbox(value="https://www.gradio.app/docs/gradio/chatinterface", label="Docs URL", render=True)
             chat = gr.ChatInterface(
-                # model.respond,
-                model.predict,
+                model.respond,
+                # model.predict,
                 # model.rag,
                 additional_inputs=[
                     url,

data.py

@@ -18,27 +18,6 @@ class Data:
             encode_kwargs={'normalize_embeddings':True}
         )
 
-    def get_context(self):
-        urls = [
-            self.url,
-        ]
-        loader = UnstructuredURLLoader(urls=urls)
-        data = loader.load()
-    
-        context = data[0].page_content # will come from 'url'
-
-        return context
-
-    def build_prompt(self, question):
-        prompt = f"""
-        Use the following piece of context to answer the question asked.
-        Please try to provide the answer only based on the context
-        {self.get_context()}
-        Question:{question}
-        Helpful Answers:
-        """
-        return prompt
-
     @property
     def retriever(self):
         loader = UnstructuredURLLoader(urls=self.urls)

model.py

@@ -21,6 +21,18 @@ class Model:
         )
         self.chat_model = ChatHuggingFace(llm=self.llm, token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
+    def build_prompt(self, question, context_urls):
+        data = Data(context_urls)
+        context = data.retriever.invoke(f"{question}")[0].page_content
+        prompt = f"""
+        Use the following piece of context to answer the question asked.
+        Please try to provide the answer only based on the context
+        {context}
+        Question:{question}
+        Helpful Answers:
+        """
+        return prompt
+
     def _build_prompt_rag(self):
         prompt_template="""
         Use the following piece of context to answer the question asked.
@@ -64,9 +76,7 @@ class Model:
         temperature,
         top_p,
     ):
-    
-        data = Data(url)
-    
+        
         messages = [{"role": "system", "content": url}]
     
         for val in history:
@@ -75,7 +85,7 @@ class Model:
             if val[1]:
                 messages.append({"role": "assistant", "content": val[1]})
     
-        messages.append({"role": "user", "content": data.build_prompt(message)})
+        messages.append({"role": "user", "content": self.build_prompt(message, [url])})
     
         response = ""