Update app.py

app.py

@@ -1,70 +1,116 @@
 import streamlit as st
 from datetime import datetime
+import pandas as pd
+from io import StringIO
 
-# ---------- Configuration ----------
+# ---------- Page Config ----------
 st.set_page_config(
     page_title="Code Assistant",
     page_icon="🧠",
-    layout="centered"
+    layout="wide"
 )
 
-# ---------- State ----------
+# ---------- Session State ----------
 if "history" not in st.session_state:
     st.session_state.history = []
 
+if "datasets" not in st.session_state:
+    st.session_state.datasets = {}
+
 # ---------- Core Logic ----------
 @st.cache_data(show_spinner=False)
-def generate_code(prompt: str) -> str:
+def generate_code(prompt: str, model_id: str, dataset_name: str | None) -> str:
     """
-    Placeholder for an LLM / Hugging Face inference call.
-    Cached to avoid recomputation for identical prompts.
+    Placeholder for LLM inference.
     """
-    # Replace with real model call later
-    return f"""# Generated code
-# Prompt: {prompt}
+    context = f"# Model: {model_id}\n"
+    if dataset_name:
+        context += f"# Dataset: {dataset_name}\n"
+
+    return f"""{context}
+# Prompt:
+# {prompt}
 
 def example():
-    print("Hello from generated code!")
+    print("Generated using {model_id}")
 """
 
-def handle_generation(prompt: str):
+def handle_generation(prompt, model_id, dataset_name):
     if not prompt.strip():
         st.warning("Prompt cannot be empty.")
         return
 
     with st.spinner("Generating code…"):
-        result = generate_code(prompt)
+        output = generate_code(prompt, model_id, dataset_name)
 
     st.session_state.history.append({
         "timestamp": datetime.now().strftime("%H:%M:%S"),
+        "model": model_id,
+        "dataset": dataset_name,
         "prompt": prompt,
-        "output": result
+        "output": output
     })
 
-# ---------- UI ----------
-st.title("🧠 Code Assistant")
-st.caption("A minimal interface for turning intent into executable structure.")
+# ---------- Sidebar ----------
+with st.sidebar:
+    st.header("⚙️ Configuration")
+
+    # Model selection (static for now)
+    model_id = st.selectbox(
+        "Model",
+        options=[
+            "mock-codegen-v1",
+            "llama-3-8b",
+            "mistral-7b",
+        ],
+        help="Select the model used for code generation"
+    )
 
-user_input = st.text_area(
-    "Enter your coding prompt",
-    height=120,
-    placeholder="e.g. Write a Python function that validates JSON input"
-)
+    st.divider()
+
+    # Dataset upload
+    uploaded_file = st.file_uploader(
+        "Upload dataset",
+        type=["csv", "json", "txt"],
+        help="Optional dataset to condition generation"
+    )
+
+    if uploaded_file:
+        dataset_name = uploaded_file.name
 
-col1, col2 = st.columns([1, 3])
+        if dataset_name not in st.session_state.datasets:
+            if uploaded_file.type == "text/csv":
+                df = pd.read_csv(uploaded_file)
+            elif uploaded_file.type == "application/json":
+                df = pd.read_json(uploaded_file)
+            else:
+                df = StringIO(uploaded_file.read().decode("utf-8")).read()
 
-with col1:
-    generate_clicked = st.button("Generate", use_container_width=True)
+            st.session_state.datasets[dataset_name] = df
 
-with col2:
-    clear_clicked = st.button("Clear History", use_container_width=True)
+        st.success(f"Loaded: {dataset_name}")
 
-if generate_clicked:
-    handle_generation(user_input)
+    # Dataset selection
+    dataset_name = st.selectbox(
+        "Active dataset",
+        options=[None] + list(st.session_state.datasets.keys()),
+        format_func=lambda x: "None" if x is None else x
+    )
 
-if clear_clicked:
-    st.session_state.history.clear()
-    st.experimental_rerun()
+# ---------- Main UI ----------
+st.title("🧠 Code Assistant")
+st.caption("Prompt → Model → (Optional) Dataset → Code")
+
+prompt = st.text_area(
+    "Enter your coding prompt",
+    height=150,
+    placeholder="e.g. Generate a Python class that validates rows in the dataset"
+)
+
+generate = st.button("Generate", type="primary")
+
+if generate:
+    handle_generation(prompt, model_id, dataset_name)
 
 # ---------- Output ----------
 if st.session_state.history:
@@ -72,5 +118,13 @@ if st.session_state.history:
     st.subheader("Generation History")
 
     for item in reversed(st.session_state.history):
-        with st.expander(f"[{item['timestamp']}] {item['prompt'][:60]}"):
+        label = f"[{item['timestamp']}] {item['model']}"
+        if item["dataset"]:
+            label += f" | {item['dataset']}"
+
+        with st.expander(label):
+            st.markdown("**Prompt**")
+            st.write(item["prompt"])
+
+            st.markdown("**Generated Code**")
             st.code(item["output"], language="python")