Update app.py

app.py

@@ -8,6 +8,16 @@ from typing import List
 import yaml
 import os
 import json
+import logging
+from functools import lru_cache
+import time
+import pygments
+from pygments.lexers import get_lexer_by_name
+from pygments.formatters import HtmlFormatter
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 def load_config():
     config_path = os.path.join(os.path.dirname(__file__), "config.yaml")
@@ -15,7 +25,7 @@ def load_config():
         with open(config_path, "r") as config_file:
             return yaml.safe_load(config_file)
     except FileNotFoundError:
-        print("Config file not found. Using default configuration.")
+        logger.warning("Config file not found. Using default configuration.")
         return {
             "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
             "hf_api_token": os.environ.get("HUGGINGFACEHUB_API_TOKEN", "your_default_token_here")
@@ -30,8 +40,9 @@ class CodeGenerationTool(BaseTool):
     name = "CodeGeneration"
     description = "Generates code based on a prompt"
 
+    @lru_cache(maxsize=100)
     def _run(self, prompt: str) -> str:
-        # Implement basic code generation logic
+        logger.info(f"Generating code for prompt: {prompt}")
         if "Flask app structure" in prompt:
             return self.generate_flask_app_structure()
         elif "binary search algorithm" in prompt:
@@ -193,22 +204,23 @@ class Agent:
     def __init__(self, name: str, description: str, tools: List[BaseTool]):
         self.name = name
         self.description = description
-        self.llm = HuggingFaceEndpoint(
-            repo_id=config["model"],
-            task="text-generation",
-            model_kwargs={"temperature": 0.7, "max_length": 1024},
-            huggingfacehub_api_token=config["hf_api_token"]
-        )
-        
-        self.prompt_template = PromptTemplate(
-            template="You are {name}, {description}. Respond to the following: {input}",
-            input_variables=["name", "description", "input"]
-        )
-        
         try:
+            self.llm = HuggingFaceEndpoint(
+                repo_id=config["model"],
+                task="text-generation",
+                model_kwargs={"temperature": 0.7, "max_length": 1024},
+                huggingfacehub_api_token=config["hf_api_token"]
+            )
+            
+            self.prompt_template = PromptTemplate(
+                template="You are {name}, {description}. Respond to the following: {input}",
+                input_variables=["name", "description", "input"]
+            )
+            
             self.agent = create_react_agent(self.llm, tools, self.prompt_template)
             self.agent_executor = AgentExecutor(agent=self.agent, tools=tools, verbose=True)
         except Exception as e:
+            logger.error(f"Failed to initialize agent: {e}")
             raise AgentInitializationError(f"Failed to initialize agent: {e}")
 
     async def run(self, input_text: str) -> str:
@@ -216,6 +228,7 @@ class Agent:
             result = await self.agent_executor.arun(input_text)
             return result
         except Exception as e:
+            logger.error(f"Error in agent execution: {e}")
             return f"Error: {str(e)}"
 
 class CodeFusion:
@@ -236,9 +249,43 @@ class CodeFusion:
 
 code_fusion = CodeFusion()
 
+def highlight_code(code: str, language: str) -> str:
+    lexer = get_lexer_by_name(language, stripall=True)
+    formatter = HtmlFormatter(style="monokai")
+    return pygments.highlight(code, lexer, formatter)
+
+def save_code_to_file(code: str, filename: str) -> str:
+    try:
+        with open(filename, 'w') as f:
+            f.write(code)
+        return f"Code saved to {filename}"
+    except Exception as e:
+        logger.error(f"Error saving code to file: {e}")
+        return f"Error saving code: {str(e)}"
+
 async def chat(message, history):
+    start_time = time.time()
     response = await code_fusion.run(message)
-    return response
+    end_time = time.time()
+    
+    # Highlight code in the response
+    highlighted_response = response
+    for lang in ['python', 'html', 'css', 'javascript']:
+        if f"```{lang}" in response:
+            code = response.split(f"```{lang}")[1].split("```")[0]
+            highlighted_code = highlight_code(code, lang)
+            highlighted_response = highlighted_response.replace(f"```{lang}{code}```", highlighted_code)
+    
+    # Save code to file if requested
+    if "save code" in message.lower():
+        filename = f"generated_code_{int(time.time())}.py"
+        save_result = save_code_to_file(response, filename)
+        highlighted_response += f"\n\n{save_result}"
+    
+    execution_time = end_time - start_time
+    highlighted_response += f"\n\nExecution time: {execution_time:.2f} seconds"
+    
+    return highlighted_response
 
 async def main():
     iface = gr.ChatInterface(
@@ -248,7 +295,8 @@ async def main():
         examples=[
             "Create a basic Flask app structure",
             "Implement a binary search algorithm in Python",
-            "Design a responsive navbar using HTML and CSS"
+            "Design a responsive navbar using HTML and CSS",
+            "Generate a simple REST API using Flask and save the code"
         ],
         retry_btn=None,
         undo_btn="Delete Previous",