Update server.py

server.py

@@ -154,148 +154,145 @@ SOURCE:
     
 # Helper functions for different file types
 async def process_pdf(file_bytes: bytes) -> str:
-    """Process PDF file and extract text/info"""
-    pdf_file = io.BytesIO(file_bytes)
-    reader = PdfReader(pdf_file)
-    
-    num_pages = len(reader.pages)
-    text = ""
-    for page in reader.pages:
-        text += page.extract_text()
-    
-    return f"PDF Analysis:\n- Pages: {num_pages}\n- Text length: {len(text)} chars\n\nFirst 500 chars:\n{text[:500]}"
+    try:
+        pdf_file = io.BytesIO(file_bytes)
+        reader = PdfReader(pdf_file)
+
+        num_pages = len(reader.pages)
+        text = ""
+        for page in reader.pages:
+            extracted = page.extract_text() or ""
+            text += extracted
+
+        return (
+            f"PDF Analysis:\n"
+            f"- Pages: {num_pages}\n"
+            f"- Text length: {len(text)} chars\n\n"
+            f"First 500 chars:\n{text[:500]}"
+        )
+    except Exception as e:
+        return f"PDF processing failed: {str(e)}\nThis might not be a valid PDF file."
+
 
 async def process_image(file_bytes: bytes) -> str:
-    """Process image file"""
-    img = Image.open(io.BytesIO(file_bytes))
-    return f"Image Analysis:\n- Format: {img.format}\n- Size: {img.size}\n- Mode: {img.mode}"
+    try:
+        img = Image.open(io.BytesIO(file_bytes))
+        return (
+            f"Image Analysis:\n"
+            f"- Format: {img.format}\n"
+            f"- Size: {img.size[0]}x{img.size[1]} pixels\n"
+            f"- Mode: {img.mode}"
+        )
+    except Exception as e:
+        return f"Image processing failed: {str(e)}"
+
 
 async def process_text(file_bytes: bytes) -> str:
-    """Process text file"""
     try:
-        text = file_bytes.decode('utf-8')
-        lines = text.split('\n')
+        text = file_bytes.decode("utf-8")
+        lines = text.split("\n")
         words = len(text.split())
-        return f"Text Analysis:\n- Lines: {len(lines)}\n- Words: {words}\n- Characters: {len(text)}\n\nFirst 500 chars:\n{text[:500]}"
+
+        return (
+            f"Text Analysis:\n"
+            f"- Lines: {len(lines)}\n"
+            f"- Words: {words}\n"
+            f"- Characters: {len(text)}\n\n"
+            f"First 500 chars:\n{text[:500]}"
+        )
     except UnicodeDecodeError:
         return "Error: File is not valid UTF-8 text"
 
-def decode_file(file_data: str | bytes) -> bytes:
-    """Decode base64 file data to bytes"""
-    if isinstance(file_data, bytes):
-        return file_data
-    
-    # If it's a string, decode from base64
-    try:
-        return base64.b64decode(file_data)
-    except Exception as e:
-        raise ValueError(f"Failed to decode base64: {e}")
 
 def detect_file_type(file_bytes: bytes) -> str:
-    """Detect file type from magic bytes"""
-    if file_bytes.startswith(b'%PDF'):
+    if file_bytes.startswith(b"%PDF"):
         return "pdf"
-    elif file_bytes.startswith(b'\x89PNG'):
+    elif file_bytes.startswith(b"\x89PNG"):
         return "image"
-    elif file_bytes.startswith(b'\xff\xd8\xff'):
+    elif file_bytes.startswith(b"\xff\xd8\xff"):
         return "image"  # JPEG
-    elif file_bytes.startswith(b'GIF87a') or file_bytes.startswith(b'GIF89a'):
-        return "image"  # GIF
-    elif file_bytes.startswith(b'\x42\x4d'):
+    elif file_bytes.startswith(b"GIF87a") or file_bytes.startswith(b"GIF89a"):
+        return "image"
+    elif file_bytes.startswith(b"\x42\x4d"):
         return "image"  # BMP
     else:
-        # Try to decode as text
         try:
-            file_bytes[:1000].decode('utf-8')
+            file_bytes[:1000].decode("utf-8")
             return "text"
         except:
             return "unknown"
 
-async def process_pdf(file_bytes: bytes) -> str:
-    """Process PDF file and extract text/info"""
-    try:
-        pdf_file = io.BytesIO(file_bytes)
-        reader = PdfReader(pdf_file)
-        
-        num_pages = len(reader.pages)
-        text = ""
-        for page in reader.pages:
-            text += page.extract_text()
-        
-        return f"PDF Analysis:\n- Pages: {num_pages}\n- Text length: {len(text)} chars\n\nFirst 500 chars:\n{text[:500]}"
-    except Exception as e:
-        return f"PDF processing failed: {str(e)}\nThis might not be a valid PDF file."
 
-async def process_image(file_bytes: bytes) -> str:
-    """Process image file"""
-    try:
-        img = Image.open(io.BytesIO(file_bytes))
-        return f"Image Analysis:\n- Format: {img.format}\n- Size: {img.size[0]}x{img.size[1]} pixels\n- Mode: {img.mode}"
-    except Exception as e:
-        return f"Image processing failed: {str(e)}"
+# ---------------- MCP TOOL ---------------- #
 
-async def process_text(file_bytes: bytes) -> str:
-    """Process text file"""
-    try:
-        text = file_bytes.decode('utf-8')
-        lines = text.split('\n')
-        words = len(text.split())
-        return f"Text Analysis:\n- Lines: {len(lines)}\n- Words: {words}\n- Characters: {len(text)}\n\nFirst 500 chars:\n{text[:500]}"
-    except UnicodeDecodeError:
-        return "Error: File is not valid UTF-8 text"
-
-@mcp.tool 
+@mcp.tool
 async def analyze_file(
-    ctx: Context, 
-    file: str,  # MCP sends base64-encoded string
-    file_type: Optional[Literal["text", "image", "pdf"]] = None
-) -> str: 
+    ctx: Context,
+    file_uri: str,  # <-- file path/uri from client
+    file_type: Optional[Literal["text", "image", "pdf"]] = None,
+) -> str:
     """
-    Analyzes a file with auto-detection.
-    
-    Args:
-        file: The file content (base64-encoded)
-        file_type: Optional. If not provided, will auto-detect.
+    Analyze a file using MCP resource URI (recommended).
+
+    file_uri examples:
+      - file:///C:/Users/Admin/Desktop/sample.pdf
+      - resource://uploads/123/sample.pdf
     """
-    # Decode base64 to bytes
+
+    # Read file bytes via MCP resource API
     try:
-        file_bytes = decode_file(file)
-    except ValueError as e:
-        return f"Error decoding file: {e}"
-    
-    info = f"Received file of size {len(file_bytes)} bytes.\n"
-    
-    # Auto-detect if not provided
+        resource_result = await ctx.read_resource(file_uri)
+    except Exception as e:
+        return f"Failed to read resource '{file_uri}': {str(e)}"
+
+    if not resource_result.contents:
+        return f"No content returned for resource: {file_uri}"
+
+    content = resource_result.contents[0]
+
+    # If it's text resource
+    if hasattr(content, "text") and content.text is not None:
+        file_bytes = content.text.encode("utf-8")
+    else:
+        # Blob resource (bytes might be base64 encoded depending on client)
+        blob_data = content.blob
+
+        # some clients give blob as base64 string
+        if isinstance(blob_data, str):
+            file_bytes = base64.b64decode(blob_data)
+        else:
+            file_bytes = blob_data
+
+    info = f"Loaded file from URI: {file_uri}\nSize: {len(file_bytes)} bytes\n"
+
+    # Auto detect if not specified
     if file_type is None:
-        file_type = detect_file_type(file_bytes)
-        info += f"Auto-detected type: {file_type}\n\n"
-        
-        if file_type == "unknown":
-            return info + "Could not detect file type. Please call again with file_type parameter (text/image/pdf)"
+        detected = detect_file_type(file_bytes)
+        info += f"Auto-detected type: {detected}\n\n"
+
+        if detected == "unknown":
+            return info + "Could not detect file type. Please specify file_type=text/image/pdf"
+
+        file_type = detected
     else:
-        # Verify the file_type matches actual content
         detected = detect_file_type(file_bytes)
-        info += f"Specified type: {file_type}\n"
-        info += f"Detected type: {detected}\n\n"
-        
+        info += f"Specified type: {file_type}\nDetected type: {detected}\n\n"
+
         if detected != file_type and detected != "unknown":
-            info += f"⚠️ Warning: Specified type '{file_type}' doesn't match detected type '{detected}'\n\n"
-    
-    # Process based on file_type
-    try:
-        if file_type == "pdf":
-            result = await process_pdf(file_bytes)
-        elif file_type == "image":
-            result = await process_image(file_bytes)
-        elif file_type == "text":
-            result = await process_text(file_bytes)
-        else:
-            return info + f"Unknown file type: {file_type}"
-        
-        return info + result
-    
-    except Exception as e:
-        return info + f"Error processing {file_type}: {str(e)}"
+            info += f"⚠️ Warning: mismatch (specified={file_type}, detected={detected})\n\n"
+
+    # Process file
+    if file_type == "pdf":
+        result = await process_pdf(file_bytes)
+    elif file_type == "image":
+        result = await process_image(file_bytes)
+    elif file_type == "text":
+        result = await process_text(file_bytes)
+    else:
+        return info + f"Unsupported file_type: {file_type}"
+
+    return info + result
+
         
 if __name__ == "__main__":
     print("Starting FastMCP server on port 7860...")