Update app.py

app.py

@@ -2,17 +2,23 @@ import gradio as gr
 import wikipedia
 from langchain_community.chat_models import ChatOpenAI
 from langchain.memory import ConversationBufferMemory
-from langchain.agents import initialize_agent
+from langchain.agents import initialize_agent, AgentType
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.tools import Tool
-from tavily import TavilyClient
+
+# Optional: Tavily client (only if installed)
+try:
+    from tavily import TavilyClient
+    tavily_available = True
+except ImportError:
+    TavilyClient = None
+    tavily_available = False
 
 
 # ----------------------
-# Define tools
+# Tools
 # ----------------------
 def create_your_own(query: str) -> str:
-    """This function can do whatever you would like once you fill it in"""
     return query[::-1]
 
 def get_current_temperature(query: str) -> str:
@@ -27,7 +33,6 @@ def search_wikipedia(query: str) -> str:
     except wikipedia.exceptions.PageError:
         return "No relevant Wikipedia page found."
 
-
 tools = [
     Tool(name="Temperature", func=get_current_temperature, description="Get current temperature"),
     Tool(name="Search Wikipedia", func=search_wikipedia, description="Search Wikipedia"),
@@ -36,56 +41,52 @@ tools = [
 
 
 # ----------------------
-# Define chatbot class
+# Chatbot class
 # ----------------------
-from langchain.agents import initialize_agent, AgentType
-
 class cbfs:
     def __init__(self, tools, openai_key: str, tavily_key: str = None):
         if not openai_key:
-            raise ValueError("⚠️ Please provide an OpenAI API key.")
+            raise ValueError("⚠️ OpenAI API key is required.")
 
         # Initialize OpenAI model
         self.model = ChatOpenAI(temperature=0, openai_api_key=openai_key)
 
-        # Initialize Tavily (optional)
-        self.tavily = TavilyClient(api_key=tavily_key) if tavily_key else None
+        # Initialize Tavily if available
+        self.tavily = None
+        if tavily_available and tavily_key:
+            self.tavily = TavilyClient(api_key=tavily_key)
 
         # Memory
         self.memory = ConversationBufferMemory(
             return_messages=True, memory_key="chat_history", ai_prefix="Assistant"
         )
 
-        # Agent
+        # Prompt
+        self.prompt = ChatPromptTemplate.from_messages([
+            ("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")
+        ])
+
+        # Initialize agent
         self.chain = initialize_agent(
             tools=tools,
             llm=self.model,
-            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # ✅ correct way
+            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
             verbose=True,
             memory=self.memory,
-            handle_parsing_errors=True  # ✅ prevents silent failure
+            handle_parsing_errors=True
         )
-    
+
     def convchain(self, query: str) -> str:
         if not query:
             return "Please enter a query."
         try:
             result = self.chain.invoke({"input": query})
-            # Debugging: show full raw result
-            print("Agent raw result:", result)
-
-            # Try both possible output keys
-            response = (
-                result.get("output")
-                or result.get("output_text")
-                or "⚠️ No response generated."
-            )
-
-            # Save memory
-            self.memory.save_context({"input": query}, {"output": response})
-            return response
+            print("Agent raw result:", result)  # Debug log in HuggingFace logs
+            return result.get("output") or result.get("output_text") or "⚠️ No response generated."
         except Exception as e:
-            print("Execution Error:", str(e))
             return f"❌ Error: {str(e)}"
 
 
@@ -94,31 +95,30 @@ class cbfs:
 # ----------------------
 with gr.Blocks() as demo:
     with gr.Row():
-        openai_key = gr.Textbox(label="🔑 OpenAI API Key", type="password", placeholder="Enter your OpenAI key")
-        tavily_key = gr.Textbox(label="🔑 Tavily API Key", type="password", placeholder="Enter your Tavily key")
+        openai_key = gr.Textbox(label="🔑 OpenAI API Key", type="password", placeholder="Paste your key")
+        tavily_key = gr.Textbox(label="🔑 Tavily API Key (optional)", type="password", placeholder="Paste your Tavily key")
 
-    chatbot_state = gr.State(None)  # will hold chatbot instance
+    chatbot_state = gr.State(None)
 
     with gr.Row():
         inp = gr.Textbox(placeholder="Enter text here…", label="User Input")
         output = gr.Textbox(placeholder="Response...", label="ChatBot Output", interactive=False)
 
-    # Initialize chatbot after keys are provided
+    status = gr.Textbox(label="Status", interactive=False)
+
     def init_chatbot(openai_key, tavily_key):
         try:
-            return cbfs(tools, openai_key, tavily_key), "✅ Chatbot initialized successfully!"
+            bot = cbfs(tools, openai_key, tavily_key)
+            return bot, "✅ Chatbot initialized successfully!"
         except Exception as e:
             return None, f"❌ Error: {str(e)}"
 
     init_btn = gr.Button("Initialize Chatbot")
-    status = gr.Textbox(label="Status", interactive=False)
-
     init_btn.click(fn=init_chatbot, inputs=[openai_key, tavily_key], outputs=[chatbot_state, status])
 
-    # Chat functionality
     def process_query(query, chatbot):
         if chatbot is None:
-            return "⚠️ Please initialize the chatbot first by providing your API keys."
+            return "⚠️ Please initialize the chatbot first by entering your API keys."
         return chatbot.convchain(query)
 
     inp.submit(process_query, inputs=[inp, chatbot_state], outputs=output)