Update app.py

app.py

@@ -6,10 +6,7 @@ import json
 import re
 import torch
 import tempfile
-import subprocess
-import ast
 import os
-import dataclasses
 from pathlib import Path
 from typing import Dict, List, Tuple, Optional, Any, Union
 from dataclasses import dataclass, field
@@ -19,7 +16,6 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 from PIL import Image
-from templates import TemplateManager, Template  # Import TemplateManager and Template
 
 # Configure logging
 logging.basicConfig(
@@ -37,72 +33,112 @@ DEFAULT_PORT = 7860
 MODEL_CACHE_DIR = Path("model_cache")
 TEMPLATE_DIR = Path("templates")
 TEMP_DIR = Path("temp")
-DATABASE_PATH = Path("code_database.json")  # Path for our simple database
+DATABASE_PATH = Path("code_database.json")
 
 # Ensure directories exist
 for directory in [MODEL_CACHE_DIR, TEMPLATE_DIR, TEMP_DIR]:
     directory.mkdir(exist_ok=True, parents=True)
 
+@dataclass
+class Template:
+    code: str
+    description: str
+    components: List[str] = field(default_factory=list)
+
+class TemplateManager:
+    def __init__(self, template_dir: Path):
+        self.template_dir = template_dir
+        self.templates: Dict[str, Template] = {}
+
+    def load_templates(self):
+        for file_path in self.template_dir.glob("*.json"):
+            try:
+                with open(file_path, 'r') as f:
+                    template_data = json.load(f)
+                    template = Template(**template_data)
+                    self.templates[template_data['description']] = template
+            except json.JSONDecodeError as e:
+                logger.error(f"Error loading template from {file_path}: {e}")
+            except KeyError as e:
+                logger.error(f"Missing key in template file {file_path}: {e}")
+
+    def save_template(self, name: str, template: Template) -> bool:
+        file_path = self.template_dir / f"{name}.json"
+        try:
+            with open(file_path, 'w') as f:
+                json.dump(dataclasses.asdict(template), f, indent=2)
+            return True
+        except Exception as e:
+            logger.error(f"Error saving template to {file_path}: {e}")
+            return False
+
+    def get_template(self, name: str) -> Optional[str]:
+        return self.templates.get(name, {}).get('code', "")
+
 class RAGSystem:
     def __init__(self, model_name: str = "gpt2", device: str = "cuda" if torch.cuda.is_available() else "cpu", embedding_model="all-mpnet-base-v2"):
+        self.device = device
+        self.embedding_model = None
+        self.code_embeddings = None
+        self.index = None
+        self.database = {'codes': [], 'embeddings': []}
+        self.pipe = None
+
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-            self.model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-            self.device = device
+            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 successfully.")
         except Exception as e:
             logger.error(f"Error loading language model or embedding model: {e}. Falling back to placeholder generation.")
-            self.pipe = None
-            self.embedding_model = None
-            self.code_embeddings = None
 
     def load_database(self):
-        """Loads or creates the code database"""
         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("Loaded code database from file.")
+                self._build_index()
             except (json.JSONDecodeError, KeyError) as e:
                 logger.error(f"Error loading code database: {e}. Creating new database.")
                 self.database = {'codes': [], 'embeddings': []}
                 self.code_embeddings = np.array([])
-
+                self._build_index()
         else:
             logger.info("Code database does not exist. Creating new database.")
             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("Mismatch between number of codes and embeddings, rebuilding embeddings.")
             self.rebuild_embeddings()
         elif self.embedding_model is None:
-            logger.warning("Embeddings are not supported in this context.")
+            logger.warning ("Embeddings are not supported in this context.")
 
-        # Index the embeddings for efficient searching
+    def _build_index(self):
         if len(self.code_embeddings) > 0 and self.embedding_model:
             self.index = faiss.IndexFlatL2(self.code_embeddings.shape[1])  # L2 distance
             self.index.add(self.code_embeddings)
 
     def add_to_database(self, code: str):
-        """Adds a code snippet to the database"""
         try:
+            if self.embedding_model is None:
+                raise ValueError("Embedding model not loaded.")
             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))
-            self.index.add(np.array([embedding]))  # update FAISS index
+            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"Added code snippet to database. Total size: {len(self.database['codes'])}.")
         except Exception as e:
             logger.error(f"Error adding to database: {e}")
 
     def save_database(self):
-        """Saves the database to a file"""
         try:
             with open(DATABASE_PATH, 'w', encoding='utf-8') as f:
                 json.dump(self.database, f, indent=2)
@@ -111,22 +147,21 @@ class RAGSystem:
             logger.error(f"Error saving database: {e}")
 
     def rebuild_embeddings(self):
-        """Rebuilds embeddings from the codes"""
         try:
+            if self.embedding_model is None:
+                raise ValueError("Embedding model not loaded.")
             embeddings = self.embedding_model.encode(self.database['codes'])
             self.code_embeddings = embeddings
             self.database['embeddings'] = embeddings.tolist()
-            self.index = faiss.IndexFlatL2(embeddings.shape[1])  # L2 distance
-            self.index.add(embeddings)
+            self._build_index()
             self.save_database()
             logger.info("Rebuilt and saved embeddings to the database.")
         except Exception as e:
             logger.error(f"Error rebuilding embeddings: {e}")
 
     def retrieve_similar_code(self, description: str, top_k: int = 3) -> List[str]:
-        """Retrieves similar code snippets from the database"""
-        if self.embedding_model is None:
-            logger.warning("Embedding model is not available. Cannot retrieve similar code.")
+        if self.embedding_model is None or self.index is None:
+            logger.warning("Embedding model or index not available. Cannot retrieve similar code.")
             return []
         try:
             embedding = self.embedding_model.encode(description)
@@ -139,7 +174,7 @@ class RAGSystem:
 
     def generate_code(self, description: str, template_code: str) -> str:
         retrieved_codes = self.retrieve_similar_code(description)
-        prompt = f"Description: {description}\nRetrieved Code Snippets:\n{''.join([f'```python\n{code}\n```\n' for code in retrieved_codes])}\nTemplate:\n```python\n{template_code}\n```\nGenerated Code:\n```python\n"
+        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']
@@ -151,58 +186,25 @@ class RAGSystem:
                 return template_code
         else:
             logger.warning("Text generation pipeline is not available. Returning placeholder code.")
-            return f"# Placeholder code generation. Description: {description}\n{template_code}"
-
-    def generate_interface(self, screenshot: Optional[Image.Image], description: str) -> str:
-        retrieved_codes = self.retrieve_similar_code(description)
-        prompt = f"Create a Gradio interface based on this description: {description}\nRetrieved Code Snippets:\n{''.join([f'```python\n{code}\n```\n' for code in retrieved_codes])}"
-        if screenshot:
-            prompt += "\nThe interface should resemble the provided screenshot."
-        prompt += "\n```python\n"
-        if self.pipe:
-            try:
-                generated_text = self.pipe(prompt, max_length=500, num_return_sequences=1)[0]['generated_text']
-                generated_code = generated_text.split("```")[1].strip()
-                logger.info("Interface code generated successfully.")
-                return generated_code
-            except Exception as e:
-                logger.error(f"Error generating interface with language model: {e}. Returning placeholder.")
-                return "import gradio as gr\n\ndemo = gr.Interface(fn=lambda x:x, inputs='text', outputs='text')\ndemo.launch()"
-        else:
-            logger.warning("Text generation pipeline is not available. Returning placeholder interface code.")
-            return "import gradio as gr\n\ndemo = gr.Interface(fn=lambda x:x, inputs='text', outputs='text')\ndemo.launch()"
-
-class PreviewManager:
-    def __init__(self):
-        self.preview_code = ""
-
-    def update_preview(self, code: str):
-        """Update the preview with the generated code."""
-        self.preview_code = code
-        logger.info("Preview updated with new code.")
+            return f"# Placeholder code generation. Description: {description} {template_code}"
 
 class GradioInterface:
     def __init__(self):
         self.template_manager = TemplateManager(TEMPLATE_DIR)
         self.template_manager.load_templates()
-        self.current_code = ""
         self.rag_system = RAGSystem()
-        self.preview_manager = PreviewManager()
 
     def _extract_components(self, code: str) -> List[str]:
-        """Extract components from the code."""
         components = []
-        function_matches = re.findall(r'def (\w+)', code components.extend(function_matches)
-        class_matches = re.findall(r'class (\w+)', code)
+        function_matches = re.findall(r'def (\w+)\(', code)  # added parenthesis for more accuracy
+        components.extend(function_matches)
+        class_matches = re.findall(r'class (\w+)\:', code)  # added colon for more accuracy
         components.extend(class_matches)
         logger.info(f"Extracted components: {components}")
         return components
 
     def _get_template_choices(self) -> List[str]:
-        """Get available template choices."""
-        choices = list(self.template_manager.templates.keys())
-        logger.info(f"Available template choices: {choices}")
-        return choices
+        return list(self.template_manager.templates.keys())
 
     def launch(self, **kwargs):
         with gr.Blocks() as interface:
@@ -212,59 +214,51 @@ class GradioInterface:
             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):
+                try:
+                    template_code = self.template_manager.get_template(template_choice) if template_choice else ""
+                    generated_code = self.rag_system.generate_code(description, template_code)
+                    return generated_code, "Code generated successfully."
+                except Exception as e:
+                    return "", f"Error generating code: {e}"
+
+            def save_template_wrapper(code, name, description):
+                try:
+                    if not name:
+                        return code, "Template name cannot be empty."
+                    if not code:
+                        return code, "Code cannot be empty."
+
+                    components = self._extract_components(code)
+                    template = Template(code=code, description=name, components=components)
+                    if self.template_manager.save_template(name, template):
+                        self.rag_system.add_to_database(code)
+                        return code, f"Template '{name}' saved successfully."
+                    else:
+                        return code, "Failed to save template."
+                except Exception as e:
+                    return code, f"Error saving template: {e}"
 
             generate_button.click(
-                fn=self.generate_code,
+                fn=generate_code_wrapper,
                 inputs=[description_input, template_choice],
-                outputs=code_output
+                outputs=[code_output, status_output]
             )
 
             save_button.click(
-                fn=self.save_template,
+                fn=save_template_wrapper,
                 inputs=[code_output, template_choice, description_input],
-                outputs=code_output
-            )
-
-            gr.Markdown("### Preview")
-            preview_output = gr.Textbox(label="Preview", interactive=False)
-            self.preview_manager.update_preview(code_output)
-
-            generate_button.click(
-                fn=lambda code: self.preview_manager.update_preview(code),
-                inputs=code_output,
-                outputs=preview_output
+                outputs=[code_output, status_output]
             )
 
         logger.info("Launching Gradio interface.")
         interface.launch(**kwargs)
 
-    def generate_code(self, description: str, template_choice: Optional[str]) -> str:
-        """Generate code based on the description and selected template."""
-        template_code = self.template_manager.get_template(template_choice) if template_choice else ""  # Get template code if selected
-        logger.info(f"Generating code for description: {description} with template: {template_choice}")
-        return self.rag_system.generate_code(description, template_code)
-
-    def save_template(self, code: str, name: str, description: str) -> str:
-        """Save the generated code as a template."""
-        try:
-            components = self._extract_components(code)
-            template = Template(code=code, description=description, components=components)
-            if self.template_manager.save_template(name, template):
-                self.rag_system.add_to_database(code)  # Add code to the database
-                logger.info(f"Template '{name}' saved successfully.")
-                return f"✅ Template '{name}' saved successfully."
-            else:
-                logger.error("Failed to save template.")
-                return "❌ Failed to save template."
-        except Exception as e:
-            logger.error(f"Error saving template: {e}")
-            return f"❌ Error saving template: {str(e)}"
-
 def main():
     logger.info("=== Application Startup ===")
-    
     try:
-        # Initialize and launch interface
         interface = GradioInterface()
         interface.launch(
             server_port=DEFAULT_PORT,