Update app.py

app.py

@@ -19,9 +19,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
 from IPython.display import display
-import json
-import html
-import re
+from tenacity import retry, stop_after_attempt, wait_fixed
 
 # Configure the Gemini API
 genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])
@@ -35,7 +33,7 @@ generation_config = {
 }
 
 model = genai.GenerativeModel(
-    model_name="gemini-2.0-pro",
+    model_name="gemini-1.5-pro",
     generation_config=generation_config,
     system_instruction="""
     You are Ath, a highly knowledgeable and advanced code assistant. Your responses are optimized for secure, high-quality, and cutting-edge code solutions.
@@ -44,29 +42,21 @@ model = genai.GenerativeModel(
 )
 chat_session = model.start_chat(history=[])
 
+@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
 def generate_response(user_input):
-    """Generate a response from the AI model."""
+    """Generate a response from the AI model with retry mechanism."""
     try:
         response = chat_session.send_message(user_input)
         return response.text
     except Exception as e:
         return f"Error: {e}"
 
-def optimize_code(code, language):
+def optimize_code(code):
     """Optimize the generated code using static analysis tools."""
-    if language == 'python':
-        with open("temp_code.py", "w") as file:
-            file.write(code)
-        result = subprocess.run(["pylint", "temp_code.py"], capture_output=True, text=True)
-        os.remove("temp_code.py")
-    elif language == 'javascript':
-        with open("temp_code.js", "w") as file:
-            file.write(code)
-        result = subprocess.run(["eslint", "temp_code.js"], capture_output=True, text=True)
-        os.remove("temp_code.js")
-    elif language == 'html':
-        # HTML optimization can be more complex due to its structure
-        code = html.escape(code)
+    with open("temp_code.py", "w") as file:
+        file.write(code)
+    result = subprocess.run(["pylint", "temp_code.py"], capture_output=True, text=True)
+    os.remove("temp_code.py")
     return code
 
 def fetch_from_github(query):
@@ -103,19 +93,12 @@ def initialize_git_repo(repo_path):
         repo = git.Repo(repo_path)
     return repo
 
-def integrate_with_git(repo_path, code, language):
+def integrate_with_git(repo_path, code):
     """Integrate the generated code with a Git repository."""
     repo = initialize_git_repo(repo_path)
-    filename = "generated_code"
-    if language == 'python':
-        filename += ".py"
-    elif language == 'javascript':
-        filename += ".js"
-    elif language == 'html':
-        filename += ".html"
-    with open(os.path.join(repo_path, filename), "w") as file:
+    with open(os.path.join(repo_path, "generated_code.py"), "w") as file:
         file.write(code)
-    repo.index.add([filename])
+    repo.index.add(["generated_code.py"])
     repo.index.commit("Added generated code")
 
 def process_user_input(user_input):
@@ -145,24 +128,16 @@ def run_tests():
     test_result = test_runner.run(test_suite)
     return test_result
 
-def execute_code_in_docker(code, language):
+def execute_code_in_docker(code):
     """Execute code in a Docker container for safety and isolation."""
     client = docker.from_env()
     try:
-        if language == 'python':
-            container = client.containers.run(
-                image="python:3.9",
-                command=f"python -c '{code}'",
-                detach=True,
-                remove=True
-            )
-        elif language == 'javascript':
-            container = client.containers.run(
-                image="node:14",
-                command=f"node -e '{code}'",
-                detach=True,
-                remove=True
-            )
+        container = client.containers.run(
+            image="python:3.9",
+            command=f"python -c '{code}'",
+            detach=True,
+            remove=True
+        )
         result = container.wait()
         logs = container.logs().decode('utf-8')
         return logs, result['StatusCode']
@@ -200,16 +175,6 @@ def display_dataframe(data):
     df = pd.DataFrame(data)
     display(df)
 
-def format_code(code, language):
-    """Format code using appropriate syntax highlighting."""
-    if language == 'python':
-        return f"```python\n{code}\n```"
-    elif language == 'javascript':
-        return f"```javascript\n{code}\n```"
-    elif language == 'html':
-        return f"```html\n{code}\n```"
-    return code
-
 # Streamlit UI setup
 st.set_page_config(page_title="Ultra AI Code Assistant", page_icon="🚀", layout="wide")
 
@@ -305,7 +270,6 @@ st.title("🚀 Ultra AI Code Assistant")
 st.markdown('<p class="subtitle">Powered by Google Gemini</p>', unsafe_allow_html=True)
 
 prompt = st.text_area("What code can I help you with today?", height=120)
-language = st.selectbox("Select language", ["python", "javascript", "html"])
 
 if st.button("Generate Code"):
     if prompt.strip() == "":
@@ -318,17 +282,18 @@ if st.button("Generate Code"):
                 if "Error" in completed_text:
                     handle_error(completed_text)
                 else:
-                    optimized_code = optimize_code(completed_text, language)
+                    optimized_code = optimize_code(completed_text)
                     st.success("Code generated and optimized successfully!")
                     
-                    formatted_code = format_code(optimized_code, language)
                     st.markdown('<div class="output-container">', unsafe_allow_html=True)
-                    st.markdown(formatted_code, unsafe_allow_html=True)
+                    st.markdown('<div class="code-block">', unsafe_allow_html=True)
+                    st.code(optimized_code)
+                    st.markdown('</div>', unsafe_allow_html=True)
                     st.markdown('</div>', unsafe_allow_html=True)
                     
                     # Integrate with Git
                     repo_path = "./repo"  # Replace with your repository path
-                    integrate_with_git(repo_path, optimized_code, language)
+                    integrate_with_git(repo_path, optimized_code)
                     
                     # Run automated tests
                     test_result = run_tests()
@@ -338,7 +303,7 @@ if st.button("Generate Code"):
                         st.error("Some tests failed. Please check the code.")
                     
                     # Execute code in Docker
-                    execution_result, status_code = execute_code_in_docker(optimized_code, language)
+                    execution_result, status_code = execute_code_in_docker(optimized_code)
                     if status_code == 0:
                         st.success("Code executed successfully in Docker!")
                         st.text(execution_result)