Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,116 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import pinecone
+import os
+import subprocess
+import json
+import requests
+from pinecone import Pinecone, ServerlessSpec
+#from streamlit_extras.add_vertical_space import add_vertical_space
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# Set page config
+st.set_page_config(page_title="StarCoder AI Code Generator", page_icon="🤖", layout="wide")
+
+# Custom CSS for Grok-like UI
+st.markdown("""
+    <style>
+        body {background-color: #121212; color: white;}
+        .stApp {background: #181818; color: white; font-family: 'Arial', sans-serif;}
+        .stTextArea textarea {background: #282828; color: white; border-radius: 10px;}
+        .stTextInput input {background: #282828; color: white; border-radius: 10px;}
+        .stButton>button {background: #3f51b5; color: white; border-radius: 8px; padding: 10px;}
+        .stDownloadButton>button {background: #4caf50; color: white; border-radius: 8px; padding: 10px;}
+        .stCode {background: #202020; padding: 15px; border-radius: 10px;}
+        .stSidebar {background: #202020; color: white;}
+        .stSidebar .stButton>button {background: #ff9800; color: white;}
+    </style>
+""", unsafe_allow_html=True)
+
+# Pinecone API Key
+PINECONE_API_KEY = "your pinecone api key here"
+pc = Pinecone(api_key=PINECONE_API_KEY)
+INDEX_NAME = "code-gen"
+if INDEX_NAME not in pc.list_indexes().names():
+    pc.create_index(name=INDEX_NAME, dimension=1536, metric='euclidean',
+                    spec=ServerlessSpec(cloud='aws', region='us-east-1'))
+index = pc.Index(INDEX_NAME)
+
+# Hugging Face API Key
+HUGGINGFACE_API_KEY = "your huggingface api key"
+
+# Initialize session state
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+if "current_chat" not in st.session_state:
+    st.session_state.current_chat = []
+
+
+def new_chat():
+    if st.session_state.current_chat:
+        st.session_state.chat_history.append(st.session_state.current_chat)
+    st.session_state.current_chat = []
+
+
+# Sidebar - Chat History
+st.sidebar.title("💬 Chat History")
+for i, chat in enumerate(st.session_state.chat_history):
+    if st.sidebar.button(f"Chat {i + 1} 🕒"):
+        st.session_state.current_chat = chat
+
+# Sidebar - New Chat Button
+if st.sidebar.button("New Chat ➕"):
+    new_chat()
+
+
+# Function to query Hugging Face model
+def generate_code(prompt, examples=[]):
+    formatted_prompt = "".join(examples) + "\n" + prompt
+    headers = {"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"}
+    data = {"inputs": formatted_prompt, "parameters": {"temperature": 0.5, "max_length": 512}}
+    response = requests.post("https://api-inference.huggingface.co/models/bigcode/starcoder2-15b", headers=headers,
+                             json=data)
+    return response.json()[0][
+        'generated_text'].strip() if response.status_code == 200 else "Error: Unable to generate code"
+
+
+def execute_python_code(code):
+    try:
+        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, timeout=5)
+        return result.stdout if result.returncode == 0 else result.stderr
+    except Exception as e:
+        return str(e)
+
+
+def store_in_pinecone(prompt, code):
+    vector = [0.1] * 1536  # Placeholder vector
+    index.upsert([(prompt, vector, {"code": code})])
+
+
+def retrieve_from_pinecone(prompt):
+    response = index.query(vector=[0.1] * 1536, top_k=2, include_metadata=True)
+    return [r["metadata"]["code"] for r in response["matches"]]
+
+
+# Main UI
+st.title("🤖 StarCoder AI Code Generator")
+
+# User Input
+prompt = st.text_area("✍️ Enter your prompt:", height=100)
+mode = st.selectbox("⚡ Choose prompting mode", ["One-shot", "Two-shot"])
+
+generate_button = st.button("🚀 Generate Code")
+
+if generate_button and prompt:
+    examples = retrieve_from_pinecone(prompt) if mode == "Two-shot" else []
+    generated_code = generate_code(prompt, examples)
+    st.session_state.current_chat.append((prompt, generated_code))
+    st.subheader("📜 Generated Code")
+    st.code(generated_code, language="python")
+    store_in_pinecone(prompt, generated_code)
+
+    execute_button = st.button("▶️ Run Code")
+    if execute_button:
+        execution_result = execute_python_code(generated_code)
+        st.subheader("🖥️ Execution Output")
+        st.text(execution_result)
+
+    st.download_button("📥 Download Code", generated_code, file_name="generated_code.py")