Update app.py

app.py

@@ -1,134 +1,132 @@
-import streamlit as st
+import gradio as gr
+from ctransformers import AutoModelForCausalLM
 import os
-from dotenv import load_dotenv
-from supplemental import EnhancedAIAgent, ProjectConfig
-from json import JSONDecodeError  # Add this line to import JSONDecodeError
-from transformers import pipeline
-
-class JSONDecodeError(Exception):
-    pass
-
-# Load environment variables
-load_dotenv()
-
-class JSONDecodeError(Exception):
-    pass
-
-def generate_project_config(project_description):
-    generator = pipeline('text-generation', model='distilgpt2', device=-1)
-    prompt = f"Generate a project configuration for the following project description: {project_description}\n\n{{"
-    
+import json
+import time
+import logging
+from threading import Lock
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Constants
+MODEL_NAME = "LoneStriker/speechless-zephyr-code-functionary-7b-GGUF"
+CACHE_DIR = "model_cache"
+CONFIG_FILE = "config.json"
+
+class EnhancedChatbot:
+    def __init__(self):
+        self.model = None
+        self.config = self.load_config()
+        self.model_lock = Lock()
+        self.load_model()
+
+    def load_config(self):
+        if os.path.exists(CONFIG_FILE):
+            with open(CONFIG_FILE, 'r') as f:
+                return json.load(f)
+        return {
+            "model_name": MODEL_NAME,
+            "max_tokens": 512,
+            "temperature": 0.7,
+            "top_p": 0.95,
+            "system_message": "You are a friendly and helpful AI assistant.",
+            "gpu_layers": 0
+        }
+
+    def save_config(self):
+        with open(CONFIG_FILE, 'w') as f:
+            json.dump(self.config, f, indent=2)
+
+    def load_model(self):
+        try:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.config["model_name"],
+                model_type="llama",
+                gpu_layers=self.config["gpu_layers"],
+                cache_dir=CACHE_DIR
+            )
+            logging.info(f"Model loaded successfully: {self.config['model_name']}")
+        except Exception as e:
+            logging.error(f"Error loading model: {str(e)}")
+            raise
+
+    def generate_response(self, message, history, system_message, max_tokens, temperature, top_p):
+        prompt = f"{system_message}\n\n"
+        for user_msg, assistant_msg in history:
+            prompt += f"Human: {user_msg}\nAssistant: {assistant_msg}\n"
+        prompt += f"Human: {message}\nAssistant: "
+
+        start_time = time.time()
+        with self.model_lock:
+            generated_text = self.model(
+                prompt,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+            )
+        end_time = time.time()
+
+        response_time = end_time - start_time
+        logging.info(f"Response generated in {response_time:.2f} seconds")
+
+        return generated_text.strip()
+
+chatbot = EnhancedChatbot()
+
+def respond(message, history, system_message, max_tokens, temperature, top_p):
     try:
-        response = generator(prompt, max_length=500, num_return_sequences=1)
-        generated_text = response[0]['generated_text']
-        
-        # Extract the JSON part
-        json_start = generated_text.index('{')
-        json_end = generated_text.rindex('}') + 1
-        config_str = generated_text[json_start:json_end]
-        
-        config = json.loads(config_str)
-        return config
-    except JSONDecodeError as e:
-        st.error(f"Error decoding JSON: {str(e)}")
-        return None
+        response = chatbot.generate_response(message, history, system_message, max_tokens, temperature, top_p)
+        yield response
     except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
-        return None
-
-st.title("Project Configuration Generator")
-
-project_description = st.text_area("Enter your project description:")
-
-if st.button("Generate Configuration"):
-    if project_description:
-        config = generate_project_config(project_description)
-        if config:
-            st.json(config)
-    else:
-        st.warning("Please enter a project description.")
-
-# Initialize EnhancedAIAgent
-agent = EnhancedAIAgent(
-    name="WebDevMaster",
-    description="Expert in full-stack web development",
-    skills=["HTML", "CSS", "JavaScript", "React", "Node.js"],
-    model_name="gpt2"
-)
-
-st.title("EnhancedAIAgent Web Development Assistant")
-
-st.write("""
-This is a powerful AI-driven web development assistant that can help you with various tasks such as:
-- Generating project configurations
-- Creating project structures
-- Implementing features
-- Reviewing and optimizing code
-- Generating documentation
-- Suggesting tests and refactoring improvements
-""")
-
-option = st.selectbox(
-    "What would you like to do?",
-    ("Generate Project Config", "Create Project Structure", "Implement Feature", 
-     "Review Code", "Optimize Code", "Generate Documentation", 
-     "Suggest Tests", "Explain Code", "Suggest Refactoring")
-)
-
-if option == "Generate Project Config":
-    project_description = st.text_area("Enter project description:")
-    if st.button("Generate"):
-        config = agent.generate_project_config(project_description)
-        st.json(config.__dict__)
-
-elif option == "Create Project Structure":
-    config_json = st.text_area("Enter ProjectConfig JSON:")
-    if st.button("Create"):
-        config_dict = eval(config_json)
-        config = ProjectConfig(**config_dict)
-        structure = agent.create_project_structure(config)
-        st.json(structure)
-
-elif option == "Implement Feature":
-    feature_description = st.text_area("Enter feature description:")
-    existing_code = st.text_area("Enter existing code (optional):")
-    if st.button("Implement"):
-        new_code = agent.implement_feature(feature_description, existing_code)
-        st.code(new_code)
-
-elif option == "Review Code":
-    code = st.text_area("Enter code to review:")
-    if st.button("Review"):
-        review = agent.review_code(code)
-        st.write(review)
-
-elif option == "Optimize Code":
-    code = st.text_area("Enter code to optimize:")
-    optimization_goal = st.text_input("Enter optimization goal:")
-    if st.button("Optimize"):
-        optimized_code = agent.optimize_code(code, optimization_goal)
-        st.code(optimized_code)
-
-elif option == "Generate Documentation":
-    code = st.text_area("Enter code to document:")
-    if st.button("Generate"):
-        documentation = agent.generate_documentation(code)
-        st.markdown(documentation)
-
-elif option == "Suggest Tests":
-    code = st.text_area("Enter code to suggest tests for:")
-    if st.button("Suggest"):
-        test_suggestions = agent.suggest_tests(code)
-        st.write(test_suggestions)
-
-elif option == "Explain Code":
-    code = st.text_area("Enter code to explain:")
-    if st.button("Explain"):
-        explanation = agent.explain_code(code)
-        st.write(explanation)
-
-elif option == "Suggest Refactoring":
-    code = st.text_area("Enter code to suggest refactoring:")
-    if st.button("Suggest"):
-        refactoring_suggestions = agent.suggest_refactoring(code)
-        st.write(refactoring_suggestions)
+        logging.error(f"Error generating response: {str(e)}")
+        yield "I apologize, but I encountered an error while processing your request. Please try again."
+
+def update_model_config(model_name, gpu_layers):
+    chatbot.config["model_name"] = model_name
+    chatbot.config["gpu_layers"] = gpu_layers
+    chatbot.save_config()
+    chatbot.load_model()
+    return f"Model updated to {model_name} with {gpu_layers} GPU layers."
+
+def update_system_message(system_message):
+    chatbot.config["system_message"] = system_message
+    chatbot .save_config()
+    return f"System message updated: {system_message}"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Enhanced AI Chatbot")
+    
+    with gr.Tab("Chat"):
+        chatbot_interface= gr.ChatInterface(
+            respond,
+            additional_inputs=[
+                gr.Textbox(value=chatbot.config["system_message"], label="System message"),
+                gr.Slider(minimum=1, maximum=2048, value=chatbot.config["max_tokens"], step=1, label="Max new tokens"),
+                gr.Slider(minimum=0.1, maximum=4.0, value=chatbot.config["temperature"], step=0.1, label="Temperature"),
+                gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=chatbot.config["top_p"],
+                    step=0.05,
+                    label="Top-p (nucleus sampling)",
+                ),
+            ],
+        )
+    
+    with gr.Tab("Settings"):
+        with gr.Group():
+            gr.Markdown("### Model Settings")
+            model_name_input = gr.Textbox(value=chatbot.config["model_name"], label="Model name")
+            gpu_layers_input = gr.Slider(minimum=0, maximum=8, value=chatbot.config["gpu_layers"], step=1, label="GPU layers")
+            update_model_button = gr.Button("Update model")
+            update_model_button.click(update_model_config, inputs=[model_name_input, gpu_layers_input], outputs="text")
+        
+        with gr.Group():
+            gr.Markdown("### System Message Settings")
+            system_message_input = gr.Textbox(value=chatbot.config["system_message"], label="System message")
+            update_system_message_button = gr.Button("Update system message")
+            update_system_message_button.click(update_system_message, inputs=[system_message_input], outputs="text")
+
+if __name__ == "__main__":
+    demo.launch()