Update app.py

app.py

@@ -19,6 +19,9 @@ import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
 from IPython.display import display
+import json
+import html
+import re
 
 # Configure the Gemini API
 genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])
@@ -32,10 +35,11 @@ generation_config = {
 }
 
 model = genai.GenerativeModel(
-    model_name="gemini-1.5-pro",
+    model_name="gemini-2.0-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.
+    Focus on generating code that is efficient, readable, and adheres to best practices. Ensure that the code is well-documented and includes error handling where necessary.
     """
 )
 chat_session = model.start_chat(history=[])
@@ -48,12 +52,21 @@ def generate_response(user_input):
     except Exception as e:
         return f"Error: {e}"
 
-def optimize_code(code):
+def optimize_code(code, language):
     """Optimize the generated code using static analysis tools."""
-    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")
+    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)
     return code
 
 def fetch_from_github(query):
@@ -90,12 +103,19 @@ def initialize_git_repo(repo_path):
         repo = git.Repo(repo_path)
     return repo
 
-def integrate_with_git(repo_path, code):
+def integrate_with_git(repo_path, code, language):
     """Integrate the generated code with a Git repository."""
     repo = initialize_git_repo(repo_path)
-    with open(os.path.join(repo_path, "generated_code.py"), "w") as file:
+    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:
         file.write(code)
-    repo.index.add(["generated_code.py"])
+    repo.index.add([filename])
     repo.index.commit("Added generated code")
 
 def process_user_input(user_input):
@@ -125,16 +145,24 @@ def run_tests():
     test_result = test_runner.run(test_suite)
     return test_result
 
-def execute_code_in_docker(code):
+def execute_code_in_docker(code, language):
     """Execute code in a Docker container for safety and isolation."""
     client = docker.from_env()
     try:
-        container = client.containers.run(
-            image="python:3.9",
-            command=f"python -c '{code}'",
-            detach=True,
-            remove=True
-        )
+        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
+            )
         result = container.wait()
         logs = container.logs().decode('utf-8')
         return logs, result['StatusCode']
@@ -172,6 +200,16 @@ 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")
 
@@ -267,6 +305,7 @@ 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() == "":
@@ -279,18 +318,17 @@ if st.button("Generate Code"):
                 if "Error" in completed_text:
                     handle_error(completed_text)
                 else:
-                    optimized_code = optimize_code(completed_text)
+                    optimized_code = optimize_code(completed_text, language)
                     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('<div class="code-block">', unsafe_allow_html=True)
-                    st.code(optimized_code)
-                    st.markdown('</div>', unsafe_allow_html=True)
+                    st.markdown(formatted_code, 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)
+                    integrate_with_git(repo_path, optimized_code, language)
                     
                     # Run automated tests
                     test_result = run_tests()
@@ -300,7 +338,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)
+                    execution_result, status_code = execute_code_in_docker(optimized_code, language)
                     if status_code == 0:
                         st.success("Code executed successfully in Docker!")
                         st.text(execution_result)