Update app.py

app.py

@@ -6,14 +6,17 @@ import gradio as gr
 from huggingface_hub import upload_file, InferenceClient
 from datetime import datetime
 import traceback
-import threading
 from typing import Optional, Dict, Any, Tuple
 
 from fastmcp import FastMCP
 
-
-HF_DATASET_REPO = "OppaAI/Robot_MCP"
-HF_VLM_MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
+# --- Configuration using Environment Variables ---
+# It is best practice to manage sensitive info outside of the code.
+# Use os.environ.get() to safely retrieve these values.
+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")
+# The token will be required in the payload, but we define the env var name here.
+# HF_TOKEN_ENV_VAR_NAME = "HF_TOKEN" 
 
 mcp = FastMCP("Robot_MCP")
 
@@ -22,6 +25,7 @@ mcp = FastMCP("Robot_MCP")
 # -----------------------------------------------------
 @mcp.tool()
 def speak(text: str, emotion: str = "neutral"):
+    """Makes the robot speak a given text with an emotion."""
     return {
         "status": "success",
         "action_executed": "speak",
@@ -31,6 +35,7 @@ def speak(text: str, emotion: str = "neutral"):
 
 @mcp.tool()
 def navigate(direction: str, distance_meters: float):
+    """Moves the robot a specified distance in a direction (max 5m)."""
     if distance_meters > 5.0:
         return {"status": "error", "message": "Safety limit exceeded"}
     return {
@@ -42,6 +47,7 @@ def navigate(direction: str, distance_meters: float):
 
 @mcp.tool()
 def scan_hazard(hazard_type: str, severity: str):
+    """Logs a potential hazard detected by the robot."""
     timestamp = datetime.now().isoformat()
     return {
         "status": "warning_logged",
@@ -51,6 +57,7 @@ def scan_hazard(hazard_type: str, severity: str):
 
 @mcp.tool()
 def analyze_human(clothing_color: str, estimated_action: str):
+    """Tracks human activity based on visual input."""
     return {
         "status": "human_tracked",
         "details": f"Human wearing {clothing_color} is {estimated_action}",
@@ -60,10 +67,13 @@ def analyze_human(clothing_color: str, estimated_action: str):
 # Save + Upload
 # -----------------------------------------------------
 def save_and_upload_image(image_b64: str, hf_token: str):
+    """Decodes a base64 image, saves it locally, and uploads to Hugging Face Hub."""
     try:
         image_bytes = base64.b64decode(image_b64)
         size_bytes = len(image_bytes)
 
+        # Ensure the /tmp directory exists
+        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"
 
@@ -83,44 +93,47 @@ def save_and_upload_image(image_b64: str, hf_token: str):
         url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/{filename}"
         return local_path, url, filename, size_bytes
 
-    except Exception:
+    except Exception as e:
+        print(f"Error during image upload: {e}")
         traceback.print_exc()
         return None, None, None, 0
 
 # -----------------------------------------------------
 # JSON Parse
 # -----------------------------------------------------
-def safe_parse_json_from_text(text: str):
+def safe_parse_json_from_text(text: str) -> Optional[Dict[str, Any]]:
+    """Attempts to safely parse JSON from potentially messy text output."""
     if not text:
         return None
     try:
         return json.loads(text)
-    except:
-        pass
+    except json.JSONDecodeError:
+        pass # Try heuristic approach
+
+    cleaned = text.strip().strip("`").strip()
+    # Remove leading 'json' if present after stripping backticks
+    if cleaned.lower().startswith("json"):
+        cleaned = cleaned[4:].strip()
 
-    cleaned = text.strip().strip("`")
     try:
         start = cleaned.find("{")
         end = cleaned.rfind("}")
         if start >= 0 and end > start:
             return json.loads(cleaned[start : end + 1])
-    except:
+    except json.JSONDecodeError:
         pass
 
     return None
 
 # -----------------------------------------------------
-# FIXED: correct MCP tool registry access (v2)
+# Validate and Call Tool
 # -----------------------------------------------------
-def validate_and_call_tool(tool_name: str, tool_args: dict):
-    # ❌ old: if tool_name not in mcp.tools:
-    # ✔ new:
+def validate_and_call_tool(tool_name: str, tool_args: dict) -> Dict[str, Any]:
+    """Validates tool access and executes the corresponding function."""
     if tool_name not in mcp._tools:
         return {"error": f"Unknown or unauthorized tool '{tool_name}'"}
 
     try:
-        # ❌ old: mcp.tools[name](...)
-        # ✔ new:
         tool_fn = mcp._tools[tool_name]["function"]
         return tool_fn(**tool_args)
 
@@ -131,39 +144,39 @@ def validate_and_call_tool(tool_name: str, tool_args: dict):
 # -----------------------------------------------------
 # Main Pipeline
 # -----------------------------------------------------
-def process_and_describe(payload):
-
+def process_and_describe(payload: Dict[str, Any]) -> Dict[str, Any]:
+    """Main pipeline function to process image, call VLM, and execute tool."""
+    
+    # Input handling for gradio.JSON input which sometimes arrives as a string
     if isinstance(payload, str):
         try:
             payload = json.loads(payload)
-        except:
-            return {"error": "Invalid JSON payload"}
+        except json.JSONDecodeError:
+            return {"error": "Invalid JSON payload provided to the function"}
 
     hf_token = payload.get("hf_token")
     if not hf_token:
-        return {"error": "hf_token missing"}
+        return {"error": "hf_token missing in payload. Cannot authenticate with HF Hub."}
 
     robot_id = payload.get("robot_id", "unknown")
     image_b64 = payload.get("image_b64")
     if not image_b64:
-        return {"error": "image_b64 missing"}
+        return {"error": "image_b64 missing in payload"}
 
     # Save + Upload
-    local_tmp_path, hf_url, filename, size_bytes = save_and_upload_image(
-        image_b64, hf_token
-    )
+    _, hf_url, _, size_bytes = save_and_upload_image(image_b64, hf_token)
 
     if not hf_url:
         return {"error": "Image upload failed"}
 
     # VLM system prompt
-    system_prompt = """
+    system_prompt = f"""
 Respond in STRICT JSON ONLY:
-{
+{{
  "description": "short visual description",
- "tool_name": "speak | navigate | scan_hazard | analyze_human",
- "arguments": { ... }
-}
+ "tool_name": "{' | '.join(mcp._tools.keys())}",
+ "arguments": {{ ... }}
+}}
 """
 
     messages = [
@@ -182,12 +195,16 @@ Respond in STRICT JSON ONLY:
 
     client = InferenceClient(token=hf_token)
 
-    response = client.chat.completions.create(
-        model=HF_VLM_MODEL,
-        messages=messages,
-        max_tokens=300,
-        temperature=0.1,
-    )
+    try:
+        response = client.chat.completions.create(
+            model=HF_VLM_MODEL,
+            messages=messages,
+            max_tokens=300,
+            temperature=0.1,
+        )
+    except Exception as e:
+        return {"status": "error", "message": f"Inference API call failed: {str(e)}"}
+
 
     vlm_output = response.choices[0].message.content.strip()
 
@@ -199,7 +216,7 @@ Respond in STRICT JSON ONLY:
             "robot_id": robot_id,
             "image_url": hf_url,
             "vlm_raw": vlm_output,
-            "message": "VLM returned invalid JSON",
+            "message": "VLM returned invalid JSON format",
         }
 
     tool_name = parsed.get("tool_name")
@@ -224,8 +241,8 @@ Respond in STRICT JSON ONLY:
 # ------------------------------
 iface = gr.Interface(
     fn=process_and_describe,
-    inputs=gr.JSON(label="Input JSON"),
-    outputs=gr.JSON(label="Output JSON"),
+    inputs=gr.JSON(label="Input JSON Payload (must contain hf_token and image_b64)"),
+    outputs=gr.JSON(label="Output JSON Result"),
     api_name="predict",
     flagging_mode="never"
 )
@@ -234,5 +251,8 @@ iface = gr.Interface(
 # Main Entry
 # ------------------------------
 if __name__ == "__main__":
+    print(f"[Config] HF_DATASET_REPO: {HF_DATASET_REPO}")
+    print(f"[Config] HF_VLM_MODEL: {HF_VLM_MODEL}")
     print("[Gradio] Launching interface...")
     iface.launch(server_name="0.0.0.0", server_port=7860)
+