Update app.py

app.py

@@ -1,5 +1,5 @@
 import threading
-from datetime import datetime
+import time
 import gradio as gr
 import logging
 import json
@@ -10,86 +10,109 @@ import os
 from pathlib import Path
 from typing import Dict, List, Tuple, Optional, Any, Union
 from dataclasses import dataclass, field
-from enum import Enum  # We're bringing in the big guns with Enum!
-
-# Our magical models and transformers
+from enum import Enum
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 from PIL import Image
+import black
 
+# Configure logging
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
-    datefmt='%Y-%m-%d %H:%M:%S',
-    handlers=[logging.StreamHandler(), logging.FileHandler('gradio_builder.log')]
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler('gradio_builder.log')
+    ]
 )
 logger = logging.getLogger(__name__)
 
-# Constants with a touch of mystery
-DEFAULT_PORT = 7860
+# Configuration
+@dataclass
+class Config:
+    port: int = 7860
+    debug: bool = False
+    share: bool = False
+    model_name: str = "gpt2"
+    embedding_model: str = "all-mpnet-base-v2"
+    theme: str = "default"
+    max_code_length: int = 5000
+    
+    @classmethod
+    def from_file(cls, path: str) -> "Config":
+        try:
+            with open(path) as f:
+                return cls(**json.load(f))
+        except Exception as e:
+            logger.warning(f"Failed to load config from {path}: {e}. Using defaults.")
+            return cls()
+
+# Constants
+CONFIG_PATH = Path("config.json")
 MODEL_CACHE_DIR = Path("model_cache")
 TEMPLATE_DIR = Path("templates")
 TEMP_DIR = Path("temp")
 DATABASE_PATH = Path("code_database.json")
 
-# Ensure our directories exist, like a well-organized wizard
+# Ensure directories exist
 for directory in [MODEL_CACHE_DIR, TEMPLATE_DIR, TEMP_DIR]:
-    directory.mkdir(parents=True, exist_ok=True)
+    directory.mkdir(exist_ok=True, parents=True)
 
 @dataclass
 class Template:
     code: str
     description: str
     components: List[str] = field(default_factory=list)
-
-    def __post_init__(self):
-        """
-        Init like a boss with some post-initialization magic.
-        """
-        self.components = self._extract_components()
+    created_at: str = field(default_factory=lambda: time.strftime("%Y-%m-%d %H:%M:%S"))
+    tags: List[str] = field(default_factory=list)
 
 class TemplateManager:
     def __init__(self, template_dir: Path):
         self.template_dir = template_dir
         self.templates: Dict[str, Template] = {}
-        self._load_templates()
 
-    def _load_templates(self):
-        """
-        Load templates with grace and elegance.
-        """
+    def load_templates(self) -> None:
         for file_path in self.template_dir.glob("*.json"):
             try:
                 with open(file_path, 'r') as f:
                     template_data = json.load(f)
-                    self.templates[template_data['description']] = Template(**template_data)
-            except (json.JSONDecodeError, KeyError) as e:
-                logger.error(f"Oh no! An error loading template from {file_path}: {e}")
+                    template = Template(**template_data)
+                    self.templates[template_data['description']] = template
+                logger.info(f"Loaded template: {file_path.stem}")
+            except Exception as e:
+                logger.error(f"Error loading template from {file_path}: {e}")
 
     def save_template(self, name: str, template: Template) -> bool:
-        """
-        Save a template with care and precision.
-        """
         file_path = self.template_dir / f"{name}.json"
         try:
             with open(file_path, 'w') as f:
                 json.dump(dataclasses.asdict(template), f, indent=2)
+            self.templates[name] = template
             return True
         except Exception as e:
-            logger.error(f"An unfortunate error saving template to {file_path}: {e}")
+            logger.error(f"Error saving template to {file_path}: {e}")
             return False
 
     def get_template(self, name: str) -> Optional[str]:
-        """
-        Retrieve a template with finesse.
-        """
-        return self.templates.get(name, {}).get('code', "")
+        template = self.templates.get(name)
+        return template.code if template else ""
+
+    def delete_template(self, name: str) -> bool:
+        file_path = self.template_dir / f"{name}.json"
+        try:
+            file_path.unlink()
+            self.templates.pop(name, None)
+            return True
+        except Exception as e:
+            logger.error(f"Error deleting template {name}: {e}")
+            return False
 
 class RAGSystem:
-    def __init__(self, model_name: str = "gpt3-incredible", device: str = "cuda" if torch.cuda.is_available() else "cpu", embedding_model="all-knowing-embedder"):
-        self.device = device
+    def __init__(self, config: Config):
+        self.config = config
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.embedding_model = None
         self.code_embeddings = None
         self.index = None
@@ -97,202 +120,183 @@ class RAGSystem:
         self.pipe = None
 
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=MODEL_CACHE_DIR)
-            self.model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=MODEL_CACHE_DIR).to(device)
-            self.pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device=self.device)
-            self.embedding_model = SentenceTransformer(embedding_model)
-            self._load_database()
-            logger.info("RAG system initialized with incredible power!")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                config.model_name, 
+                cache_dir=MODEL_CACHE_DIR
+            )
+            self.model = AutoModelForCausalLM.from_pretrained(
+                config.model_name, 
+                cache_dir=MODEL_CACHE_DIR
+            ).to(self.device)
+            self.pipe = pipeline(
+                "text-generation", 
+                model=self.model, 
+                tokenizer=self.tokenizer, 
+                device=self.device
+            )
+            self.embedding_model = SentenceTransformer(config.embedding_model)
+            self.load_database()
+            logger.info("RAG system initialized successfully.")
         except Exception as e:
-            logger.error(f"A dark force prevented loading the language model or embedding model: {e}. The placeholder generation shall be used.")
+            logger.error(f"Error initializing RAG system: {e}")
 
-    def _load_database(self):
-        """
-        Load the code database with ancient knowledge.
-        """
+    def load_database(self) -> None:
         if DATABASE_PATH.exists():
             try:
                 with open(DATABASE_PATH, 'r', encoding='utf-8') as f:
                     self.database = json.load(f)
                     self.code_embeddings = np.array(self.database['embeddings'])
-                logger.info("Ancient code database has been loaded.")
-                self._build_index()
-            except (json.JSONDecodeError, KeyError) as e:
-                logger.error(f"A curse has been cast upon the code database: {e}. A new database shall be created.")
-                self.database = {'codes': [], 'embeddings': []}
-                self.code_embeddings = np.array([])
+                logger.info(f"Loaded {len(self.database['codes'])} code snippets from database.")
                 self._build_index()
+            except Exception as e:
+                logger.error(f"Error loading database: {e}")
+                self._initialize_empty_database()
         else:
-            logger.info("No code database has been found. A new one shall be created.")
-            self.database = {'codes': [], 'embeddings': []}
-            self.code_embeddings = np.array([])
-            self._build_index()
-
-        if self.embedding_model and len(self.database['codes']) != len(self.database['embeddings']):
-            logger.warning("A mysterious mismatch between codes and embeddings has occurred. The embeddings shall be rebuilt.")
-            self.rebuild_embeddings()
-        elif self.embedding_model is None:
-            logger.warning("Embeddings are not supported in this realm. Proceed with caution.")
-
-    def _build_index(self):
-        """
-        Construct an index with magical efficiency.
-        """
+            logger.info("Creating new database.")
+            self._initialize_empty_database()
+
+    def _initialize_empty_database(self) -> None:
+        self.database = {'codes': [], 'embeddings': []}
+        self.code_embeddings = np.array([])
+        self._build_index()
+
+    def _build_index(self) -> None:
         if len(self.code_embeddings) > 0 and self.embedding_model:
-            self.index = faiss.IndexFlatL2(self.code_embeddings.shape[1])  # L2 distance, the measure of true similarity
+            dim = self.code_embeddings.shape[1]
+            self.index = faiss.IndexFlatL2(dim)
             self.index.add(self.code_embeddings)
+            logger.info(f"Built FAISS index with {len(self.code_embeddings)} vectors")
 
-    def add_to_database(self, code: str):
-        """
-        Add a code snippet to the database with care.
-        """
-        try:
-            if self.embedding_model is None:
-                raise ValueError("The embedding model has not been summoned.")
-            embedding = self.embedding_model.encode(code)
-            self.database['codes'].append(code)
-            self.database['embeddings'].append(embedding.tolist())
-            self.code_embeddings = np.vstack((self.code_embeddings, embedding)) if len(self.code_embeddings) > 0 else np.array([embedding])
-            self.index.add(np.array([embedding]))
-            self._save_database()
-            logger.info(f"A new code snippet has been added to the ancient database. Total size: {len(self.database['codes'])}.")
-        except Exception as e:
-            logger.error(f"A dark force prevented adding to the database: {e}")
+class GradioInterface:
+    def __init__(self, config: Config):
+        self.config = config
+        self.template_manager = TemplateManager(TEMPLATE_DIR)
+        self.template_manager.load_templates()
+        self.rag_system = RAGSystem(config)
 
-    def _save_database(self):
-        """
-        Save the database with eternal preservation.
-        """
+    def format_code(self, code: str) -> str:
         try:
-            with open(DATABASE_PATH, 'w', encoding='utf-8') as f:
-                json.dump(self.database, f, indent=2)
-            logger.info(f"The ancient database has been saved to {DATABASE_PATH}.")
+            return black.format_str(code, mode=black.FileMode())
         except Exception as e:
-            logger.error(f"A curse has been cast upon saving the database: {e}")
+            logger.warning(f"Code formatting failed: {e}")
+            return code
 
-    def rebuild_embeddings(self):
-        """
-        Rebuild embeddings with renewed power.
-        """
-        try:
-            if self.embedding_model is None:
-                raise ValueError("The embedding model has not been summoned.")
-            embeddings = self.embedding_model.encode(self.database['codes'])
-            self.code_embeddings = embeddings
-            self.database['embeddings'] = embeddings.tolist()
-            self._build_index()
-            self._save_database()
-            logger.info("The embeddings have been rebuilt and saved with enhanced power.")
-        except Exception as e:
-            logger.error(f"A dark force prevented rebuilding the embeddings: {e}")
-
-    def retrieve_similar_code(self, description: str, top_k: int = 3) -> List[str]:
-        """
-        Retrieve similar code with uncanny accuracy.
-        """
-        if self.embedding_model is None or self.index is None:
-            logger.warning("The embedding model or index is missing. Similar code retrieval is beyond our reach.")
-            return []
+    def _extract_components(self, code: str) -> List[str]:
+        components = []
         try:
-            embedding = self.embedding_model.encode(description)
-            D, I = self.index.search(np.array([embedding]), top_k)
-            logger.info(f"{top_k} similar code snippets have been retrieved for the description: {description}. Prepare to be amazed!")
-            return [self.database['codes'][i] for i in I[0]]
+            function_matches = re.findall(r'def (\w+)\(', code)
+            class_matches = re.findall(r'class (\w+):', code)
+            components.extend(function_matches)
+            components.extend(class_matches)
         except Exception as e:
-            logger.error(f"A dark force prevented retrieving similar code: {e}. The retrieval shall be attempted again.")
-            return []
-
-    def generate_code(self, description: str, template_code: str) -> str:
-        """
-        Generate code with incredible creativity.
-        """
-        retrieved_codes = self.retrieve_similar_code(description)
-        prompt = f"Description: {description} Retrieved Code Snippets: {''.join([f'```python {code} ```' for code in retrieved_codes])} Template: ```python {template_code} ``` Generated Code: ```python "
-        if self.pipe:
-            try:
-                generated_text = self.pipe(prompt, max_length=500, num_return_sequences=1)[0]['generated_text']
-                generated_code = generated_text.split("Generated Code:")[1].strip().split('```')[0]
-                logger.info("Incredible code has been generated!")
-                return generated_code
-            except Exception as e:
-                logger.error(f"A dark force prevented code generation with the language model: {e}. The template code shall be returned.")
-                return template_code
-        else:
-            logger.warning("The text generation pipeline is beyond our reach. A placeholder code shall be returned.")
-            return f"# Placeholder code generation. Description: {description} {template_code}"
+            logger.error(f"Error extracting components: {e}")
+        return list(set(components))
 
-class GradioInterface:
-    def __init__(self):
-        self.template_manager = TemplateManager(TEMPLATE_DIR)
-        self.rag_system = RAGSystem()
-        self.interface = self._build_interface()
-
-    def _extract_components(self, code: str) -> List[str]:
-        """
-        Extract components with precision and clarity.
-        """
-        components = []
-        function_matches = re.findall(r'def (\w+)\(', code)  # Parenthesis, the key to accuracy
-        components.extend(function_matches)
-        class_matches = re.findall(r'class (\w+)\:', code)  # Colon, the revealer of classes
-        components.extend(class_matches)
-        logger.info(f"Components have been extracted: {components}")
-        return components
-
-    def _get_template_choices(self) -> List[str]:
-        """
-        Present template choices with elegance.
-        """
-        return list(self.template_manager.templates.keys())
-
-    def _build_interface(self) -> gr.Blocks:
-        """
-        Construct the Gradio interface with style and functionality.
-        """
-        with gr.Blocks() as interface:
-            gr.Markdown("## Code Generation Interface")
-            description_input = gr.Textbox(label="Description", placeholder="Enter a description for the code you wish to bring to life.")
-            code_output = gr.Textbox(label="Generated Code", interactive=False)
-            generate_button = gr.Button("Generate Code")
-            template_choice = gr.Dropdown(label="Select Template", choices=self._get_template_choices(), value=None)
-            save_button = gr.Button("Save as Template")
-            status_output = gr.Textbox(label="Status", interactive=False)
-
-            def generate_code_wrapper(description, template_choice):
+    def launch(self) -> None:
+        with gr.Blocks(theme=gr.themes.Base()) as interface:
+            # Custom CSS
+            gr.Markdown(
                 """
-                Generate code with a simple button click.
+                <style>
+                .header {
+                    text-align: center;
+                    background-color: #f0f0f0;
+                    padding: 20px;
+                    border-radius: 10px;
+                    margin-bottom: 20px;
+                }
+                .container {
+                    max-width: 1200px;
+                    margin: 0 auto;
+                }
+                </style>
+                <div class="header">
+                <h1>Code Generation Interface</h1>
+                <p>Generate and manage code templates easily</p>
+                </div>
                 """
+            )
+
+            with gr.Row():
+                with gr.Column(scale=2):
+                    description_input = gr.Textbox(
+                        label="Description",
+                        placeholder="Enter a description for the code you want to generate",
+                        lines=3
+                    )
+                    template_choice = gr.Dropdown(
+                        label="Select Template",
+                        choices=list(self.template_manager.templates.keys()),
+                        value=None
+                    )
+
+            with gr.Row():
+                generate_button = gr.Button("Generate Code", variant="primary")
+                save_button = gr.Button("Save as Template", variant="secondary")
+                clear_button = gr.Button("Clear", variant="stop")
+
+            with gr.Row():
+                code_output = gr.Code(
+                    label="Generated Code",
+                    language="python",
+                    interactive=True
+                )
+
+            status_output = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
+
+            def generate_code_wrapper(description: str, template_choice: str) -> Tuple[str, str]:
+                if not description.strip():
+                    return "", "Please provide a description"
+                
                 try:
                     template_code = self.template_manager.get_template(template_choice) if template_choice else ""
-                    generated_code, status = self.rag_system.generate_code(description, template_code)
-                    return generated_code, status
+                    generated_code = self.rag_system.generate_code(description, template_code)
+                    formatted_code = self.format_code(generated_code)
+                    
+                    if not formatted_code:
+                        return "", "Failed to generate code. Please try again."
+                        
+                    return formatted_code, "Code generated successfully."
+                    
                 except Exception as e:
-                    return "", f"A dark force prevented code generation: {e}"
+                    logger.error(f"Error in code generation: {str(e)}")
+                    return "", f"Error: {str(e)}"
 
-            def save_template_wrapper(code, name, description):
-                """
-                Save a template with ease and security.
-                """
+            def save_template_wrapper(code: str, name: str, description: str) -> Tuple[str, str]:
                 try:
-                    if not name:
-                        return code, "A template name must be provided to seal its destiny."
-                    if not code:
-                        return code, "Code cannot be empty. It must be filled with potential."
+                    if not name or not code:
+                        return code, "Template name and code are required."
 
                     components = self._extract_components(code)
-                    template = Template(code=code, description=name, components=components)
+                    template = Template(
+                        code=code,
+                        description=name,
+                        components=components,
+                        tags=[t.strip() for t in description.split(',') if t.strip()]
+                    )
+
                     if self.template_manager.save_template(name, template):
                         self.rag_system.add_to_database(code)
-                        return code, f"Template '{name}' has been saved for eternity."
+                        template_choice.choices = list(self.template_manager.templates.keys())
+                        return code, f"Template '{name}' saved successfully."
                     else:
-                        return code, "A mysterious force prevented saving the template."
+                        return code, "Failed to save template."
                 except Exception as e:
-                    return code, f"An error occurred while saving the template: {e}"
+                    return code, f"Error saving template: {e}"
 
+            def clear_outputs() -> Tuple[str, str, str]:
+                return "", "", ""
+
+            # Event handlers
             generate_button.click(
                 fn=generate_code_wrapper,
                 inputs=[description_input, template_choice],
-                outputs=[code_output, status_output]
+                outputs=[code_output, status_output],
+                api_name="generate_code",
+                show_progress=True
             )
 
             save_button.click(
@@ -301,38 +305,30 @@ class GradioInterface:
                 outputs=[code_output, status_output]
             )
 
-        logger.info("The Gradio interface is ready to be unveiled.")
-        return interface
-
-    def launch(self, **kwargs):
-        """
-        Launch the Gradio interface with a flourish.
-        """
-        logger.info("=== Application Startup ===")
-        try:
-            self.interface.launch(
-                server_port=DEFAULT_PORT,
-                share=False,
-                debug=True,
-                **kwargs
+            clear_button.click(
+                fn=clear_outputs,
+                inputs=[],
+                outputs=[description_input, code_output, status_output]
             )
-        except Exception as e:
-            logger.error(f"An unexpected error has occurred: {e}. The application shall be shut down.")
-            raise
-        finally:
-            logger.info("=== Application Shutdown ===")
+
+        # Launch the interface
+        interface.launch(
+            server_port=self.config.port,
+            share=self.config.share,
+            debug=self.config.debug
+        )
 
 def main():
-    """
-    The main function, where the magic begins.
-    """
-    logger.info("=== Application Initiation ===")
+    logger.info("=== Application Startup ===")
     try:
-        interface = GradioInterface()
+        config = Config.from_file(CONFIG_PATH)
+        interface = GradioInterface(config)
         interface.launch()
     except Exception as e:
-        logger.error(f"A critical error has occurred: {e}. The application shall be terminated.")
+        logger.error(f"Application error: {e}")
         raise
+    finally:
+        logger.info("=== Application Shutdown ===")
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()