app.py

import gradio as gr
import os
import json
import requests
from bs4 import BeautifulSoup
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import io
import base64
from huggingface_hub import InferenceClient
import re
from urllib.parse import urlparse

def fetch_content(url_or_text):
    """Fetch content from URL or return text directly.
    
    Args:
        url_or_text: Either a URL to fetch content from, or direct text input
        
    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

def extract_entities(text):
    """Extract entities and relationships using Mistral.
    
    Args:
        text: Input text to analyze
        
    Returns:
        Dictionary containing entities and relationships
    """
    try:
        client = InferenceClient(
            provider="together",
            api_key=os.environ.get("HF_TOKEN"),
        )
        
        prompt = f"""
        Analyze the following text and extract:
        1. Named entities (people, organizations, locations, concepts)
        2. Relationships between these entities
        
        Return the result as a JSON object with this structure:
        {{
            "entities": [
                {{"name": "entity_name", "type": "PERSON|ORG|LOCATION|CONCEPT", "description": "brief description"}}
            ],
            "relationships": [
                {{"source": "entity1", "target": "entity2", "relation": "relationship_type", "description": "brief description"}}
            ]
        }}
        
        Text to analyze:
        {text[:2000]}
        
        JSON:"""
        
        completion = client.chat.completions.create(
            model="mistralai/Mistral-Small-24B-Instruct-2501",
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            max_tokens=1500,
            temperature=0.3,
        )
        
        response_text = completion.choices[0].message.content
        
        # Extract JSON from response
        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"})
            
            return {"entities": entities, "relationships": []}
            
    except Exception as e:
        return {"entities": [{"name": "Error", "type": "ERROR", "description": str(e)}], "relationships": []}

def build_knowledge_graph(entities_data):
    """Build and visualize knowledge graph.
    
    Args:
        entities_data: Dictionary containing entities and relationships
        
    Returns:
        PIL Image object of the knowledge graph
    """
    try:
        # Create networkx graph
        G = nx.Graph()
        
        # 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", ""))
        
        # 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"], 
                          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")
        
        # 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))
        
        # 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,))
        
        from PIL import Image
        pil_image = Image.fromarray(img_array)
        plt.close(fig)
        
        return pil_image
        
    except Exception as e:
        # Create error image
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.text(0.5, 0.5, f"Error creating graph:\n{str(e)}", 
                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,))
        
        from PIL import Image
        pil_image = Image.fromarray(img_array)
        plt.close(fig)
        
        return pil_image

def knowledge_graph_builder(url_or_text):
    """Main function to build knowledge graph from URL or text.
    
    Args:
        url_or_text: URL to analyze or direct text input
        
    Returns:
        Tuple of (entities_json, graph_image, summary)
    """
    try:
        # Step 1: Fetch content
        content = fetch_content(url_or_text)
        
        if content.startswith("Error"):
            return content, None, "Failed to fetch content"
        
        # Step 2: Extract entities
        entities_data = extract_entities(content)
        
        # Step 3: Build knowledge graph
        graph_image = build_knowledge_graph(entities_data)
        
        # Step 4: Create summary
        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
        
        🔍 **Extracted Entities:**
        """
        
        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", [])) > 10:
            summary += f"\n... and {len(entities_data.get('entities', [])) - 10} more entities"
        
        return json.dumps(entities_data, indent=2), graph_image, summary
            
    except Exception as e:
        return f"Error: {str(e)}", None, "Analysis failed"

# 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!**
    
    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
    
    **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)