Update app.py

app.py

@@ -335,6 +335,42 @@ def build_knowledge_graph(entities_data):
         
         return pil_image
 
+def build_ascii_diagram(entities, relationships):
+    """Create simple ASCII diagram of knowledge graph"""
+    if not entities:
+        return "No entities to visualize"
+    
+    diagram = "KNOWLEDGE GRAPH DIAGRAM:\n"
+    diagram += "=" * 50 + "\n\n"
+    
+    # Show entities by type
+    entity_types = {}
+    for entity in entities[:10]:  # Limit to 10
+        etype = entity.get("type", "UNKNOWN")
+        if etype not in entity_types:
+            entity_types[etype] = []
+        entity_types[etype].append(entity.get("name", "Unknown"))
+    
+    for etype, names in entity_types.items():
+        diagram += f"📍 {etype}:\n"
+        for name in names:
+            diagram += f"   • {name}\n"
+        diagram += "\n"
+    
+    # Show relationships
+    if relationships:
+        diagram += "🔗 RELATIONSHIPS:\n"
+        for rel in relationships[:8]:  # Limit to 8
+            source = rel.get("source", "?")
+            target = rel.get("target", "?")
+            relation = rel.get("relation", "related")
+            diagram += f"   {source} --[{relation}]--> {target}\n"
+        
+        if len(relationships) > 8:
+            diagram += f"   ... and {len(relationships) - 8} more relationships\n"
+    
+    return diagram
+
 def build_kg(url_or_text):
     """Main function to build knowledge graph from URL or text.
     
@@ -342,7 +378,7 @@ def build_kg(url_or_text):
         url_or_text: URL to analyze or direct text input
         
     Returns:
-        String: JSON containing entities, relationships, and summary
+        String: JSON containing entities, relationships, and simple diagram
     """
     try:
         if not url_or_text or len(url_or_text.strip()) == 0:
@@ -357,100 +393,61 @@ def build_kg(url_or_text):
         # Step 2: Extract entities
         entities_data = extract_entities(content)
         
-        # Step 3: Create summary (no image for MCP compatibility)
-        num_entities = len(entities_data.get("entities", []))
-        num_relationships = len(entities_data.get("relationships", []))
+        # Step 3: Create simple ASCII diagram
+        entities = entities_data.get("entities", [])
+        relationships = entities_data.get("relationships", [])
+        ascii_diagram = build_ascii_diagram(entities, relationships)
+        
+        # Step 4: Create summary
+        num_entities = len(entities)
+        num_relationships = len(relationships)
         
         summary = f"""Knowledge Graph Analysis Complete!
 
 Statistics:
-- Entities found: {num_entities}
-- Relationships found: {num_relationships}
+- Entities: {num_entities}
+- Relationships: {num_relationships}
 - Content length: {len(content)} characters
 
-Extracted Entities:"""
+Top Entities:"""
         
-        for entity in entities_data.get("entities", [])[:8]:  # Show first 8
+        for entity in entities[:5]:  # Show first 5
             name = entity.get('name', 'Unknown')
             entity_type = entity.get('type', 'UNKNOWN')
-            desc = entity.get('description', 'No description')
-            summary += f"\n• {name} ({entity_type}): {desc}"
+            summary += f"\n• {name} ({entity_type})"
         
-        if len(entities_data.get("entities", [])) > 8:
-            summary += f"\n... and {len(entities_data.get('entities', [])) - 8} more entities"
+        if len(entities) > 5:
+            summary += f"\n... and {len(entities) - 5} more"
         
-        # Create complete response
+        # Create MCP-friendly response
         response = {
             "success": True,
-            "entities": entities_data.get("entities", []),
-            "relationships": entities_data.get("relationships", []),
+            "entities": entities,
+            "relationships": relationships,
+            "diagram": ascii_diagram,
             "summary": summary,
-            "statistics": {
+            "stats": {
                 "entity_count": num_entities,
                 "relationship_count": num_relationships,
                 "content_length": len(content)
             }
         }
         
-        return json.dumps(response, indent=2, ensure_ascii=True)
+        return json.dumps(response, ensure_ascii=True)
             
     except Exception as e:
         error_msg = f"Analysis failed: {str(e)}"
-        print(f"Function error: {error_msg}")
         return json.dumps({"success": False, "error": error_msg})
 
 # Create Gradio interface with error handling
 try:
-    def parse_kg_output(json_result):
-        """Parse the JSON result and return separate components for UI"""
-        try:
-            data = json.loads(json_result) if isinstance(json_result, str) else json_result
-            if data.get("success"):
-                entities_json = {
-                    "entities": data.get("entities", []),
-                    "relationships": data.get("relationships", [])
-                }
-                summary = data.get("summary", "No summary available")
-                return json.dumps(entities_json, indent=2), summary
-            else:
-                error_msg = data.get("error", "Unknown error")
-                return json.dumps({"error": error_msg}), f"Error: {error_msg}"
-        except Exception as e:
-            return json.dumps({"error": str(e)}), f"Parse error: {str(e)}"
-    
-    def kg_interface(url_or_text):
-        """Interface wrapper for the knowledge graph builder"""
-        result = build_kg(url_or_text)
-        return parse_kg_output(result)
-    
     demo = gr.Interface(
-        fn=kg_interface,
-        inputs=[
-            gr.Textbox(
-                label="URL or Text Input", 
-                placeholder="Enter a URL (https://example.com) or paste text directly...",
-                lines=3,
-                info="Enter a website URL to analyze, or paste text content directly"
-            )
-        ],
-        outputs=[
-            gr.JSON(label="Extracted Entities & Relationships"),
-            gr.Markdown(label="Analysis Summary")
-        ],
+        fn=build_kg,
+        inputs="text",
+        outputs="text",
         title="🧠 KG Builder",
-        description="""
-        **Transform any text or webpage into knowledge graphs!**
-        
-        This tool:
-        1. 📖 Extracts content from URLs or analyzes your text
-        2. 🤖 Uses AI to identify entities and relationships  
-        3. 📊 Provides detailed analysis summaries
-        
-        **Try with:** news articles, Wikipedia pages, or any text content
-        """,
-        theme=gr.themes.Soft(),
-        allow_flagging="never",
-        cache_examples=False  # Disable example caching to prevent startup errors
+        description="Transform text or URLs into knowledge graphs for AI agents.",
+        allow_flagging="never"
     )
     
 except Exception as e: