Update app.py

app.py

@@ -8,28 +8,31 @@ 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
+from pydantic import BaseModel, 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
+mcp = FastMCP("Robot_MCP_Server")
+
 
 # ---------------------------------------------------
-#  Define Pydantic Schema for the input payload
+#  Payload Schema
 # ---------------------------------------------------
-# 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")
+    robot_id: str = Field(description="Robot identifier.", default="unknown")
     image_b64: str = Field(description="Base64 encoded image data.")
 
 
+# ---------------------------------------------------
+# Upload Helper
+# ---------------------------------------------------
 def upload_image(image_b64: str, hf_token: str):
     try:
         image_bytes = base64.b64decode(image_b64)
-
         os.makedirs("/tmp", exist_ok=True)
+
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
         local_path = f"/tmp/robot_img_{timestamp}.jpg"
 
@@ -38,6 +41,7 @@ def upload_image(image_b64: str, hf_token: str):
 
         filename = f"robot_{timestamp}.jpg"
         api = HfApi()
+
         api.upload_file(
             path_or_fileobj=local_path,
             path_in_repo=f"tmp/{filename}",
@@ -46,14 +50,17 @@ def upload_image(image_b64: str, hf_token: str):
             token=hf_token
         )
 
-        url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/tmp/{filename}"
-        return local_path, url, filename, len(image_bytes)
+        hf_url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/tmp/{filename}"
+        return local_path, hf_url, filename, len(image_bytes)
 
     except Exception:
         traceback.print_exc()
         return None, None, None, 0
 
 
+# ---------------------------------------------------
+# JSON Cleaning Helper
+# ---------------------------------------------------
 def safe_parse_json_from_text(text: str):
     if not text:
         return None
@@ -75,20 +82,14 @@ def safe_parse_json_from_text(text: str):
 
 
 # ---------------------------------------------------
-#  TRUE MCP TOOL — THIS must be exposed to MCP client
+# TRUE MCP TOOL
 # ---------------------------------------------------
 @mcp.tool("robot_watch", description="Analyze a base64 image using Qwen VLM and return structured JSON.")
-def robot_watch(payload: RobotWatchPayload): # <-- Type hint with Pydantic model
-    # The payload is already validated and typed correctly by fastmcp/pydantic
+def robot_watch_tool(payload: RobotWatchPayload):
     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"}
-    # image_b64 existence is also guaranteed by the schema
-
     _, hf_url, _, size_bytes = upload_image(image_b64, hf_token)
     if not hf_url:
         return {"error": "Image upload failed"}
@@ -136,21 +137,19 @@ Respond in STRICT JSON ONLY:
 
 
 # ---------------------------------------------------
-# Gradio UI — Use a simple placeholder function for the UI
+# Gradio UI Placeholder
 # ---------------------------------------------------
-def robot_watch(payload):
-    # This is just for the interactive UI, the real API call goes to /robot_watch endpoint
-    return {"message": "Use the MCP Client to call the robot_watch tool."}
+def robot_watch_ui(payload):
+    return {"message": "Use an MCP Client to call the robot_watch tool."}
 
 
 app = gr.Interface(
-    fn=robot_watch,
+    fn=robot_watch_ui,
     inputs=gr.JSON(),
     outputs=gr.JSON(),
     title="Robot MCP Server",
-    description="A MCP Server to describe image obtained from the CV of a Robot/Webcam"
+    description="A MCP Server to describe image obtained from the CV of a robot/webcam."
 )
 
 if __name__ == "__main__":
-    # Launch Gradio, which automatically hooks up the 'mcp' instance's APIs
     app.launch(mcp_server=True)