add tool calling and user wise chat

app.py

@@ -1,67 +1,144 @@
 import gradio as gr
 import uuid
 from langchain_core.messages import HumanMessage
-# from chatbot import chatbot
-from memory_chatbot import chatbot
+from memory_chatbot import chatbot,get_last_ai_message,extract_tool_logs
 
 
 # -----------------------
-# Generate Thread ID
+# Generate IDs
 # -----------------------
 def new_thread_id():
-    return str(uuid.uuid4())[:8]
+    return "thread_" + str(uuid.uuid4())[:8]
+
+def new_user_id():
+    return "user_" + str(uuid.uuid4())[:6]
 
 
 # -----------------------
-# First Thread
+# Initial User + Thread
 # -----------------------
+first_user = new_user_id()
 first_thread = new_thread_id()
 
 
 # -----------------------
 # Chat Function
 # -----------------------
-def chat_fn(message, chat_history, thread_id, all_threads):
+def chat_fn(message, chat_history, thread_id, user_id, all_threads):
 
-    config = {"configurable": {"thread_id": thread_id, "user_id": "user_id"}}
+    config = {
+        "configurable": {
+            "thread_id": thread_id,
+            "user_id": user_id
+        }
+    }
 
     res = chatbot.invoke(
         {"messages": [HumanMessage(content=message)]},
         config=config
     )
 
-    ai_message = res["messages"][-1].content
+    ai_message = get_last_ai_message(res["messages"])
+    tool_logs = extract_tool_logs(res["messages"], message)
 
-    # Append in NEW format
+    chat_history = chat_history or []
     chat_history.append({"role": "user", "content": message})
     chat_history.append({"role": "assistant", "content": ai_message})
 
-    all_threads[thread_id] = chat_history
+    all_threads[(user_id, thread_id)] = {
+        "chat": chat_history,
+        "tools": tool_logs
+    }
 
-    return "", chat_history, all_threads
+    return "", chat_history, all_threads, tool_logs
 
 
 # -----------------------
 # Create New Thread
 # -----------------------
-def create_new_thread(all_threads):
-    new_id = new_thread_id()
-    all_threads[new_id] = []
+
+def create_new_thread(user_id, all_threads):
+
+    new_tid = new_thread_id()
+    all_threads[(user_id, new_tid)] = {
+        "chat": [],
+        "tools": "_No tools used yet._"
+    }
 
     return (
         all_threads,
-        gr.update(choices=list(all_threads.keys()), value=new_id),
-        new_id,
-        []
+        gr.update(choices=[t for (u, t) in all_threads if u == user_id], value=new_tid),
+        new_tid,
+        [],
+        "_No tools used yet._"
     )
 
 
 # -----------------------
 # Switch Thread
 # -----------------------
-def switch_thread(thread_id, all_threads):
-    history = all_threads.get(thread_id, [])
-    return thread_id, history
+
+def switch_thread(thread_id, user_id, all_threads):
+
+    data = all_threads.get(
+        (user_id, thread_id),
+        {"chat": [], "tools": "_No tools used yet._"}
+    )
+
+    return thread_id, data["chat"], data["tools"]
+
+# -----------------------
+# Create New User
+# -----------------------
+def create_new_user(all_threads):
+
+    new_uid = new_user_id()
+    new_tid = new_thread_id()
+
+    all_threads[(new_uid, new_tid)] = {
+        "chat": [],
+        "tools": "_No tools used yet._"
+    }
+
+    return (
+        all_threads,
+        gr.update(choices=list(set([u for (u, _) in all_threads])), value=new_uid),
+        new_uid,
+        gr.update(choices=[new_tid], value=new_tid),
+        new_tid,
+        [],
+        "_No tools used yet._"
+    )
+
+
+# -----------------------
+# Switch User
+# -----------------------
+def switch_user(user_id, all_threads):
+
+    user_threads = [t for (u, t) in all_threads if u == user_id]
+
+    if not user_threads:
+        new_tid = new_thread_id()
+        all_threads[(user_id, new_tid)] = {
+            "chat": [],
+            "tools": "_No tools used yet._"
+        }
+        user_threads = [new_tid]
+
+    first_thread = user_threads[0]
+    data = all_threads.get(
+        (user_id, first_thread),
+        {"chat": [], "tools": "_No tools used yet._"}
+    )
+
+    return (
+        user_id,
+        gr.update(choices=user_threads, value=first_thread),
+        first_thread,
+        data["chat"],
+        data["tools"]
+    )
 
 
 # -----------------------
@@ -69,13 +146,31 @@ def switch_thread(thread_id, all_threads):
 # -----------------------
 with gr.Blocks() as demo:
 
+    user_state = gr.State(first_user)
     thread_state = gr.State(first_thread)
-    threads_state = gr.State({first_thread: []})
+
+    threads_state = gr.State({
+    (first_user, first_thread): {
+        "chat": [],
+        "tools": "_No tools used yet._"
+        }
+    })
 
     with gr.Row():
 
         # Sidebar
         with gr.Column(scale=1):
+
+            gr.Markdown("## 👤 Users")
+
+            user_dropdown = gr.Dropdown(
+                choices=[first_user],
+                value=first_user,
+                label="Select User"
+            )
+
+            new_user_btn = gr.Button("➕ New User")
+
             gr.Markdown("## 🧵 Threads")
 
             thread_radio = gr.Radio(
@@ -86,36 +181,74 @@ with gr.Blocks() as demo:
 
             new_chat_btn = gr.Button("➕ New Chat")
 
-        # Main Chat
-        with gr.Column(scale=4):
+        # Chat Area
+        with gr.Column(scale=7):
 
             chatbot_ui = gr.Chatbot()
             msg_box = gr.Textbox(placeholder="Type message and press Enter")
+            tool_panel = gr.Markdown("## 🛠 Tool Activity\n_No tools used yet._")
 
             msg_box.submit(
                 chat_fn,
-                inputs=[msg_box, chatbot_ui, thread_state, threads_state],
-                outputs=[msg_box, chatbot_ui, threads_state],
+                inputs=[msg_box, chatbot_ui, thread_state, user_state, threads_state],
+                outputs=[msg_box, chatbot_ui, threads_state, tool_panel],
             )
 
-    # New Chat Button
+    # -----------------------
+    # New User
+    # -----------------------
+    new_user_btn.click(
+    create_new_user,
+    inputs=threads_state,
+    outputs=[
+        threads_state,
+        user_dropdown,
+        user_state,
+        thread_radio,
+        thread_state,
+        chatbot_ui,
+        tool_panel
+    ]
+)
+
+    # -----------------------
+    # Switch User
+    # -----------------------
+    user_dropdown.change(
+    switch_user,
+    inputs=[user_dropdown, threads_state],
+    outputs=[
+        user_state,
+        thread_radio,
+        thread_state,
+        chatbot_ui,
+        tool_panel
+    ]
+)
+
+    # -----------------------
+    # New Thread
+    # -----------------------
     new_chat_btn.click(
-        create_new_thread,
-        inputs=threads_state,
-        outputs=[
-            threads_state,
-            thread_radio,
-            thread_state,
-            chatbot_ui
-        ]
-    )
-
+    create_new_thread,
+    inputs=[user_state, threads_state],
+    outputs=[
+        threads_state,
+        thread_radio,
+        thread_state,
+        chatbot_ui,
+        tool_panel
+    ]
+)
+
+    # -----------------------
     # Switch Thread
+    # -----------------------
     thread_radio.change(
-        switch_thread,
-        inputs=[thread_radio, threads_state],
-        outputs=[thread_state, chatbot_ui]
-    )
+    switch_thread,
+    inputs=[thread_radio, user_state, threads_state],
+    outputs=[thread_state, chatbot_ui, tool_panel]
+)
 
 
 demo.launch()

app_backup.py

@@ -0,0 +1,321 @@
+import gradio as gr
+import uuid
+from langchain_core.messages import HumanMessage
+from memory_chatbot import chatbot
+
+
+# -----------------------
+# Generate IDs
+# -----------------------
+def new_thread_id():
+    return "thread_" + str(uuid.uuid4())[:8]
+
+def new_user_id():
+    return "user_" + str(uuid.uuid4())[:6]
+
+
+# -----------------------
+# Initial User + Thread
+# -----------------------
+first_user = new_user_id()
+first_thread = new_thread_id()
+
+
+# -----------------------
+# Chat Function
+# -----------------------
+def chat_fn(message, chat_history, thread_id, user_id, all_threads):
+
+    # tool_logs = []
+
+    config = {
+        "configurable": {
+            "thread_id": thread_id,
+            "user_id": user_id
+        }
+    }
+
+    res = chatbot.invoke(
+        {"messages": [HumanMessage(content=message)]},
+        config=config
+    )
+
+    ai_message = res["messages"][-1].content
+
+
+    chat_history.append({"role": "user", "content": message})
+    chat_history.append({"role": "assistant", "content": ai_message})
+
+    all_threads[(user_id, thread_id)] = chat_history
+
+    return "", chat_history, all_threads
+
+
+# -----------------------
+# Create New Thread
+# -----------------------
+def create_new_thread(user_id, all_threads):
+
+    new_tid = new_thread_id()
+    all_threads[(user_id, new_tid)] = []
+
+    return (
+        all_threads,
+        gr.update(choices=[t for (u, t) in all_threads if u == user_id], value=new_tid),
+        new_tid,
+        []
+    )
+
+
+# -----------------------
+# Switch Thread
+# -----------------------
+def switch_thread(thread_id, user_id, all_threads):
+
+    history = all_threads.get((user_id, thread_id), [])
+    return thread_id, history
+
+
+# -----------------------
+# Create New User
+# -----------------------
+def create_new_user(all_threads):
+
+    new_uid = new_user_id()
+    new_tid = new_thread_id()
+
+    all_threads[(new_uid, new_tid)] = []
+
+    return (
+        all_threads,
+        gr.update(choices=list(set([u for (u, _) in all_threads])), value=new_uid),
+        new_uid,
+        gr.update(choices=[new_tid], value=new_tid),
+        new_tid,
+        []
+    )
+
+
+# -----------------------
+# Switch User
+# -----------------------
+def switch_user(user_id, all_threads):
+
+    user_threads = [t for (u, t) in all_threads if u == user_id]
+
+    if not user_threads:
+        new_tid = new_thread_id()
+        all_threads[(user_id, new_tid)] = []
+        user_threads = [new_tid]
+
+    first_thread = user_threads[0]
+    history = all_threads.get((user_id, first_thread), [])
+
+    return (
+        user_id,
+        gr.update(choices=user_threads, value=first_thread),
+        first_thread,
+        history
+    )
+
+
+# -----------------------
+# UI
+# -----------------------
+with gr.Blocks() as demo:
+
+    user_state = gr.State(first_user)
+    thread_state = gr.State(first_thread)
+
+    threads_state = gr.State({
+        (first_user, first_thread): []
+    })
+
+    with gr.Row():
+
+        # Sidebar
+        with gr.Column(scale=1):
+
+            gr.Markdown("## 👤 Users")
+
+            user_dropdown = gr.Dropdown(
+                choices=[first_user],
+                value=first_user,
+                label="Select User"
+            )
+
+            new_user_btn = gr.Button("➕ New User")
+
+            gr.Markdown("## 🧵 Threads")
+
+            thread_radio = gr.Radio(
+                choices=[first_thread],
+                value=first_thread,
+                label="Select Thread"
+            )
+
+            new_chat_btn = gr.Button("➕ New Chat")
+
+        # Chat Area
+        with gr.Column(scale=4):
+
+            chatbot_ui = gr.Chatbot()
+            msg_box = gr.Textbox(placeholder="Type message and press Enter")
+
+            msg_box.submit(
+                chat_fn,
+                inputs=[msg_box, chatbot_ui, thread_state, user_state, threads_state],
+                outputs=[msg_box, chatbot_ui, threads_state],
+            )
+
+    # -----------------------
+    # New User
+    # -----------------------
+    new_user_btn.click(
+        create_new_user,
+        inputs=threads_state,
+        outputs=[
+            threads_state,
+            user_dropdown,
+            user_state,
+            thread_radio,
+            thread_state,
+            chatbot_ui
+        ]
+    )
+
+    # -----------------------
+    # Switch User
+    # -----------------------
+    user_dropdown.change(
+        switch_user,
+        inputs=[user_dropdown, threads_state],
+        outputs=[
+            user_state,
+            thread_radio,
+            thread_state,
+            chatbot_ui
+        ]
+    )
+
+    # -----------------------
+    # New Thread
+    # -----------------------
+    new_chat_btn.click(
+        create_new_thread,
+        inputs=[user_state, threads_state],
+        outputs=[
+            threads_state,
+            thread_radio,
+            thread_state,
+            chatbot_ui
+        ]
+    )
+
+    # -----------------------
+    # Switch Thread
+    # -----------------------
+    thread_radio.change(
+        switch_thread,
+        inputs=[thread_radio, user_state, threads_state],
+        outputs=[thread_state, chatbot_ui]
+    )
+
+
+demo.launch()
+
+
+
+
+
+
+
+
+
+# import gradio as gr
+# import uuid
+# from langchain_core.messages import HumanMessage
+# from memory_chatbot import chatbot
+
+
+# # -----------------------
+# # Generate IDs
+# # -----------------------
+# def new_thread_id():
+#     return "thread_" + str(uuid.uuid4())[:8]
+
+# def new_user_id():
+#     return "user_" + str(uuid.uuid4())[:6]
+
+
+# # -----------------------
+# # Chat Function
+# # -----------------------
+# def chat_fn(message, chat_history, thread_id, user_id):
+
+#     config = {
+#         "configurable": {
+#             "thread_id": thread_id,
+#             "user_id": user_id
+#         }
+#     }
+
+#     res = chatbot.invoke(
+#         {"messages": [HumanMessage(content=message)]},
+#         config=config
+#     )
+
+#     ai_message = res["messages"][-1].content
+
+#     chat_history.append({"role": "user", "content": message})
+#     chat_history.append({"role": "assistant", "content": ai_message})
+
+#     return "", chat_history
+
+
+# # -----------------------
+# # New Chat (Reset Thread)
+# # -----------------------
+# def new_chat():
+#     return new_thread_id(), []
+
+
+# # -----------------------
+# # UI
+# # -----------------------
+# with gr.Blocks() as demo:
+
+#     # Hidden states
+#     user_state = gr.State(new_user_id())
+#     thread_state = gr.State(new_thread_id())
+
+#     gr.Markdown("# 🤖 AI Chatbot")
+
+#     chatbot_ui = gr.Chatbot(height=500)
+#     msg_box = gr.Textbox(placeholder="Type your message...")
+
+#     with gr.Row():
+#         send_btn = gr.Button("Send")
+#         new_chat_btn = gr.Button("🧹 New Chat")
+
+#     # Send message
+#     send_btn.click(
+#         chat_fn,
+#         inputs=[msg_box, chatbot_ui, thread_state, user_state],
+#         outputs=[msg_box, chatbot_ui],
+#     )
+
+#     msg_box.submit(
+#         chat_fn,
+#         inputs=[msg_box, chatbot_ui, thread_state, user_state],
+#         outputs=[msg_box, chatbot_ui],
+#     )
+
+#     # New chat (reset thread)
+#     new_chat_btn.click(
+#         new_chat,
+#         outputs=[thread_state, chatbot_ui]
+#     )
+
+
+# demo.launch()

memory_chatbot.py

@@ -2,7 +2,7 @@ import json
 from typing import TypedDict, Annotated
 
 from dotenv import load_dotenv
-from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.messages import BaseMessage, HumanMessage,ToolMessage,AIMessage
 from langchain_groq import ChatGroq
 
 from langgraph.graph import START, END, StateGraph
@@ -44,6 +44,59 @@ llm = ChatGroq(
 # State
 # -----------------------------
 
+def get_last_ai_message(messages):
+    """Get the final AI response safely"""
+    for msg in reversed(messages):
+        if isinstance(msg, AIMessage) and msg.content:
+            return msg.content if isinstance(msg.content, str) else str(msg.content)
+    return "No response"
+
+
+def extract_tool_logs(messages, current_user_message):
+    """
+    Extract tool calls and tool outputs from the latest turn only
+    """
+    latest_turn = []
+    found_user_message = False
+
+    # Walk backwards until we find the current user message
+    for msg in reversed(messages):
+        latest_turn.append(msg)
+        if isinstance(msg, HumanMessage) and msg.content == current_user_message:
+            found_user_message = True
+            break
+
+    if found_user_message:
+        latest_turn = list(reversed(latest_turn))
+    else:
+        latest_turn = messages
+
+    logs = []
+
+    for msg in latest_turn:
+        # AI asked to call tool
+        if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
+            for call in msg.tool_calls:
+                tool_name = call.get("name", "unknown_tool")
+                tool_args = call.get("args", {})
+
+                logs.append(
+                    f"### 🔧 Calling `{tool_name}`\n"
+                    f"```json\n{json.dumps(tool_args, indent=2)}\n```"
+                )
+
+        # Tool returned output
+        elif isinstance(msg, ToolMessage):
+            tool_name = getattr(msg, "name", "unknown_tool")
+            tool_output = msg.content if isinstance(msg.content, str) else str(msg.content)
+
+            logs.append(
+                f"### ✅ Output from `{tool_name}`\n"
+                f"```\n{tool_output[:1000]}\n```"
+            )
+
+    return "\n\n".join(logs) if logs else "_No tools used in this response._"
+
 class ChatState(TypedDict):
     messages: Annotated[list[BaseMessage], add_messages]
 
@@ -147,6 +200,38 @@ def memory_retrieval(state: ChatState, config, store):
 
     return {"messages": state["messages"]}
 
+
+# def memory_retrieval(state: ChatState, config, store):
+
+#     user_id = config["configurable"]["user_id"]
+#     query = state["messages"][-1].content
+
+#     categories = [
+#         "preferences",
+#         "profile",
+#         "interests",
+#         "project"
+#     ]
+
+#     all_memories = []
+
+#     for category in categories:
+#         namespace = (user_id, category)
+#         results = store.search(namespace, query=query, limit=3)
+#         all_memories.extend(results)
+
+#     memory_text = "\n".join([f"{m.namespace[1]}: {m.value}" for m in all_memories])
+
+#     # FIX: Only return the new message(s) added in this step, not the whole state.
+#     if memory_text:
+#         return {
+#             "messages": [
+#                 HumanMessage(content=f"Relevant user memories:\n{memory_text}")
+#             ]
+#         }
+    
+#     return {}
+
 # -----------------------------
 # Chat Node
 # -----------------------------