Update app.py

app.py

@@ -11,6 +11,22 @@ import base64
 from huggingface_hub import InferenceClient
 import re
 from urllib.parse import urlparse
+import warnings
+
+# Configure matplotlib for better font handling
+plt.rcParams['font.family'] = ['DejaVu Sans', 'Arial', 'Liberation Sans']
+plt.rcParams['font.size'] = 10
+warnings.filterwarnings('ignore', category=UserWarning)
+
+def clean_text_for_display(text):
+    """Clean text to remove characters that might cause font issues."""
+    if not isinstance(text, str):
+        return str(text)
+    
+    # Remove or replace problematic characters
+    text = re.sub(r'[^\x00-\x7F]+', '', text)  # Remove non-ASCII characters
+    text = re.sub(r'\s+', ' ', text).strip()   # Normalize whitespace
+    return text[:50] if len(text) > 50 else text  # Limit length for display
 
 def fetch_content(url_or_text):
     """Fetch content from URL or return text directly.
@@ -21,38 +37,79 @@ def fetch_content(url_or_text):
     Returns:
         Extracted text content
     """
-    # Check if input looks like a URL
-    parsed = urlparse(url_or_text)
-    if parsed.scheme in ['http', 'https']:
-        try:
-            headers = {
-                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
-            }
-            response = requests.get(url_or_text, headers=headers, timeout=10)
-            response.raise_for_status()
-            
-            # Parse HTML and extract text
-            soup = BeautifulSoup(response.content, 'html.parser')
-            
-            # Remove script and style elements
-            for script in soup(["script", "style"]):
-                script.decompose()
-            
-            # Get text and clean it up
-            text = soup.get_text()
-            lines = (line.strip() for line in text.splitlines())
-            chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
-            text = ' '.join(chunk for chunk in chunks if chunk)
-            
-            return text[:5000]  # Limit to first 5000 characters
-        except Exception as e:
-            return f"Error fetching URL: {str(e)}"
-    else:
-        # It's direct text input
-        return url_or_text
+    try:
+        # Check if input looks like a URL
+        parsed = urlparse(url_or_text)
+        if parsed.scheme in ['http', 'https']:
+            try:
+                headers = {
+                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+                }
+                response = requests.get(url_or_text, headers=headers, timeout=10)
+                response.raise_for_status()
+                
+                # Parse HTML and extract text
+                soup = BeautifulSoup(response.content, 'html.parser')
+                
+                # Remove script and style elements
+                for script in soup(["script", "style"]):
+                    script.decompose()
+                
+                # Get text and clean it up
+                text = soup.get_text()
+                lines = (line.strip() for line in text.splitlines())
+                chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
+                text = ' '.join(chunk for chunk in chunks if chunk)
+                
+                return text[:5000]  # Limit to first 5000 characters
+            except Exception as e:
+                return f"Error fetching URL: {str(e)}"
+        else:
+            # It's direct text input
+            return url_or_text
+    except Exception as e:
+        return f"Error processing input: {str(e)}"
+
+def simple_entity_extraction(text):
+    """Fallback entity extraction when AI is not available."""
+    try:
+        words = text.split()
+        entities = []
+        
+        # Simple heuristic: words that are capitalized and longer than 2 characters
+        seen = set()
+        for word in words[:30]:  # Limit to first 30 words
+            clean_word = re.sub(r'[^\w]', '', word)
+            if (clean_word.istitle() and len(clean_word) > 2 and 
+                clean_word.lower() not in seen and 
+                clean_word not in ['The', 'This', 'That', 'When', 'Where', 'How']):
+                entities.append({
+                    "name": clean_text_for_display(clean_word),
+                    "type": "CONCEPT",
+                    "description": "Auto-detected entity"
+                })
+                seen.add(clean_word.lower())
+        
+        # Create some basic relationships
+        relationships = []
+        if len(entities) > 1:
+            for i in range(min(len(entities) - 1, 5)):  # Max 5 relationships
+                relationships.append({
+                    "source": entities[i]["name"],
+                    "target": entities[i + 1]["name"],
+                    "relation": "related_to",
+                    "description": "Sequential relationship"
+                })
+        
+        return {"entities": entities[:10], "relationships": relationships}
+    except Exception as e:
+        return {
+            "entities": [{"name": "Error", "type": "ERROR", "description": str(e)}],
+            "relationships": []
+        }
 
 def extract_entities(text):
-    """Extract entities and relationships using Mistral.
+    """Extract entities and relationships using Mistral AI with fallback.
     
     Args:
         text: Input text to analyze
@@ -61,9 +118,15 @@ def extract_entities(text):
         Dictionary containing entities and relationships
     """
     try:
+        # Check if HF_TOKEN is available
+        hf_token = os.environ.get("HF_TOKEN")
+        if not hf_token:
+            print("No HF_TOKEN found, using simple extraction")
+            return simple_entity_extraction(text)
+        
         client = InferenceClient(
             provider="together",
-            api_key=os.environ.get("HF_TOKEN"),
+            api_key=hf_token,
         )
         
         prompt = f"""
@@ -71,52 +134,51 @@ def extract_entities(text):
         1. Named entities (people, organizations, locations, concepts)
         2. Relationships between these entities
         
-        Return the result as a JSON object with this structure:
+        Return ONLY a valid JSON object with this structure:
         {{
             "entities": [
-                {{"name": "entity_name", "type": "PERSON|ORG|LOCATION|CONCEPT", "description": "brief description"}}
+                {{"name": "entity_name", "type": "PERSON", "description": "brief description"}}
             ],
             "relationships": [
                 {{"source": "entity1", "target": "entity2", "relation": "relationship_type", "description": "brief description"}}
             ]
         }}
         
-        Text to analyze:
-        {text[:2000]}
-        
-        JSON:"""
+        Text to analyze: {text[:1500]}
+        """
         
         completion = client.chat.completions.create(
             model="mistralai/Mistral-Small-24B-Instruct-2501",
-            messages=[
-                {
-                    "role": "user",
-                    "content": prompt
-                }
-            ],
-            max_tokens=1500,
-            temperature=0.3,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=1000,
+            temperature=0.1,
         )
         
         response_text = completion.choices[0].message.content
         
-        # Extract JSON from response
+        # Clean and extract JSON
         json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
         if json_match:
             json_str = json_match.group()
-            return json.loads(json_str)
-        else:
-            # Fallback: create simple entities from text
-            words = text.split()
-            entities = []
-            for i, word in enumerate(words[:20]):  # Limit to first 20 words
-                if word.istitle() and len(word) > 2:
-                    entities.append({"name": word, "type": "CONCEPT", "description": "Extracted entity"})
+            # Clean the JSON string
+            json_str = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', json_str)  # Remove control characters
+            
+            parsed_data = json.loads(json_str)
+            
+            # Clean entity names for display
+            if "entities" in parsed_data:
+                for entity in parsed_data["entities"]:
+                    if "name" in entity:
+                        entity["name"] = clean_text_for_display(entity["name"])
             
-            return {"entities": entities, "relationships": []}
+            return parsed_data
+        else:
+            print("No valid JSON found in AI response, using fallback")
+            return simple_entity_extraction(text)
             
     except Exception as e:
-        return {"entities": [{"name": "Error", "type": "ERROR", "description": str(e)}], "relationships": []}
+        print(f"AI extraction failed: {e}, using fallback")
+        return simple_entity_extraction(text)
 
 def build_knowledge_graph(entities_data):
     """Build and visualize knowledge graph.
@@ -133,70 +195,71 @@ def build_knowledge_graph(entities_data):
         
         # Add nodes (entities)
         entities = entities_data.get("entities", [])
-        for entity in entities:
-            G.add_node(entity["name"], 
-                      type=entity.get("type", "UNKNOWN"),
-                      description=entity.get("description", ""))
+        for entity in entities[:15]:  # Limit to 15 entities for better visualization
+            clean_name = clean_text_for_display(entity.get("name", "Unknown"))
+            if clean_name and len(clean_name.strip()) > 0:
+                G.add_node(clean_name, 
+                          type=entity.get("type", "UNKNOWN"),
+                          description=entity.get("description", ""))
         
         # Add edges (relationships)
         relationships = entities_data.get("relationships", [])
         for rel in relationships:
-            if rel["source"] in G.nodes and rel["target"] in G.nodes:
-                G.add_edge(rel["source"], rel["target"], 
+            source = clean_text_for_display(rel.get("source", ""))
+            target = clean_text_for_display(rel.get("target", ""))
+            if source in G.nodes and target in G.nodes:
+                G.add_edge(source, target, 
                           relation=rel.get("relation", "related"),
                           description=rel.get("description", ""))
         
         # If no relationships found, create some connections between entities
-        if len(relationships) == 0 and len(entities) > 1:
-            entity_names = [e["name"] for e in entities[:10]]  # Limit to 10
-            for i in range(len(entity_names) - 1):
-                G.add_edge(entity_names[i], entity_names[i + 1], relation="related")
+        if len(relationships) == 0 and len(list(G.nodes())) > 1:
+            node_list = list(G.nodes())
+            for i in range(min(len(node_list) - 1, 5)):
+                G.add_edge(node_list[i], node_list[i + 1], relation="related")
         
         # Create visualization
-        fig, ax = plt.subplots(figsize=(12, 8))
-        
-        # Position nodes using spring layout
-        pos = nx.spring_layout(G, k=2, iterations=50)
-        
-        # Color nodes by type
-        node_colors = []
-        type_colors = {
-            "PERSON": "#FF6B6B",
-            "ORG": "#4ECDC4", 
-            "LOCATION": "#45B7D1",
-            "CONCEPT": "#96CEB4",
-            "ERROR": "#FF0000",
-            "UNKNOWN": "#DDA0DD"
-        }
-        
-        for node in G.nodes():
-            node_type = G.nodes[node].get('type', 'UNKNOWN')
-            node_colors.append(type_colors.get(node_type, "#DDA0DD"))
-        
-        # Draw the graph
-        nx.draw(G, pos, 
-                node_color=node_colors,
-                node_size=1000,
-                font_size=8,
-                font_weight='bold',
-                with_labels=True,
-                edge_color='gray',
-                width=2,
-                alpha=0.7,
-                ax=ax)
-        
-        # Add title
-        ax.set_title("Knowledge Graph", size=16, weight='bold')
-        
-        # Add legend
-        legend_elements = []
-        for type_name, color in type_colors.items():
-            if any(G.nodes[node].get('type') == type_name for node in G.nodes()):
-                legend_elements.append(plt.Line2D([0], [0], marker='o', color='w', 
-                                                markerfacecolor=color, markersize=10, label=type_name))
-        
-        if legend_elements:
-            ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.15, 1))
+        fig, ax = plt.subplots(figsize=(10, 8))
+        
+        # Skip if no nodes
+        if len(G.nodes()) == 0:
+            ax.text(0.5, 0.5, "No entities found to visualize", 
+                   ha='center', va='center', fontsize=14, transform=ax.transAxes)
+            ax.set_title("Knowledge Graph")
+            ax.axis('off')
+        else:
+            # Position nodes using spring layout
+            pos = nx.spring_layout(G, k=1, iterations=50)
+            
+            # Color nodes by type
+            node_colors = []
+            type_colors = {
+                "PERSON": "#FF6B6B",
+                "ORG": "#4ECDC4", 
+                "LOCATION": "#45B7D1",
+                "CONCEPT": "#96CEB4",
+                "ERROR": "#FF0000",
+                "UNKNOWN": "#DDA0DD"
+            }
+            
+            for node in G.nodes():
+                node_type = G.nodes[node].get('type', 'UNKNOWN')
+                node_colors.append(type_colors.get(node_type, "#DDA0DD"))
+            
+            # Draw the graph
+            nx.draw(G, pos, 
+                    node_color=node_colors,
+                    node_size=800,
+                    font_size=8,
+                    font_weight='bold',
+                    with_labels=True,
+                    edge_color='gray',
+                    width=1.5,
+                    alpha=0.8,
+                    ax=ax)
+            
+            # Add title
+            ax.set_title("Knowledge Graph", size=14, weight='bold')
         
         # Convert to PIL Image
         fig.canvas.draw()
@@ -210,14 +273,13 @@ def build_knowledge_graph(entities_data):
         return pil_image
         
     except Exception as e:
-        # Create error image
+        # Create simple error image
         fig, ax = plt.subplots(figsize=(8, 6))
-        ax.text(0.5, 0.5, f"Error creating graph:\n{str(e)}", 
+        ax.text(0.5, 0.5, f"Error creating graph", 
                 ha='center', va='center', fontsize=12, transform=ax.transAxes)
         ax.set_title("Knowledge Graph Error")
         ax.axis('off')
         
-        # Convert to PIL Image
         fig.canvas.draw()
         img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
         img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
@@ -238,11 +300,14 @@ def knowledge_graph_builder(url_or_text):
         Tuple of (entities_json, graph_image, summary)
     """
     try:
+        if not url_or_text or len(url_or_text.strip()) == 0:
+            return "{}", None, "Please provide some text or a URL to analyze."
+        
         # Step 1: Fetch content
         content = fetch_content(url_or_text)
         
         if content.startswith("Error"):
-            return content, None, "Failed to fetch content"
+            return json.dumps({"error": content}), None, content
         
         # Step 2: Extract entities
         entities_data = extract_entities(content)
@@ -254,65 +319,89 @@ def knowledge_graph_builder(url_or_text):
         num_entities = len(entities_data.get("entities", []))
         num_relationships = len(entities_data.get("relationships", []))
         
-        summary = f"""
-        Knowledge Graph Analysis Complete!
-        
-        📊 **Statistics:**
-        - Entities found: {num_entities}
-        - Relationships found: {num_relationships}
-        - Content length: {len(content)} characters
+        summary = f"""## Knowledge Graph Analysis Complete!
+
+📊 **Statistics:**
+- Entities found: {num_entities}
+- Relationships found: {num_relationships}
+- Content length: {len(content)} characters
+
+🔍 **Extracted Entities:**"""
         
-        🔍 **Extracted Entities:**
-        """
+        for entity in entities_data.get("entities", [])[:8]:  # Show first 8
+            name = entity.get('name', 'Unknown')
+            entity_type = entity.get('type', 'UNKNOWN')
+            desc = entity.get('description', 'No description')
+            summary += f"\n• **{name}** ({entity_type}): {desc}"
         
-        for entity in entities_data.get("entities", [])[:10]:  # Show first 10
-            summary += f"\n• **{entity['name']}** ({entity.get('type', 'UNKNOWN')}): {entity.get('description', 'No description')}"
+        if len(entities_data.get("entities", [])) > 8:
+            summary += f"\n\n... and {len(entities_data.get('entities', [])) - 8} more entities"
         
-        if len(entities_data.get("entities", [])) > 10:
-            summary += f"\n... and {len(entities_data.get('entities', [])) - 10} more entities"
+        # Ensure valid JSON output
+        try:
+            json_output = json.dumps(entities_data, indent=2, ensure_ascii=True)
+        except Exception as e:
+            json_output = json.dumps({"error": f"JSON serialization failed: {str(e)}"})
         
-        return json.dumps(entities_data, indent=2), graph_image, summary
+        return json_output, graph_image, summary
             
     except Exception as e:
-        return f"Error: {str(e)}", None, "Analysis failed"
+        error_msg = f"Analysis failed: {str(e)}"
+        return json.dumps({"error": error_msg}), None, error_msg
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=knowledge_graph_builder,
-    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.Image(label="Knowledge Graph Visualization"),
-        gr.Markdown(label="Analysis Summary")
-    ],
-    title="🧠 AI Knowledge Graph Builder",
-    description="""
-    **Transform any text or webpage into an interactive knowledge graph!**
+# Create Gradio interface with error handling
+try:
+    demo = gr.Interface(
+        fn=knowledge_graph_builder,
+        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.Image(label="Knowledge Graph Visualization", type="pil"),
+            gr.Markdown(label="Analysis Summary")
+        ],
+        title="🧠 AI Knowledge Graph Builder",
+        description="""
+        **Transform any text or webpage into an interactive knowledge graph!**
+        
+        This tool:
+        1. 📖 Extracts content from URLs or analyzes your text
+        2. 🤖 Uses AI to identify entities and relationships  
+        3. 🕸️ Builds and visualizes knowledge graphs
+        4. 📊 Provides detailed analysis summaries
+        
+        **Examples to try:**
+        - News articles, Wikipedia pages, or any text content
+        """,
+        examples=[
+            ["Artificial intelligence companies like OpenAI, Google, and Microsoft are developing large language models for various applications."],
+            ["https://en.wikipedia.org/wiki/Machine_learning"],
+        ],
+        theme=gr.themes.Soft(),
+        allow_flagging="never"
+    )
     
-    This tool uses AI to:
-    1. 📖 Extract content from URLs or analyze your text
-    2. 🤖 Use Mistral AI to identify entities and relationships  
-    3. 🕸️ Build and visualize knowledge graphs
-    4. 📊 Provide detailed analysis summaries
+    if __name__ == "__main__":
+        demo.launch(mcp_server=True, share=False)
+        
+except Exception as e:
+    print(f"Failed to create Gradio interface: {e}")
+    # Create a simple fallback interface
+    def simple_demo(text):
+        return f"Error: {e}", None, "Application failed to start properly."
     
-    **Examples to try:**
-    - News articles: `https://www.bbc.com/news`
-    - Wikipedia pages: `https://en.wikipedia.org/wiki/Artificial_intelligence`
-    - Direct text: Copy and paste any article or document
-    """,
-    examples=[
-        ["https://en.wikipedia.org/wiki/Machine_learning"],
-        ["Artificial intelligence is transforming the world. Companies like OpenAI, Google, and Microsoft are leading the development of large language models. These models are being used in applications ranging from chatbots to code generation."],
-        ["https://www.nature.com/articles/d41586-023-00057-9"]
-    ],
-    theme=gr.themes.Soft()
-)
-
-demo.launch(mcp_server=True)
+    demo = gr.Interface(
+        fn=simple_demo,
+        inputs="text",
+        outputs=["text", "image", "text"],
+        title="Error - Knowledge Graph Builder"
+    )
+    
+    if __name__ == "__main__":
+        demo.launch(mcp_server=True, share=False)