Update app.py

app.py

@@ -414,158 +414,174 @@ def build_kg(url_or_text):
         url_or_text: URL to analyze or direct text input
         
     Returns:
-        String: JSON containing entities, relationships, and simple diagram
+        String: Simple JSON response optimized for MCP streaming
     """
     try:
+        # Quick validation
         if not url_or_text or len(url_or_text.strip()) == 0:
-            return json.dumps({"error": "Please provide some text or a URL to analyze."}, ensure_ascii=True)
+            return '{"error":"Please provide text or URL to analyze"}'
         
-        # Step 1: Fetch content
-        content = fetch_content(url_or_text)
+        # Limit input size immediately to prevent timeouts
+        input_text = url_or_text[:2000] if len(url_or_text) > 2000 else url_or_text
         
-        if content.startswith("Error"):
-            return json.dumps({"error": content}, ensure_ascii=True)
-        
-        # Step 2: Extract entities
-        entities_data = extract_entities(content)
-        
-        # Step 3: Create simple ASCII diagram (truncated for MCP)
-        entities = entities_data.get("entities", [])
-        relationships = entities_data.get("relationships", [])
-        
-        # Limit diagram size for MCP compatibility
-        limited_entities = entities[:8]  # Reduce from 10 to 8
-        limited_relationships = relationships[:6]  # Reduce from 8 to 6
-        ascii_diagram = build_ascii_diagram(limited_entities, limited_relationships)
-        
-        # Truncate diagram if too long
-        if len(ascii_diagram) > 1000:
-            ascii_diagram = ascii_diagram[:950] + "\n... (truncated for MCP compatibility)"
-        
-        # Step 4: Create summary (more concise for MCP)
-        num_entities = len(entities)
-        num_relationships = len(relationships)
-        
-        summary = f"Knowledge Graph Analysis Complete!\n\nStats: {num_entities} entities, {num_relationships} relationships\n\nTop Entities:"
-        
-        for entity in entities[:3]:  # Show only first 3 for MCP
-            name = entity.get('name', 'Unknown')
-            entity_type = entity.get('type', 'UNKNOWN')
-            summary += f"\n• {name} ({entity_type})"
+        # Step 1: Fetch content (with timeout protection)
+        try:
+            content = fetch_content(input_text)
+            if content.startswith("Error"):
+                return f'{{"error":"{content}"}}'
+        except Exception:
+            content = input_text  # Use input directly if fetch fails
         
-        if len(entities) > 3:
-            summary += f"\n... and {len(entities) - 3} more"
+        # Limit content size for fast processing
+        content = content[:1500] if len(content) > 1500 else content
         
-        # Create MCP-friendly response with size limits
+        # Step 2: Quick entity extraction (simplified for speed)
+        try:
+            entities_data = simple_entity_extraction(content)  # Always use simple extraction for MCP
+        except Exception:
+            entities_data = {"entities": [], "relationships": []}
+        
+        # Step 3: Minimal response
+        entities = entities_data.get("entities", [])[:5]  # Max 5 entities
+        relationships = entities_data.get("relationships", [])[:3]  # Max 3 relationships
+        
+        # Create minimal ASCII summary
+        diagram_parts = []
+        if entities:
+            diagram_parts.append("ENTITIES:")
+            for entity in entities:
+                name = str(entity.get("name", "Unknown"))[:20]  # Truncate names
+                diagram_parts.append(f"  - {name}")
+        
+        if relationships:
+            diagram_parts.append("RELATIONSHIPS:")
+            for rel in relationships:
+                source = str(rel.get("source", ""))[:15]
+                target = str(rel.get("target", ""))[:15]
+                diagram_parts.append(f"  {source} -> {target}")
+        
+        diagram = "\n".join(diagram_parts) if diagram_parts else "No entities found"
+        
+        # Ultra-minimal response
         response = {
             "success": True,
-            "entities": entities[:15],  # Limit entities for MCP
-            "relationships": relationships[:10],  # Limit relationships for MCP
-            "diagram": ascii_diagram,
-            "summary": summary,
-            "stats": {
-                "entity_count": num_entities,
-                "relationship_count": num_relationships,
-                "content_length": min(len(content), 5000)  # Cap reported length
-            }
+            "entity_count": len(entities),
+            "relationship_count": len(relationships), 
+            "entities": [{"name": e.get("name", "")[:20], "type": e.get("type", "UNKNOWN")} for e in entities],
+            "relationships": [{"source": r.get("source", "")[:15], "target": r.get("target", "")[:15]} for r in relationships],
+            "diagram": diagram[:500]  # Strict limit
         }
         
-        # Validate and sanitize response for MCP compatibility
-        validated_response = validate_mcp_response(response)
-        
-        # Return compact JSON
-        return json.dumps(validated_response, ensure_ascii=True, separators=(',', ':'))
+        # Return ultra-compact JSON
+        return json.dumps(response, separators=(',', ':'))[:2000]  # Hard size limit
             
     except Exception as e:
-        error_msg = f"Analysis failed: {str(e)}"
-        return json.dumps({"success": False, "error": error_msg}, ensure_ascii=True, separators=(',', ':'))
+        # Ultra-simple error response
+        error_msg = str(e)[:100]  # Truncate error message
+        return f'{{"success":false,"error":"{error_msg}"}}'
+
+# Wrapper function with timeout protection for MCP
+def mcp_safe_build_kg(url_or_text):
+    """MCP-safe wrapper with timeout protection"""
+    try:
+        import signal
+        import functools
+        
+        def timeout_handler(signum, frame):
+            raise TimeoutError("Function timed out")
+        
+        # Set timeout for 10 seconds
+        signal.signal(signal.SIGALRM, timeout_handler)
+        signal.alarm(10)
+        
+        try:
+            result = build_kg(url_or_text)
+            signal.alarm(0)  # Cancel timeout
+            return result
+        except TimeoutError:
+            return '{"success":false,"error":"Request timed out"}'
+        except Exception as e:
+            signal.alarm(0)  # Cancel timeout
+            return f'{{"success":false,"error":"Function error: {str(e)[:50]}"}}'
+            
+    except Exception:
+        # Fallback if signal not available (Windows, etc.)
+        try:
+            return build_kg(url_or_text)
+        except Exception as e:
+            return f'{{"success":false,"error":"Fallback error: {str(e)[:50]}"}}'
 
 # Create Gradio interface with error handling
 try:
     demo = gr.Interface(
-        fn=build_kg,
-        inputs="text",
-        outputs="text",
-        title="🧠 KG Builder",
-        description="Transform text or URLs into knowledge graphs for AI agents.",
-        allow_flagging="never"
+        fn=mcp_safe_build_kg,  # Use the timeout-protected version
+        inputs=gr.Textbox(
+            label="Input Text or URL",
+            placeholder="Enter text to analyze or paste a URL...",
+            max_lines=5
+        ),
+        outputs=gr.Textbox(
+            label="Knowledge Graph JSON",
+            show_copy_button=True
+        ),
+        title="KG Builder - MCP Edition",
+        description="Lightweight knowledge graph builder optimized for MCP servers.",
+        allow_flagging="never",
+        cache_examples=False
     )
     
 except Exception as e:
     print(f"Failed to create Gradio interface: {e}")
-    # Create a simple fallback interface
-    def simple_demo(text):
-        return json.dumps({"error": f"Startup failed: {str(e)}"}), None, "Application failed to start properly."
-    
-    try:
-        demo = gr.Interface(
-            fn=simple_demo,
-            inputs=[gr.Textbox(label="Input", placeholder="Enter text...")],
-            outputs=[
-                gr.JSON(label="Error Output"),
-                gr.Image(label="No Image", type="filepath"),
-                gr.Markdown(label="Error Message")
-            ],
-            title="⚠️ KG Builder - Startup Error",
-            allow_flagging="never",
-            cache_examples=False
-        )
-    except Exception as e2:
-        # Ultimate fallback - minimal interface
-        def minimal_demo(text):
-            return {"error": str(e2)}
-        
-        demo = gr.Interface(
-            fn=minimal_demo,
-            inputs="text",
-            outputs="json",
-            title="KG Builder Error"
-        )
+    # Create minimal fallback
+    def error_demo(text):
+        return f'{{"error":"Interface creation failed: {str(e)[:100]}"}}'
+        
+    demo = gr.Interface(
+        fn=error_demo,
+        inputs="text",
+        outputs="text",
+        title="KG Builder - Error Mode",
+        allow_flagging="never"
+    )
 
 # Launch the demo
 if __name__ == "__main__":
+    print("Starting KG Builder MCP Server...")
+    
     try:
-        # Check if running in HuggingFace Spaces or other cloud environment
-        is_hf_space = os.environ.get("SPACE_ID") is not None
-        
-        if is_hf_space:
-            # HuggingFace Spaces configuration
-            demo.launch(
-                mcp_server=True,
-                share=False,
-                show_error=True,
-                quiet=False,
-                server_name="0.0.0.0",  # Allow external connections
-                server_port=7860       # Standard HF port
-            )
-        else:
-            # Local development configuration
-            demo.launch(
-                mcp_server=True, 
-                share=False,
-                show_error=True,
-                quiet=False
-            )
+        demo.launch(
+            mcp_server=True,
+            share=False,
+            show_error=False,  # Reduce error verbosity for MCP
+            quiet=True,        # Reduce logging to prevent SSE issues
+            server_name="0.0.0.0",
+            server_port=7860,
+            max_threads=1,     # Limit concurrency to prevent resource issues
+            show_api=False     # Disable API docs to reduce overhead
+        )
     except Exception as e:
-        print(f"MCP Launch failed: {e}")
-        # Try without MCP server as fallback
+        print(f"MCP server launch failed: {e}")
+        print("Trying fallback mode...")
         try:
+            # Fallback without MCP
             demo.launch(
                 mcp_server=False,
                 share=False,
-                show_error=True,
-                quiet=True
+                quiet=True,
+                show_error=False
             )
         except Exception as e2:
-            print(f"Complete failure: {e2}")
-            # Create minimal working demo
-            def error_demo(text):
-                return json.dumps({"error": f"Server startup failed: {str(e2)}"}, ensure_ascii=True)
+            print(f"All launch attempts failed: {e2}")
+            print("Creating emergency fallback...")
+            
+            # Create absolute minimal demo
+            def emergency_demo(text):
+                return '{"error":"Server in emergency mode"}'
             
-            minimal_demo = gr.Interface(
-                fn=error_demo,
-                inputs="text",
+            emergency = gr.Interface(
+                fn=emergency_demo,
+                inputs="text", 
                 outputs="text",
-                title="KG Builder - Error Mode"
+                title="KG Builder Emergency Mode"
             )
-            minimal_demo.launch(share=False, quiet=True)
+            emergency.launch(quiet=True, share=False)