Update app.py

app.py

@@ -5,30 +5,59 @@ import subprocess
 import os
 import pylint
 import pandas as pd
+import numpy as np
 from sklearn.model_selection import train_test_split
 from sklearn.ensemble import RandomForestClassifier
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from transformers import AutoTokenizer, AutoModel
+import ast
+import networkx as nx
+import matplotlib.pyplot as plt
 
 # Configure the Gemini API
 genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])
 
 # Create the model with optimized parameters and enhanced system instructions
 generation_config = {
-    "temperature": 0.6,  # Lower temperature for more deterministic responses
-    "top_p": 0.8,        # Adjusted for better diversity
-    "top_k": 30,         # Increased for more diverse tokens
-    "max_output_tokens": 16384,  # Increased for longer responses
+    "temperature": 0.6,
+    "top_p": 0.8,
+    "top_k": 30,
+    "max_output_tokens": 16384,
 }
 
 model = genai.GenerativeModel(
     model_name="gemini-1.5-pro",
     generation_config=generation_config,
     system_instruction="""
-    You are Ath, a highly knowledgeable code assistant. You speak in a friendly and casual tone, using bro words, and you never mention you're an AI. 
-    Your responses should contain optimized, secure, and high-quality code only, without explanations. You are designed to provide accurate, efficient, and cutting-edge code solutions.
+    You are Ath, a highly advanced code assistant with deep knowledge in AI, machine learning, and software engineering. You provide cutting-edge, optimized, and secure code solutions. Speak casually and use tech jargon when appropriate.
     """
 )
 chat_session = model.start_chat(history=[])
 
+# Load pre-trained BERT model for code understanding
+tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
+codebert_model = AutoModel.from_pretrained("microsoft/codebert-base")
+
+class CodeImprovement(nn.Module):
+    def __init__(self, input_dim):
+        super(CodeImprovement, self).__init__()
+        self.fc1 = nn.Linear(input_dim, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, 128)
+        self.fc4 = nn.Linear(128, 2)  # Binary classification: needs improvement or not
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = torch.relu(self.fc3(x))
+        return torch.sigmoid(self.fc4(x))
+
+code_improvement_model = CodeImprovement(768)  # 768 is BERT's output dimension
+optimizer = optim.Adam(code_improvement_model.parameters())
+criterion = nn.BCELoss()
+
 def generate_response(user_input):
     try:
         response = chat_session.send_message(user_input)
@@ -37,151 +66,129 @@ def generate_response(user_input):
         return f"Error: {e}"
 
 def optimize_code(code):
-    # Placeholder for advanced code optimization logic
-    # This could involve using external tools or libraries for static analysis and optimization
+    # Use abstract syntax tree for advanced code analysis
+    tree = ast.parse(code)
+    analyzer = CodeAnalyzer()
+    analyzer.visit(tree)
+    
+    # Apply code transformations based on analysis
+    transformer = CodeTransformer(analyzer.get_optimizations())
+    optimized_tree = transformer.visit(tree)
+    
+    optimized_code = ast.unparse(optimized_tree)
+    
+    # Run pylint for additional suggestions
     with open("temp_code.py", "w") as file:
-        file.write(code)
+        file.write(optimized_code)
     result = subprocess.run(["pylint", "temp_code.py"], capture_output=True, text=True)
     os.remove("temp_code.py")
-    return code
+    
+    return optimized_code, result.stdout
 
 def fetch_from_github(query):
-    # Placeholder for fetching code snippets from GitHub
-    # This could involve using the GitHub API to search for relevant code
-    return ""
+    headers = {"Authorization": f"token {st.secrets['GITHUB_TOKEN']}"}
+    response = requests.get(f"https://api.github.com/search/code?q={query}", headers=headers)
+    if response.status_code == 200:
+        return response.json()['items'][:5]  # Return top 5 results
+    return []
 
 def interact_with_api(api_url):
-    # Placeholder for interacting with external APIs
     response = requests.get(api_url)
     return response.json()
 
 def train_ml_model(code_data):
-    # Placeholder for training a machine learning model to predict code improvements
     df = pd.DataFrame(code_data)
     X = df.drop('target', axis=1)
     y = df['target']
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
-    model = RandomForestClassifier()
+    model = RandomForestClassifier(n_estimators=100, max_depth=10)
     model.fit(X_train, y_train)
     return model
 
-# Streamlit UI setup
-st.set_page_config(page_title="Sleek AI Code Assistant", page_icon="💻", layout="wide")
+def analyze_code_quality(code):
+    # Tokenize and encode the code
+    inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512, padding="max_length")
+    
+    # Get BERT embeddings
+    with torch.no_grad():
+        outputs = codebert_model(**inputs)
+    
+    # Use the [CLS] token embedding for classification
+    cls_embedding = outputs.last_hidden_state[:, 0, :]
+    
+    # Pass through our code improvement model
+    prediction = code_improvement_model(cls_embedding)
+    
+    return prediction.item()  # Return the probability of needing improvement
 
-st.markdown("""
-<style>
-    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap');
+def visualize_code_structure(code):
+    tree = ast.parse(code)
+    graph = nx.DiGraph()
     
-    body {
-        font-family: 'Inter', sans-serif;
-        background-color: #f0f4f8;
-        color: #1a202c;
-    }
-    .stApp {
-        max-width: 1000px;
-        margin: 0 auto;
-        padding: 2rem;
-    }
-    .main-container {
-        background: #ffffff;
-        border-radius: 16px;
-        padding: 2rem;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
-    }
-    h1 {
-        font-size: 2.5rem;
-        font-weight: 700;
-        color: #2d3748;
-        text-align: center;
-        margin-bottom: 1rem;
-    }
-    .subtitle {
-        font-size: 1.1rem;
-        text-align: center;
-        color: #4a5568;
-        margin-bottom: 2rem;
-    }
-    .stTextArea textarea {
-        border: 2px solid #e2e8f0;
-        border-radius: 8px;
-        font-size: 1rem;
-        padding: 0.75rem;
-        transition: all 0.3s ease;
-    }
-    .stTextArea textarea:focus {
-        border-color: #4299e1;
-        box-shadow: 0 0 0 3px rgba(66, 153, 225, 0.5);
-    }
-    .stButton button {
-        background-color: #4299e1;
-        color: white;
-        border: none;
-        border-radius: 8px;
-        font-size: 1.1rem;
-        font-weight: 600;
-        padding: 0.75rem 2rem;
-        transition: all 0.3s ease;
-        width: 100%;
-    }
-    .stButton button:hover {
-        background-color: #3182ce;
-    }
-    .output-container {
-        background: #f7fafc;
-        border-radius: 8px;
-        padding: 1rem;
-        margin-top: 2rem;
-    }
-    .code-block {
-        background-color: #2d3748;
-        color: #e2e8f0;
-        font-family: 'Fira Code', monospace;
-        font-size: 0.9rem;
-        border-radius: 8px;
-        padding: 1rem;
-        margin-top: 1rem;
-        overflow-x: auto;
-    }
-    .stAlert {
-        background-color: #ebf8ff;
-        color: #2b6cb0;
-        border-radius: 8px;
-        border: none;
-        padding: 0.75rem 1rem;
-    }
-    .stSpinner {
-        color: #4299e1;
-    }
-</style>
-""", unsafe_allow_html=True)
+    def add_nodes_edges(node, parent=None):
+        node_id = id(node)
+        graph.add_node(node_id, label=type(node).__name__)
+        if parent:
+            graph.add_edge(id(parent), node_id)
+        for child in ast.iter_child_nodes(node):
+            add_nodes_edges(child, node)
+    
+    add_nodes_edges(tree)
+    
+    plt.figure(figsize=(12, 8))
+    pos = nx.spring_layout(graph)
+    nx.draw(graph, pos, with_labels=True, node_color='lightblue', node_size=1000, font_size=8, font_weight='bold')
+    labels = nx.get_node_attributes(graph, 'label')
+    nx.draw_networkx_labels(graph, pos, labels, font_size=6)
+    
+    return plt
+
+# Streamlit UI setup
+st.set_page_config(page_title="Advanced AI Code Assistant", page_icon="🚀", layout="wide")
+
+# ... (keep the existing CSS styles) ...
 
 st.markdown('<div class="main-container">', unsafe_allow_html=True)
-st.title("💻 Sleek AI Code Assistant")
-st.markdown('<p class="subtitle">Powered by Google Gemini</p>', unsafe_allow_html=True)
+st.title("🚀 Advanced AI Code Assistant")
+st.markdown('<p class="subtitle">Powered by Google Gemini & Deep Learning</p>', unsafe_allow_html=True)
 
-prompt = st.text_area("What code can I help you with today?", height=120)
+prompt = st.text_area("What advanced code task can I help you with today?", height=120)
 
-if st.button("Generate Code"):
+if st.button("Generate Advanced Code"):
     if prompt.strip() == "":
         st.error("Please enter a valid prompt.")
     else:
-        with st.spinner("Generating code..."):
+        with st.spinner("Generating and analyzing code..."):
             completed_text = generate_response(prompt)
             if "Error" in completed_text:
                 st.error(completed_text)
             else:
-                optimized_code = optimize_code(completed_text)
-                st.success("Code generated and optimized successfully!")
+                optimized_code, lint_results = optimize_code(completed_text)
+                quality_score = analyze_code_quality(optimized_code)
+                
+                st.success(f"Code generated and optimized successfully! Quality Score: {quality_score:.2f}")
                 
                 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)
+                
+                with st.expander("View Code Structure Visualization"):
+                    st.pyplot(visualize_code_structure(optimized_code))
+                
+                with st.expander("View Lint Results"):
+                    st.text(lint_results)
+                
+                with st.expander("Fetch Similar Code from GitHub"):
+                    github_results = fetch_from_github(prompt)
+                    for item in github_results:
+                        st.markdown(f"[{item['name']}]({item['html_url']})")
+                
                 st.markdown('</div>', unsafe_allow_html=True)
 
 st.markdown("""
 <div style='text-align: center; margin-top: 2rem; color: #4a5568;'>
-    Created with ❤️ by Your Sleek AI Code Assistant
+    Crafted with 🚀 by Your Advanced AI Code Assistant
 </div>
 """, unsafe_allow_html=True)