Update app.py

app.py

@@ -1,18 +1,25 @@
 import streamlit as st
 import pandas as pd
 import os
-import requests
+import openai
+import numpy as np
 import json
 import re
 
 # — Page config
 st.set_page_config(page_title="CSV-Backed AI Agent", layout="wide")
 
-# — Title & (optional) image
+# — Title & image
 st.title("CSV-Backed AI Agent")
 st.image("./nadi-lok-image.png")
 
-# — Sidebar: CSV upload & preview (identical to your working code)
+# — Load your OpenAI key
+openai.api_key = os.getenv("OPENAI_API_KEY")
+if not openai.api_key:
+    st.error("❌ OPENAI_API_KEY not set in Settings → Secrets.")
+    st.stop()
+
+# — Sidebar CSV upload
 st.sidebar.header("Upload CSV File")
 uploaded_file = st.sidebar.file_uploader("Choose a CSV file", type="csv")
 
@@ -20,96 +27,98 @@ if uploaded_file:
     try:
         df = pd.read_csv(uploaded_file)
         st.sidebar.success("File uploaded successfully!")
-        st.sidebar.write("Preview of the uploaded file:")
+        st.sidebar.write("Preview:")
         st.sidebar.dataframe(df.head())
-        # Convert DataFrame back to CSV text to feed into the agent
-        csv_text = df.to_csv(index=False)
     except Exception as e:
         st.sidebar.error(f"Error reading file: {e}")
         df = None
-        csv_text = None
 else:
     df = None
-    csv_text = None
 
-# — Show basic info about the loaded CSV
 if df is not None:
     st.markdown(f"**Loaded CSV:** {df.shape[0]} rows × {df.shape[1]} columns")
 
+    @st.cache_data(show_spinner=False)
+    def build_row_embeddings(df: pd.DataFrame):
+        # Serialize each row to a compact JSON string
+        texts = df.apply(lambda row: row.to_json(), axis=1).tolist()
+
+        # Batch‐call embeddings
+        all_embs = []
+        for i in range(0, len(texts), 100):
+            batch = texts[i : i + 100]
+            resp = openai.Embedding.create(model="text-embedding-ada-002", input=batch)
+            all_embs.extend([d["embedding"] for d in resp["data"]])
+
+        return np.array(all_embs), texts
+
+    embeddings, row_texts = build_row_embeddings(df)
+
 # — Prompt input
 prompt = st.text_area(
     "Enter your prompt for the agent",
     placeholder="e.g. Which products have price > 100?",
-    height=150
+    height=150,
 )
 
-# — Run the AI Agent
+# — Run Agent
 if st.button("Run Agent"):
-    # 1️⃣ Validation
     if df is None:
         st.error("Please upload a CSV file first.")
     elif not prompt.strip():
         st.error("Please enter a prompt.")
     else:
-        # 2️⃣ Build the "memory + prompt" messages
+        # 1) Embed the prompt
+        q_resp = openai.Embedding.create(model="text-embedding-ada-002", input=[prompt])
+        q_emb = np.array(q_resp["data"][0]["embedding"])
+        # 2) Compute cosine similarities
+        row_norms = np.linalg.norm(embeddings, axis=1)
+        q_norm = np.linalg.norm(q_emb)
+        sims = embeddings.dot(q_emb) / (row_norms * q_norm + 1e-8)
+        # 3) Pick top K rows (e.g. 5)
+        K = min(5, len(sims))
+        top_idxs = sims.argsort()[-K:][::-1]
+        relevant_rows = [row_texts[i] for i in top_idxs]
+
+        # 4) Build the messages
         system_msg = {
             "role": "system",
             "content": (
-                "You are an AI agent that reads the provided CSV data and answers the user's query. "
-                "Return your answer strictly as JSON (no additional explanation)."
-            )
+                "You are an AI agent that reads the provided CSV rows and answers the user's query. "
+                "Return your answer strictly as JSON (no extra explanation)."
+            ),
         }
-        csv_msg = {
+        memory_msg = {
             "role": "system",
-            "content": f"CSV Data:\n{csv_text}\n<End of CSV>"
+            "content": "Relevant CSV rows:\n" + "\n".join(relevant_rows) + "\n<end of rows>",
         }
         user_msg = {"role": "user", "content": prompt}
 
-        # 3️⃣ Call OpenAI
-        api_key = os.getenv("OPENAI_API_KEY")
-        if not api_key:
-            st.error("❌ OPENAI_API_KEY not set in Secrets.")
-            st.stop()
-
-        response = requests.post(
-            "https://api.openai.com/v1/chat/completions",
-            headers={
-                "Authorization": f"Bearer {api_key}",
-                "Content-Type": "application/json",
-            },
-            json={
-                "model": "gpt-3.5-turbo",
-                "messages": [system_msg, csv_msg, user_msg],
-                "temperature": 0,
-                "max_tokens": 1500,
-            },
-            timeout=60,
+        # 5) Call ChatCompletion
+        chat = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[system_msg, memory_msg, user_msg],
+            temperature=0,
+            max_tokens=1500,
         )
+        answer = chat.choices[0].message.content
+
+        # 6) Extract JSON
+        txt = re.sub(r"```(?:json)?", "", answer).strip()
+        start = txt.find("{")
+        end = txt.rfind("}") + 1
+        parsed = None
+        if start >= 0 and end > 0:
+            frag = re.sub(r",\s*([}\]])", r"\1", txt[start:end])
+            try:
+                parsed = json.loads(frag)
+            except json.JSONDecodeError:
+                parsed = None
 
-        # 4️⃣ Handle API response
-        if response.status_code != 200:
-            st.error(f"API Error {response.status_code}: {response.text}")
+        # 7) Display
+        if parsed is not None:
+            st.subheader("✅ JSON Output")
+            st.json(parsed)
         else:
-            answer = response.json()["choices"][0]["message"]["content"]
-
-            # Strip any ``` fences and pull out the JSON object
-            txt = re.sub(r'```(?:json)?', '', answer).strip()
-            start = txt.find("{")
-            end = txt.rfind("}") + 1
-            parsed = None
-            if start >= 0 and end > 0:
-                frag = txt[start:end]
-                # remove trailing commas
-                frag = re.sub(r',\s*([}\]])', r'\1', frag)
-                try:
-                    parsed = json.loads(frag)
-                except json.JSONDecodeError:
-                    parsed = None
-
-            # 5️⃣ Display
-            if parsed is not None:
-                st.subheader("✅ JSON Output")
-                st.json(parsed)
-            else:
-                st.subheader("🔍 Raw Output")
-                st.text(answer)
+            st.subheader("🔍 Raw Output")
+            st.text(answer)