Update app.py

app.py

@@ -8,12 +8,22 @@ from typing import Dict, Any
 import gradio as gr
 from huggingface_hub import HfApi, InferenceClient
 from fastmcp import FastMCP
+from pydantic import BaseModel, Field # Import Pydantic BaseModel and Field
 
 HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "OppaAI/Robot_MCP")
 HF_VLM_MODEL = os.environ.get("HF_VLM_MODEL", "Qwen/Qwen2.5-VL-7B-Instruct")
 
 mcp = FastMCP("Robot_MCP_Server")   # <-- Important
 
+# ---------------------------------------------------
+#  Define Pydantic Schema for the input payload
+# ---------------------------------------------------
+# This defines the expected structure and automatically generates the valid JSON schema
+class RobotWatchPayload(BaseModel):
+    hf_token: str = Field(description="Your Hugging Face API token.")
+    robot_id: str = Field(description="The unique identifier for the robot.", default="unknown")
+    image_b64: str = Field(description="Base64 encoded image data.")
+
 
 def upload_image(image_b64: str, hf_token: str):
     try:
@@ -69,17 +79,19 @@ def safe_parse_json_from_text(text: str):
 # ---------------------------------------------------
 @mcp.tool(
     name="robot_watch",
-    description="Analyze a base64 image using Qwen VLM and return structured JSON."
+    description="Analyze a base64 image using Qwen VLM and return structured JSON.",
+    input_schema=RobotWatchPayload # <-- Explicitly use the Pydantic schema here
 )
-def robot_watch(payload: Dict[str, Any]):
-    hf_token = payload.get("hf_token")
-    image_b64 = payload.get("image_b64")
-    robot_id = payload.get("robot_id", "unknown")
+def robot_watch(payload: RobotWatchPayload): # <-- Type hint with Pydantic model
+    # The payload is already validated and typed correctly by fastmcp/pydantic
+    hf_token = payload.hf_token
+    image_b64 = payload.image_b64
+    robot_id = payload.robot_id
 
     if not hf_token:
+        # This check is technically redundant if the schema demands it, but safe.
         return {"error": "Missing hf_token"}
-    if not image_b64:
-        return {"error": "Missing image_b64"}
+    # image_b64 existence is also guaranteed by the schema
 
     _, hf_url, _, size_bytes = upload_image(image_b64, hf_token)
     if not hf_url:
@@ -130,8 +142,11 @@ Respond in STRICT JSON ONLY:
 # ---------------------------------------------------
 # Gradio UI — separate from MCP tool layer
 # ---------------------------------------------------
-def process_json(payload):
-    return robot_watch(payload)
+# The process_json function will still work with the Pydantic model input
+def process_json(payload: Dict[str, Any]):
+    # When called via Gradio UI (not MCP), input will be Dict, so handle type conversion
+    pydantic_payload = RobotWatchPayload(**payload)
+    return robot_watch(pydantic_payload)
 
 
 app = gr.Interface(
@@ -143,4 +158,5 @@ app = gr.Interface(
 )
 
 if __name__ == "__main__":
+    # Gradio will use the 'mcp' instance defined globally to host the MCP server endpoints
     app.launch(mcp_server=True)