Update app.py

app.py

@@ -41,7 +41,11 @@ class Agent1:
         chain = LLMChain(llm=self.model, prompt=rephrase_prompt)
         response = chain.run(query=user_input).strip()
         
-        return [q.strip() for q in response.split('\n') if q.strip()]
+        # Remove any lines that contain instructions or explanations
+        rephrased_queries = [q.strip() for q in response.split('\n') if q.strip() and not q.startswith("Rephrase") and "query" not in q.lower()]
+        
+        # If no valid rephrased queries, return the original input
+        return rephrased_queries if rephrased_queries else [user_input]
 
     def process(self, user_input: str) -> Dict[str, List[Dict[str, str]]]:
         queries = self.rephrase_and_split(user_input)
@@ -50,39 +54,6 @@ class Agent1:
             results[query] = google_search(query)
         return results
 
-class Agent2:
-    def __init__(self, model):
-        self.model = model
-
-    def validate_response(self, user_query: str, response: str) -> bool:
-        validation_prompt = PromptTemplate(
-            input_variables=["query", "response"],
-            template="""
-            Evaluate if the following response fully answers the user's query.
-            User query: {query}
-            Response: {response}
-            
-            Does the response fully answer the query? Answer with Yes or No:"""
-        )
-        
-        chain = LLMChain(llm=self.model, prompt=validation_prompt)
-        result = chain.run(query=user_query, response=response).strip().lower()
-        return result == 'yes'
-
-    def generate_follow_up_query(self, user_query: str, response: str) -> str:
-        follow_up_prompt = PromptTemplate(
-            input_variables=["query", "response"],
-            template="""
-            The following response did not fully answer the user's query.
-            User query: {query}
-            Response: {response}
-            
-            Generate a follow-up query to get more relevant information:"""
-        )
-        
-        chain = LLMChain(llm=self.model, prompt=follow_up_prompt)
-        return chain.run(query=user_query, response=response).strip()
-
 def load_document(file: NamedTemporaryFile) -> List[Document]:
     """Loads and splits the document into pages."""
     loader = PyPDFLoader(file.name)
@@ -270,7 +241,6 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
     model = get_model(temperature, top_p, repetition_penalty)
     embed = get_embeddings()
     agent1 = Agent1(model)
-    agent2 = Agent2(model)
 
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -279,7 +249,6 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
 
     max_attempts = 3
     context_reduction_factor = 0.7
-    agent2_max_attempts = 2
 
     for attempt in range(max_attempts):
         try:
@@ -350,20 +319,6 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
             else:
                 answer = full_response.strip()
 
-            for agent2_attempt in range(agent2_max_attempts):
-                if agent2.validate_response(question, answer):
-                    break
-                
-                if agent2_attempt < agent2_max_attempts - 1:
-                    follow_up_query = agent2.generate_follow_up_query(question, answer)
-                    follow_up_results = agent1.process(follow_up_query)
-                    follow_up_docs = [Document(page_content=result["text"], metadata={"source": result["link"], "query": follow_up_query}) for results in follow_up_results.values() for result in results if result["text"]]
-                    database.add_documents(follow_up_docs)
-                    context_str += "\n" + "\n".join([f"Follow-up Query: {doc.metadata['query']}\nSource: {doc.metadata['source']}\nContent: {doc.page_content}" for doc in follow_up_docs])
-                    formatted_prompt = prompt_val.format(context=context_str, question=question)
-                    full_response = generate_chunked_response(model, formatted_prompt)
-                    answer = full_response.strip()
-
             if web_search:
                 sources = set(doc.metadata['source'] for doc in web_docs)
                 sources_section = "\n\nSources:\n" + "\n".join(f"- {source}" for source in sources)