Update app.py

app.py

@@ -1,8 +1,11 @@
 # app.py
 # Hugging Face Space: Gradio app that chats about Gradio docs via the Gradio Docs MCP server.
-# - Uses huggingface_hub.MCPClient as an MCP client
-# - Streams responses while the model calls MCP tools under the hood
-# - Make sure to set HUGGING_FACE_HUB_TOKEN (or HF_TOKEN) in Space secrets.
+# Requirements:
+#   - gradio
+#   - huggingface_hub
+#
+# Space secret needed:
+#   - HUGGING_FACE_HUB_TOKEN (or HF_TOKEN)
 
 import os
 import asyncio
@@ -70,23 +73,25 @@ async def ensure_init():
 
 
 # ----------------------------
-# Chat plumbing
+# Helpers for messages
 # ----------------------------
-def to_messages(history: List[List[str]], user_msg: str) -> List[Dict[str, Any]]:
+def to_llm_messages(history_msgs: List[Dict[str, Any]], user_msg: str) -> List[Dict[str, Any]]:
     """
-    Convert gr.Chatbot history into an LLM-style message list.
+    Convert Chatbot messages list (role/content dicts) to the LLM format,
+    with a system message prepended and the new user message appended.
     """
-    messages: List[Dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for u, a in history:
-        if u:
-            messages.append({"role": "user", "content": str(u)})
-        if a:
-            messages.append({"role": "assistant", "content": str(a)})
-    messages.append({"role": "user", "content": user_msg})
-    return messages
-
-
-async def stream_answer(messages: List[Dict[str, Any]]) -> Iterable[str]:
+    msgs: List[Dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
+    # Keep only role/content keys
+    for m in history_msgs:
+        role = m.get("role")
+        content = m.get("content")
+        if role in ("user", "assistant") and isinstance(content, str):
+            msgs.append({"role": role, "content": content})
+    msgs.append({"role": "user", "content": user_msg})
+    return msgs
+
+
+async def stream_answer(messages_for_llm: List[Dict[str, Any]]) -> Iterable[str]:
     """
     Stream deltas and tool logs from MCPClient.process_single_turn_with_tools.
     """
@@ -102,38 +107,26 @@ async def stream_answer(messages: List[Dict[str, Any]]) -> Iterable[str]:
         return
 
     try:
-        partial = ""
-        async for chunk in client.process_single_turn_with_tools(messages):
+        async for chunk in client.process_single_turn_with_tools(messages_for_llm):
             # chunk is a dict describing text deltas and tool activity
             if isinstance(chunk, dict):
                 ctype = chunk.get("type")
                 if ctype == "tool_log":
                     name = chunk.get("tool", "tool")
                     status = chunk.get("status", "")
-                    text = f"\n\n_(using **{name}** {status})_"
-                    partial += text
-                    yield text
+                    yield f"\n\n_(using **{name}** {status})_"
                 elif ctype == "text_delta":
-                    delta = chunk.get("delta", "")
-                    partial += delta
-                    yield delta
+                    yield chunk.get("delta", "")
                 elif ctype == "text":
-                    text = chunk.get("text", "")
-                    partial += text
-                    yield text
+                    yield chunk.get("text", "")
                 elif ctype == "tool_result":
                     content = chunk.get("content")
                     if isinstance(content, str) and content.strip():
-                        text = f"\n\n**Result:**\n{content}"
-                        partial += text
-                        yield text
+                        yield f"\n\n**Result:**\n{content}"
             else:
                 # Fallback if provider yields plain strings
-                s = str(chunk)
-                partial += s
-                yield s
+                yield str(chunk)
     except Exception as e:
-        # Surface common errors nicely
         msg = str(e)
         if "401" in msg or "Unauthorized" in msg:
             yield (
@@ -146,12 +139,6 @@ async def stream_answer(messages: List[Dict[str, Any]]) -> Iterable[str]:
             yield f"❌ Error: {msg}"
 
 
-async def respond(user_msg: str, history: List[List[str]]):
-    messages = to_messages(history, user_msg)
-    async for piece in stream_answer(messages):
-        yield piece
-
-
 # ----------------------------
 # Gradio UI
 # ----------------------------
@@ -161,11 +148,11 @@ with gr.Blocks(fill_height=True) as demo:
         "Ask anything about **Gradio**. Answers are grounded in the official docs via MCP."
     )
 
+    # Using the new messages format
     chat = gr.Chatbot(
         label="Gradio Docs Assistant",
         height=520,
-        type="messages",  # keeps history as list[[user, assistant], ...]
-        avatar_images=(None, None)
+        type="messages",   # expects a list of dicts: {"role": "...", "content": "..."}
     )
 
     with gr.Row():
@@ -179,18 +166,26 @@ with gr.Blocks(fill_height=True) as demo:
     with gr.Row():
         clear = gr.ClearButton(components=[chat, msg], value="Clear")
         info = gr.Markdown(
-            f"**Model:** `{MODEL_ID}` · **Provider:** `{PROVIDER}` · "
-            f"**MCP:** Gradio Docs SSE",
-            elem_classes=["text-sm", "opacity-70"],
+            f"**Model:** `{MODEL_ID}` · **Provider:** `{PROVIDER}` · **MCP:** Gradio Docs SSE",
         )
 
-    async def on_submit(user_msg: str, history: List[List[str]]):
-        # Append a placeholder response to stream into
-        history = (history or []) + [[user_msg, ""]]
-        # Stream chunks into the last assistant message
-        async for chunk in respond(user_msg, history[:-1]):
-            history[-1][1] = chunk
-            yield history
+    async def on_submit(user_msg: str, history_msgs: List[Dict[str, Any]]):
+        """
+        history_msgs is a list of {"role": ..., "content": ...} dicts.
+        We append the user's message, then stream the assistant reply by
+        updating/overwriting the last assistant message content.
+        """
+        history_msgs = (history_msgs or []) + [{"role": "user", "content": user_msg}]
+        # Add a placeholder assistant message to stream into
+        history_msgs.append({"role": "assistant", "content": ""})
+        yield history_msgs
+
+        # Build LLM messages and stream chunks
+        messages_for_llm = to_llm_messages(history_msgs[:-1], user_msg)
+        async for delta in stream_answer(messages_for_llm):
+            # Append the delta to the last assistant message
+            history_msgs[-1]["content"] += delta
+            yield history_msgs
 
     # Wire both Enter and button click
     msg.submit(on_submit, inputs=[msg, chat], outputs=chat, queue=True)