Update app.py

app.py

@@ -5,11 +5,10 @@ import pandas as pd
 from datetime import datetime
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain_community.llms import HuggingFacePipeline
-from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
+from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain.agents import Tool
 from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
-from langchain_community.utilities import TextRequestsWrapper
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 
@@ -30,83 +29,90 @@ pipe = pipeline(
 )
 llm = HuggingFacePipeline(pipeline=pipe)
 
-# --- System Message ---
-system_prompt = """You are a helpful assistant tasked with answering questions using a set of tools. 
-Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
-FINAL ANSWER: [YOUR FINAL ANSWER]. 
-YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
-system_message_prompt = SystemMessagePromptTemplate.from_template(system_prompt)
-
-# --- Tools ---
+# --- Tools Setup ---
 ddg = DuckDuckGoSearchAPIWrapper()
-requests_wrapper = TextRequestsWrapper()
 
-def wiki_search(query):
-    """Search Wikipedia for a query and return maximum 2 results."""
-    search_results = ddg.run(query)
-    return {"wiki_results": search_results}
+def enhanced_search(query):
+    """Enhanced search combining multiple sources"""
+    try:
+        # Web search
+        web_results = ddg.results(query, 3)
+        # Wikipedia search
+        wiki_results = ddg.results(f"wikipedia {query}", 2)
+        return {
+            "web": [r["snippet"] for r in web_results],
+            "wikipedia": [r["snippet"] for r in wiki_results]
+        }
+    except Exception as e:
+        print(f"Search error: {e}")
+        return {}
+
+# --- Prompt Engineering ---
+PROMPT_TEMPLATE = """Use the following context to answer the question. 
+If you don't know the answer, say you don't know. Keep answers very short.
+
+Context:
+{search_results}
+
+Question: {question}
 
-def web_search(query):
-    """Search Tavily for a query and return maximum 3 results."""
-    search_results = ddg.run(query)
-    return {"web_results": search_results}
+Think step by step, then write the final answer starting with FINAL ANSWER:"""
 
-def arxiv_search(query):
-    """Search Arxiv for a query and return maximum 3 results."""
-    url = f"https://export.arxiv.org/api/query?search_query=all:{query}&start=0&max_results=3"
-    response = requests_wrapper.get(url)
-    return {"arxiv_results": response.text}
+prompt = PromptTemplate(
+    template=PROMPT_TEMPLATE,
+    input_variables=["search_results", "question"]
+)
+
+# --- Answer Processing ---
+def process_answer(raw_answer: str) -> str:
+    """Extract and clean the final answer"""
+    if "FINAL ANSWER:" in raw_answer:
+        answer = raw_answer.split("FINAL ANSWER:")[-1].strip()
+        answer = answer.split('\n')[0].strip()
+        answer = answer[:MAX_ANSWER_LENGTH]
+        return answer
+    return raw_answer.strip()[:MAX_ANSWER_LENGTH]
 
 # --- Chroma DB Setup ---
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-
 vector_store = Chroma(
     embedding_function=embeddings,
     persist_directory="./chroma_db"
 )
 
-def create_retriever_tool(query):
-    """A tool to retrieve similar questions from a vector store."""
-    similar_question = vector_store.similarity_search(query)
-    return {"retriever_results": similar_question[0].page_content}
-
-tools = [
-    Tool(
-        name="Wikipedia Search",
-        func=wiki_search,
-        description="Search Wikipedia for a query and return maximum 2 results."
-    ),
-    Tool(
-        name="Web Search",
-        func=web_search,
-        description="Search Tavily for a query and return maximum 3 results."
-    ),
-    Tool(
-        name="Arxiv Search",
-        func=arxiv_search,
-        description="Search Arxiv for a query and return maximum 3 results."
-    ),
-    Tool(
-        name="Retriever",
-        func=create_retriever_tool,
-        description="A tool to retrieve similar questions from a vector store."
-    )
-]
-
-def create_agent(llm, tools):
-    """Create an agent with the specified tools."""
-    prompt = ChatPromptTemplate.from_messages([
-        system_message_prompt,
-        HumanMessagePromptTemplate.from_template("{input}")
-    ])
-    llm_chain = LLMChain(llm=llm, prompt=prompt)
-    return llm_chain
-
-def extract_final_answer(full_response):
-    """Extract only the final answer from the agent's response."""
-    if "FINAL ANSWER:" in full_response:
-        return full_response.split("FINAL ANSWER:")[1].strip()
-    return full_response.strip()
+# --- Core Agent Logic ---
+def get_agent_response(question: str) -> str:
+    """Get agent response with integrated search"""
+    try:
+        # Step 1: Search for relevant information
+        search_results = enhanced_search(question)
+        
+        # Step 2: Format context
+        context = []
+        if search_results.get("web"):
+            context.append("Web results:\n- " + "\n- ".join(search_results["web"]))
+        if search_results.get("wikipedia"):
+            context.append("Wikipedia results:\n- " + "\n- ".join(search_results["wikipedia"]))
+        
+        # Step 3: Retrieve similar questions
+        similar = vector_store.similarity_search(question, k=1)
+        if similar:
+            context.append(f"Similar question: {similar[0].page_content}")
+        
+        full_context = "\n\n".join(context) if context else "No search results found"
+        
+        # Step 4: Generate answer
+        chain = LLMChain(llm=llm, prompt=prompt)
+        response = chain.run({
+            "search_results": full_context,
+            "question": question
+        })
+        
+        return process_answer(response)
+    
+    except Exception as e:
+        print(f"Agent error: {e}")
+        return f"Error processing question: {e}"
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """