Update app.py

app.py

@@ -1,13 +1,16 @@
 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
+from langchain.agents import initialize_agent
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.tools import Tool
+from tavily import TavilyClient
 
+
+# ----------------------
 # Define tools
+# ----------------------
 def create_your_own(query: str) -> str:
     """This function can do whatever you would like once you fill it in"""
     return query[::-1]
@@ -24,16 +27,29 @@ 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"),
     Tool(name="Create Your Own", func=create_your_own, description="Custom tool for processing input")
 ]
 
+
+# ----------------------
 # Define chatbot class
+# ----------------------
 class cbfs:
-    def __init__(self, tools):
-        self.model = ChatOpenAI(temperature=0, openai_api_key=os.getenv("OPENAI_API_KEY"))
+    def __init__(self, tools, openai_key: str, tavily_key: str):
+        if not openai_key or not tavily_key:
+            raise ValueError("⚠️ Please provide both OpenAI and Tavily API keys.")
+
+        # Initialize OpenAI model with user key
+        self.model = ChatOpenAI(temperature=0, openai_api_key=openai_key)
+
+        # Initialize Tavily client (for future tools or expansion)
+        self.tavily = TavilyClient(api_key=tavily_key)
+
+        # Memory + prompt
         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. Remember what the user tells you in the conversation."),
@@ -41,6 +57,8 @@ class cbfs:
             ("user", "{input}"),
             MessagesPlaceholder(variable_name="agent_scratchpad")
         ])
+
+        # Initialize agent
         self.chain = initialize_agent(
             tools=tools,
             llm=self.model,
@@ -49,34 +67,55 @@ class cbfs:
             memory=self.memory
         )
     
-    def convchain(self, query):
+    def convchain(self, query: str) -> str:
+        """Run a single query through the agent."""
         if not query:
             return "Please enter a query."
         try:
             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)}"
-
-# Create chatbot instance
-cb = cbfs(tools)
+            return f"❌ Error: {str(e)}"
 
-def process_query(query):
-    return cb.convchain(query)
 
-# Define Gradio UI
+# ----------------------
+# Gradio UI
+# ----------------------
 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")
+
+    chatbot_state = gr.State(None)  # will hold chatbot instance
+
     with gr.Row():
         inp = gr.Textbox(placeholder="Enter text here…", label="User Input")
         output = gr.Textbox(placeholder="Response...", label="ChatBot Output", interactive=False)
-    inp.submit(process_query, inputs=inp, outputs=output)
 
-demo.launch(share=True)
+    # Initialize chatbot after keys are provided
+    def init_chatbot(openai_key, tavily_key):
+        try:
+            return cbfs(tools, openai_key, tavily_key), "✅ 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 chatbot.convchain(query)
+
+    inp.submit(process_query, inputs=[inp, chatbot_state], outputs=output)
+
+
+demo.launch(share=True)