Update app.py

app.py

@@ -43,29 +43,6 @@ else:
 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):
-        # 1) Convert each row to a compact JSON string
-        texts = df.apply(lambda r: r.to_json(), axis=1).tolist()
-
-        # 2) Batch‐call the embeddings endpoint
-        embeddings = []
-        for i in range(0, len(texts), 100):
-            batch = texts[i : i + 100]
-            resp = requests.post(
-                "https://api.openai.com/v1/embeddings",
-                headers=HEADERS,
-                json={"model": "text-embedding-ada-002", "input": batch},
-                timeout=60,
-            )
-            resp.raise_for_status()
-            data = resp.json()["data"]
-            embeddings.extend(d["embedding"] for d in data)
-
-        return np.array(embeddings), texts
-
-    embeddings, row_texts = build_row_embeddings(df)
-
 # — Prompt input
 prompt = st.text_area(
     "Enter your prompt for the agent",
@@ -73,15 +50,23 @@ prompt = st.text_area(
     height=150,
 )
 
-# — Define function for OpenAI function calling
+# — Functions for function calling
 def search_csv(query: str):
-    # Run a Pandas query safely
     try:
         result_df = df.query(query)
-        return result_df.to_dict(orient="records")
+        # Limit output to 50 rows for large results
+        return result_df.head(50).to_dict(orient="records")
     except Exception as e:
-        return {"error": f"Query error: {str(e)}"}
+        return {"error": f"Invalid query. Example: 'price > 100'. Details: {str(e)}"}
 
+def count_unique(column: str):
+    try:
+        n = df[column].nunique()
+        return {"column": column, "unique_count": int(n)}
+    except Exception as e:
+        return {"error": f"Column '{column}' not found or not countable. Details: {str(e)}"}
+
+# — Function schemas for OpenAI
 function_schema = [
     {
         "name": "search_csv",
@@ -96,9 +81,29 @@ function_schema = [
             },
             "required": ["query"],
         },
+    },
+    {
+        "name": "count_unique",
+        "description": "Count the number of unique values in a column.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "column": {
+                    "type": "string",
+                    "description": "The column name to count unique values."
+                },
+            },
+            "required": ["column"],
+        },
     }
 ]
 
+# — Map function names to Python functions
+function_map = {
+    "search_csv": search_csv,
+    "count_unique": count_unique,
+}
+
 # — Run Agent
 if st.button("Run Agent"):
     if df is None:
@@ -106,24 +111,28 @@ if st.button("Run Agent"):
     elif not prompt.strip():
         st.error("Please enter a prompt.")
     else:
-        # 1) First call: ask OpenAI if it wants to use a function
+        # Build dynamic system message with column info
+        columns = ", ".join(df.columns)
+        system_message = {
+            "role": "system",
+            "content": (
+                f"You are an AI agent helping users analyze a CSV file with these columns: {columns}. "
+                "If you need to search or filter the CSV, call the 'search_csv' function. "
+                "If the user wants to know how many unique values are in a column, call the 'count_unique' function. "
+                "If you use 'search_csv', use Pandas query syntax."
+            ),
+        }
         messages = [
-            {
-                "role": "system",
-                "content": (
-                    "You are an AI agent helping users analyze a CSV file. "
-                    "If you need to search or filter the CSV, call the 'search_csv' function. "
-                    "Only use the function when you need data from the CSV."
-                ),
-            },
+            system_message,
             {"role": "user", "content": prompt}
         ]
 
+        # First OpenAI call: See if a function call is needed
         chat_resp = requests.post(
             "https://api.openai.com/v1/chat/completions",
             headers=HEADERS,
             json={
-                "model": "gpt-3.5-turbo-1106",  # or "gpt-4-1106-preview" if available
+                "model": "gpt-3.5-turbo-1106",  # or "gpt-4-1106-preview" if you have it
                 "messages": messages,
                 "functions": function_schema,
                 "function_call": "auto",
@@ -136,27 +145,25 @@ if st.button("Run Agent"):
         response_json = chat_resp.json()
         msg = response_json["choices"][0]["message"]
 
-        # 2) Check if function call is requested
+        # If OpenAI requests a function call
         if msg.get("function_call"):
             func_name = msg["function_call"]["name"]
             args_json = msg["function_call"]["arguments"]
             args = json.loads(args_json)
-            # Only one function: search_csv
-            search_result = search_csv(args["query"])
+            # Call the appropriate function
+            if func_name in function_map:
+                function_result = function_map[func_name](**args)
+            else:
+                function_result = {"error": f"Unknown function: {func_name}"}
 
-            # 3) Pass function result back to OpenAI for final answer
+            # Send function result back for final answer
             followup_messages = [
-                {
-                    "role": "system",
-                    "content": (
-                        "You are an AI agent helping users analyze a CSV file."
-                    ),
-                },
+                system_message,
                 {"role": "user", "content": prompt},
                 {
                     "role": "function",
                     "name": func_name,
-                    "content": json.dumps(search_result),
+                    "content": json.dumps(function_result),
                 }
             ]
 
@@ -164,7 +171,7 @@ if st.button("Run Agent"):
                 "https://api.openai.com/v1/chat/completions",
                 headers=HEADERS,
                 json={
-                    "model": "gpt-3.5-turbo-1106",  # or "gpt-4-1106-preview"
+                    "model": "gpt-3.5-turbo-1106",
                     "messages": followup_messages,
                     "temperature": 0,
                     "max_tokens": 1500,
@@ -176,8 +183,8 @@ if st.button("Run Agent"):
 
             st.subheader("✅ Agent Answer")
             st.markdown(answer)
-            st.subheader("📊 Filtered CSV Data")
-            st.json(search_result)
+            st.subheader("🔎 Function Output")
+            st.json(function_result)
         else:
             # No function call: model answered directly
             st.subheader("✅ Agent Answer")