Update app.py

app.py

@@ -1,44 +1,59 @@
+# app.py
+
 import os
 import streamlit as st
 from dotenv import load_dotenv
-from langchain_core.messages import AIMessage, HumanMessage
-from langchain_community.tools.tavily.tool import TavilySearchResults
-
-# from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
+from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
+from langchain_community.tools.tavily.search import TavilySearchResults  # ✅ Corrected import
 from langchain_groq import ChatGroq 
 from langgraph.graph import StateGraph, START, END
 from langgraph.prebuilt import ToolNode, tools_condition
 from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.messages import AIMessage, HumanMessage
 from langgraph.graph.message import add_messages
 from typing import Annotated
 from typing_extensions import TypedDict
-from langchain_tavily.tool import TavilySearchResults  # ✅ Corrected import
 from tools import get_all_tools
 
-# ---------------------- Environment ----------------------
+# Load environment
 load_dotenv()
 os.environ["TAVILY_API_KEY"] = os.getenv("TAVILY_API_KEY")
 os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
 
-# ---------------------- Tool Map ----------------------
-tool_map = {
-    "Execute Code": "execute_code",
-    "Explain Code": "explain_code",
-    "Web Search": "web_search",
-    "Deep Think": "deep_think",
-    "Analyze Code": "analyze_code",
-    "Code Generator": "code_generator"
-}
+tools = get_all_tools()
 
-# ---------------------- LangGraph State ----------------------
+# LangGraph State
 class State(TypedDict):
     messages: Annotated[list, add_messages]
     name: str
     birthday: str
 
-# ---------------------- Streamlit UI ----------------------
+# LLM
+llm = ChatGroq(model="deepseek-r1-distill-llama-70b")
+llm_with_tools = llm.bind_tools(tools=tools)
+
+# Memory
+memory = MemorySaver()
+
+# Define node
+def ai_assistance(state: State):
+    result = llm_with_tools.invoke(state["messages"])
+    return {"messages": [result]}
+
+# Build LangGraph
+builder = StateGraph(State)
+builder.add_node("AI_Assistance", ai_assistance)
+builder.add_node("tools", ToolNode(tools))
+
+builder.add_edge(START, "AI_Assistance")
+builder.add_conditional_edges("AI_Assistance", tools_condition)
+builder.add_edge("tools", "AI_Assistance")
+graph = builder.compile(checkpointer=memory)
+
+# Streamlit UI setup
 st.set_page_config(page_title="LangGraph Chatbot", layout="wide")
 
+# --- Custom Chat UI ---
 st.markdown("""
 <style>
 .stChatMessage {
@@ -47,24 +62,42 @@ st.markdown("""
     border-radius: 12px;
     max-width: 90%;
 }
-.user { background-color: #dcf8c6; align-self: flex-end; }
-.bot { background-color: #f1f0f0; align-self: flex-start; }
-.input-box { display: flex; align-items: center; gap: 0.5rem; }
+.user {
+    background-color: #dcf8c6;
+    align-self: flex-end;
+}
+.bot {
+    background-color: #f1f0f0;
+    align-self: flex-start;
+}
+.input-box {
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+}
 #floating-container {
-    display: flex; align-items: center; justify-content: space-between;
-    padding: 0.25rem 0.75rem; background-color: #f9f9f9;
-    border-radius: 0.75rem; margin-top: 1rem; border: 1px solid #ccc;
+    display: flex;
+    align-items: center;
+    justify-content: space-between;
+    padding: 0.25rem 0.75rem;
+    background-color: #f9f9f9;
+    border-radius: 0.75rem;
+    margin-top: 1rem;
+    border: 1px solid #ccc;
 }
 .floating-popup {
-    margin-top: 0.5rem; padding: 0.5rem;
-    border-radius: 0.5rem; border: 1px solid #ccc; background-color: white;
+    margin-top: 0.5rem;
+    padding: 0.5rem;
+    border-radius: 0.5rem;
+    border: 1px solid #ccc;
+    background-color: white;
 }
 </style>
 """, unsafe_allow_html=True)
 
 st.title("🧠 LangGraph Chatbot (Groq + Tools)")
 
-# ---------------------- Session Init ----------------------
+# Session State Init
 if "thread_id" not in st.session_state:
     st.session_state.thread_id = "1"
 if "chat_history" not in st.session_state:
@@ -72,13 +105,14 @@ if "chat_history" not in st.session_state:
 if "show_upload" not in st.session_state:
     st.session_state.show_upload = False
 if "selected_tools" not in st.session_state:
-    st.session_state.selected_tools = ["Web Search", "Explain Code", "Analyze Code"]
+    st.session_state.selected_tools = ["Arxiv", "Wikipedia", "Tavily"]
 
-# ---------------------- Upload/Tool UI ----------------------
+# Chat History
 for msg in st.session_state.chat_history:
     role = "user" if isinstance(msg, HumanMessage) else "bot"
     st.markdown(f"<div class='stChatMessage {role}'>{msg.content}</div>", unsafe_allow_html=True)
 
+# Chat Input Box
 with st.container():
     with st.form("chat_form", clear_on_submit=True):
         st.markdown('<div id="floating-container">', unsafe_allow_html=True)
@@ -91,6 +125,7 @@ with st.container():
 
         st.markdown('</div>', unsafe_allow_html=True)
 
+        # Tool Selector & Upload Panel
         if toggle:
             st.session_state.show_upload = not st.session_state.show_upload
 
@@ -101,41 +136,15 @@ with st.container():
 
             st.session_state.selected_tools = st.multiselect(
                 "🛠 Select Tools:",
-                list(tool_map.keys()),
+                ["Arxiv", "Wikipedia", "Tavily"],
                 default=st.session_state.selected_tools
             )
 
+        # Submit Action
         submitted = st.form_submit_button(label="Send")
-
         if submitted and user_input:
             st.session_state.chat_history.append(HumanMessage(content=user_input))
-
-            # ---------------------- Get Selected Tools ----------------------
-            all_tools = get_all_tools()
-            selected_tool_names = [tool_map[name] for name in st.session_state.selected_tools]
-            tools = [tool for tool in all_tools if tool.name in selected_tool_names]
-
-            # ---------------------- LLM & Graph ----------------------
-            llm = ChatGroq(model="deepseek-r1-distill-llama-70b")
-            llm_with_tools = llm.bind_tools(tools=tools)
-
-            def ai_assistance(state: State):
-                result = llm_with_tools.invoke(state["messages"])
-                return {"messages": [result]}
-
-            builder = StateGraph(State)
-            builder.add_node("AI_Assistance", ai_assistance)
-            builder.add_node("tools", ToolNode(tools))
-
-            builder.add_edge(START, "AI_Assistance")
-            builder.add_conditional_edges("AI_Assistance", tools_condition)
-            builder.add_edge("tools", "AI_Assistance")
-
-            memory = MemorySaver()
-            graph = builder.compile(checkpointer=memory)
-
             config = {"configurable": {"thread_id": st.session_state.thread_id}}
             result = graph.invoke({"messages": st.session_state.chat_history}, config=config)
-
             st.session_state.chat_history.append(result["messages"][-1])
             st.rerun()