Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import os
+import wikipedia
 from langchain_community.chat_models import ChatOpenAI
 from langchain.memory import ConversationBufferMemory
 from langchain.agents import AgentExecutor, initialize_agent
@@ -15,7 +16,13 @@ def get_current_temperature(query: str) -> str:
     return "It's sunny and 75°F."
 
 def search_wikipedia(query: str) -> str:
-    return "Wikipedia search results for: " + query
+    try:
+        summary = wikipedia.summary(query, sentences=2)
+        return summary
+    except wikipedia.exceptions.DisambiguationError as e:
+        return f"Multiple results found: {', '.join(e.options[:5])}"
+    except wikipedia.exceptions.PageError:
+        return "No relevant Wikipedia page found."
 
 tools = [
     Tool(name="Temperature", func=get_current_temperature, description="Get current temperature"),
@@ -27,9 +34,9 @@ tools = [
 class cbfs:
     def __init__(self, tools):
         self.model = ChatOpenAI(temperature=0, openai_api_key=os.getenv("OPENAI_API_KEY"))
-        self.memory = ConversationBufferMemory(return_messages=True, memory_key="chat_history")
+        self.memory = ConversationBufferMemory(return_messages=True, memory_key="chat_history", ai_prefix="Assistant")
         self.prompt = ChatPromptTemplate.from_messages([
-            ("system", "You are a helpful but sassy assistant"),
+            ("system", "You are a helpful but sassy assistant. Remember what the user tells you in the conversation."),
             MessagesPlaceholder(variable_name="chat_history"),
             ("user", "{input}"),
             MessagesPlaceholder(variable_name="agent_scratchpad")
@@ -46,10 +53,11 @@ class cbfs:
         if not query:
             return "Please enter a query."
         try:
-            result = self.chain.invoke({"input": query, "chat_history": self.memory.load_memory_variables({})["chat_history"]})
-            self.memory.save_context({"input": query}, {"output": result.get("output", "No response generated.")})
-            print("Agent Execution Result:", result)  # Debugging output
-            return result.get("output", "No response generated.")
+            result = self.chain.invoke({"input": query})
+            response = result.get("output", "No response generated.")
+            self.memory.save_context({"input": query}, {"output": response})
+            print("Agent Execution Result:", response)  # Debugging output
+            return response
         except Exception as e:
             print("Execution Error:", str(e))
             return f"Error: {str(e)}"
@@ -71,3 +79,4 @@ demo.launch(share=True)
 
 
 
+