Create app.py

app.py

@@ -0,0 +1,398 @@
+import gradio as gr
+import torch
+import logging
+import yaml
+import os
+import json
+import jwt
+import redis
+import sqlite3
+from datetime import datetime, timedelta
+from pathlib import Path
+from transformers import AutoTokenizer, T5ForConditionalGeneration
+import networkx as nx
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+import io
+import traceback
+from typing import Tuple, Optional, Dict, List, Union
+import colorama
+from colorama import Fore, Style
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+import plotly.graph_objects as go
+import hashlib
+import asyncio
+import aiohttp
+from fastapi import FastAPI, HTTPException, Depends, status
+from fastapi.security import OAuth2PasswordBearer
+from pydantic import BaseModel, EmailStr
+import uvicorn
+
+# Advanced Configuration Models
+@dataclass
+class ModelConfig:
+    name: str
+    max_length: int
+    num_beams: int
+    temperature: float
+    top_k: int
+    top_p: float
+
+@dataclass
+class StyleConfig:
+    node_color: str
+    edge_color: str
+    node_size: int
+    font_size: int
+    layout: str
+
+@dataclass
+class OutputConfig:
+    width: int
+    height: int
+    dpi: int
+    format: str
+    quality: int
+
+# Abstract Base Classes for Extensibility
+class DiagramStrategy(ABC):
+    @abstractmethod
+    def create_diagram(self, components: List[str], style: StyleConfig) -> Image.Image:
+        pass
+
+class NetworkDiagram(DiagramStrategy):
+    def create_diagram(self, components: List[str], style: StyleConfig) -> Image.Image:
+        G = nx.DiGraph()
+        for i in range(len(components)-1):
+            G.add_edge(components[i], components[i+1])
+
+        plt.figure(figsize=(12, 8))
+        
+        if style.layout == "spring":
+            pos = nx.spring_layout(G)
+        elif style.layout == "circular":
+            pos = nx.circular_layout(G)
+        else:
+            pos = nx.kamada_kawai_layout(G)
+
+        nx.draw_networkx_nodes(G, pos,
+                             node_color=style.node_color,
+                             node_size=style.node_size)
+        nx.draw_networkx_edges(G, pos,
+                             edge_color=style.edge_color,
+                             arrows=True)
+        nx.draw_networkx_labels(G, pos,
+                              font_size=style.font_size)
+
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', dpi=300)
+        plt.close()
+        buf.seek(0)
+        return Image.open(buf)
+
+class PlotlyDiagram(DiagramStrategy):
+    def create_diagram(self, components: List[str], style: StyleConfig) -> Image.Image:
+        G = nx.DiGraph()
+        for i in range(len(components)-1):
+            G.add_edge(components[i], components[i+1])
+
+        pos = nx.spring_layout(G)
+        
+        edge_x = []
+        edge_y = []
+        for edge in G.edges():
+            x0, y0 = pos[edge[0]]
+            x1, y1 = pos[edge[1]]
+            edge_x.extend([x0, x1, None])
+            edge_y.extend([y0, y1, None])
+
+        node_x = [pos[node][0] for node in G.nodes()]
+        node_y = [pos[node][1] for node in G.nodes()]
+
+        fig = go.Figure()
+        fig.add_trace(go.Scatter(x=edge_x, y=edge_y,
+                                line=dict(width=0.5, color=style.edge_color),
+                                hoverinfo='none',
+                                mode='lines'))
+        
+        fig.add_trace(go.Scatter(x=node_x, y=node_y,
+                                mode='markers+text',
+                                marker=dict(size=style.node_size/100,
+                                          color=style.node_color),
+                                text=list(G.nodes()),
+                                textposition="bottom center"))
+
+        fig.update_layout(showlegend=False,
+                         hovermode='closest',
+                         margin=dict(b=0,l=0,r=0,t=0))
+
+        img_bytes = fig.to_image(format="png")  # Requires kaleido
+        return Image.open(io.BytesIO(img_bytes))
+
+# Database Manager
+class DatabaseManager:
+    def __init__(self, db_path: str = "diagrams.db"):
+        self.conn = sqlite3.connect(db_path)
+        self.create_tables()
+
+    def create_tables(self):
+        cursor = self.conn.cursor()
+        cursor.execute('''
+        CREATE TABLE IF NOT EXISTS users (
+            id INTEGER PRIMARY KEY,
+            username TEXT UNIQUE,
+            email TEXT UNIQUE,
+            password_hash TEXT,
+            created_at TIMESTAMP
+        )
+        ''')
+        
+        cursor.execute('''
+        CREATE TABLE IF NOT EXISTS diagrams (
+            id INTEGER PRIMARY KEY,
+            user_id INTEGER,
+            title TEXT,
+            description TEXT,
+            created_at TIMESTAMP,
+            image_path TEXT,
+            FOREIGN KEY (user_id) REFERENCES users (id)
+        )
+        ''')
+        self.conn.commit()
+
+# Cache Manager using Redis
+class CacheManager:
+    def __init__(self, redis_url: str = "redis://localhost"):
+        self.redis_client = redis.from_url(redis_url)
+
+    def get_cached_diagram(self, key: str) -> Optional[bytes]:
+        return self.redis_client.get(key)
+
+    def cache_diagram(self, key: str, diagram: bytes, expire: int = 3600):
+        self.redis_client.set(key, diagram, ex=expire)
+
+# Advanced Diagram Generator
+class AdvancedDiagramGenerator:
+    def __init__(self, config_path: str = "config.yaml"):
+        # Initialize logging first
+        self.setup_logging()
+        
+        # Load configuration
+        self.load_config(config_path)
+        
+        # Setup components (tokenizer, model, etc.)
+        self.setup_components()
+        
+        # Initialize diagram strategies
+        self.strategies = {
+            "network": NetworkDiagram(),
+            "plotly": PlotlyDiagram()
+        }
+
+    def load_config(self, config_path: str):
+        try:
+            if not os.path.exists(config_path):
+                self.logger.warning(f"Config file not found at {config_path}. Using default configuration.")
+                # Default configuration
+                config_data = {
+                    'model': {
+                        'name': 't5-small',
+                        'max_length': 512,
+                        'num_beams': 4,
+                        'temperature': 1.0,
+                        'top_k': 50,
+                        'top_p': 0.9
+                    },
+                    'styles': {
+                        'network': {
+                            'node_color': '#1f77b4',
+                            'edge_color': '#7f7f7f',
+                            'node_size': 3000,
+                            'font_size': 12,
+                            'layout': 'spring'
+                        }
+                    },
+                    'output': {
+                        'width': 1200,
+                        'height': 800,
+                        'dpi': 300,
+                        'format': 'png',
+                        'quality': 95
+                    }
+                }
+            else:
+                with open(config_path) as f:
+                    config_data = yaml.safe_load(f)
+
+            # Ensure all required sections exist with defaults
+            if 'model' not in config_data:
+                config_data['model'] = {}
+            
+            # Set default values for model configuration
+            model_config_data = config_data['model']
+            model_config_data.setdefault('name', 't5-small')
+            model_config_data.setdefault('max_length', 512)
+            model_config_data.setdefault('num_beams', 4)
+            model_config_data.setdefault('temperature', 1.0)
+            model_config_data.setdefault('top_k', 50)
+            model_config_data.setdefault('top_p', 0.9)
+            
+            # Create ModelConfig instance
+            self.model_config = ModelConfig(**model_config_data)
+
+            # Handle styles configuration
+            if 'styles' not in config_data:
+                config_data['styles'] = {
+                    'network': {
+                        'node_color': '#1f77b4',
+                        'edge_color': '#7f7f7f',
+                        'node_size': 3000,
+                        'font_size': 12,
+                        'layout': 'spring'
+                    }
+                }
+            
+            # Create StyleConfig instances
+            self.style_configs = {}
+            for style_name, style_data in config_data['styles'].items():
+                style_data.setdefault('node_color', '#1f77b4')
+                style_data.setdefault('edge_color', '#7f7f7f')
+                style_data.setdefault('node_size', 3000)
+                style_data.setdefault('font_size', 12)
+                style_data.setdefault('layout', 'spring')
+                self.style_configs[style_name] = StyleConfig(**style_data)
+
+            # Handle output configuration
+            if 'output' not in config_data:
+                config_data['output'] = {}
+            
+            output_config_data = config_data['output']
+            output_config_data.setdefault('width', 1200)
+            output_config_data.setdefault('height', 800)
+            output_config_data.setdefault('dpi', 300)
+            output_config_data.setdefault('format', 'png')
+            output_config_data.setdefault('quality', 95)
+            
+            # Create OutputConfig instance
+            self.output_config = OutputConfig(**output_config_data)
+
+            self.config = config_data
+            self.logger.info("Configuration loaded successfully")
+
+        except Exception as e:
+            self.logger.error(f"Error loading configuration: {str(e)}")
+            raise RuntimeError(f"Failed to load configuration: {str(e)}")
+
+    def setup_components(self):
+        # Initialize tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_config.name)
+        self.model = T5ForConditionalGeneration.from_pretrained(self.model_config.name)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+    def setup_logging(self):
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(__name__)
+
+    async def extract_components(self, text: str) -> List[str]:
+        inputs = self.tokenizer(
+            f"convert to diagram: {text}",
+            return_tensors="pt",
+            max_length=self.model_config.max_length,
+            truncation=True
+        ).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                inputs.input_ids,
+                max_length=150,
+                num_beams=self.model_config.num_beams,
+                temperature=self.model_config.temperature,
+                top_k=self.model_config.top_k,
+                top_p=self.model_config.top_p
+            )
+
+        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        components = [comp.strip() for comp in decoded.replace('->', ',').split(',')]
+        return [comp for comp in components if comp]
+
+    def save_to_database(self, user_id: int, description: str, diagram: Image.Image):
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        image_path = f"diagrams/user_{user_id}/{timestamp}.png"
+        os.makedirs(os.path.dirname(image_path), exist_ok=True)
+        diagram.save(image_path)
+
+        cursor = self.db.conn.cursor()
+        cursor.execute('''
+        INSERT INTO diagrams (user_id, description, created_at, image_path)
+        VALUES (?, ?, ?, ?)
+        ''', (user_id, description, datetime.now(), image_path))
+        self.db.conn.commit()
+
+    async def generate_diagram(self, text: str, style: str, strategy: str) -> Tuple[Optional[Image.Image], str]:
+        try:
+            components = await self.extract_components(text)
+            diagram = self.strategies[strategy].create_diagram(components, self.style_configs[style])
+            return diagram, "Diagram generated successfully!"
+        except Exception as e:
+            self.logger.error(f"Error generating diagram: {str(e)}")
+            return None, f"Error: {str(e)}"
+
+# FastAPI Integration
+app = FastAPI()
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+class DiagramRequest(BaseModel):
+    text: str
+    style: str = "network"
+    strategy: str = "network"
+
+# Initialize the generator
+generator = AdvancedDiagramGenerator()
+
+# Gradio Interface
+def create_gradio_interface():
+    def generate(text: str, style: str, strategy: str) -> Tuple[Optional[Image.Image], str]:
+        return asyncio.run(generator.generate_diagram(text, style, strategy))
+
+    iface = gr.Interface(
+        fn=generate,
+        inputs=[
+            gr.Textbox(
+                label="Enter your diagram description",
+                placeholder="e.g., 'Create a flowchart for software development lifecycle'",
+                lines=3
+            ),
+            gr.Dropdown(
+                choices=list(generator.style_configs.keys()),
+                label="Diagram Style",
+                value="network"
+            ),
+            gr.Dropdown(
+                choices=list(generator.strategies.keys()),
+                label="Visualization Strategy",
+                value="network"
+            )
+        ],
+        outputs=[
+            gr.Image(label="Generated Diagram", type="pil"),
+            gr.Textbox(label="Status")
+        ],
+        title="Advanced Enterprise Diagram Generator",
+        description="""
+        Enterprise-grade diagram generation tool with advanced features:
+        - Multiple visualization strategies
+        - Caching system
+        - Database storage
+        - Multiple output formats
+        - Custom styling options
+        """,
+        theme=gr.themes.Glass()
+    )
+    return iface
+
+if __name__ == "__main__":
+    iface = create_gradio_interface()
+    iface.launch(share=True)