Update app.py

app.py

@@ -7,10 +7,10 @@ import json
 import re
 
 # — Page config
-st.set_page_config(page_title="CSV-Backed AI Agent", layout="wide")
+st.set_page_config(page_title="CSV-Backed AI Agent with Function Calling", layout="wide")
 
 # — Title & image
-st.title("CSV-Backed AI Agent")
+st.title("CSV-Backed AI Agent with Function Calling")
 st.image("./nadi-lok-image.png")
 
 # — Load API key
@@ -73,6 +73,32 @@ prompt = st.text_area(
     height=150,
 )
 
+# — Define function for OpenAI 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")
+    except Exception as e:
+        return {"error": f"Query error: {str(e)}"}
+
+function_schema = [
+    {
+        "name": "search_csv",
+        "description": "Filter the CSV rows by a Pandas query. Example: price > 100",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "A Pandas query string, e.g. 'price > 100 and category == \"A\"'"
+                },
+            },
+            "required": ["query"],
+        },
+    }
+]
+
 # — Run Agent
 if st.button("Run Agent"):
     if df is None:
@@ -80,71 +106,79 @@ if st.button("Run Agent"):
     elif not prompt.strip():
         st.error("Please enter a prompt.")
     else:
-        # 1) Embed the prompt
-        q_resp = requests.post(
-            "https://api.openai.com/v1/embeddings",
-            headers=HEADERS,
-            json={"model": "text-embedding-ada-002", "input": [prompt]},
-            timeout=60,
-        )
-        q_resp.raise_for_status()
-        q_emb = np.array(q_resp.json()["data"][0]["embedding"])
-
-        # 2) Compute cosine similarities
-        norms = np.linalg.norm(embeddings, axis=1)
-        q_norm = np.linalg.norm(q_emb)
-        sims = embeddings.dot(q_emb) / (norms * q_norm + 1e-8)
-
-        # 3) Select top-K relevant rows
-        K = min(5, len(sims))
-        top_idxs = sims.argsort()[-K:][::-1]
-        relevant_rows = [row_texts[i] for i in top_idxs]
-
-        # 4) Build the chat messages
-        system_msg = {
-            "role": "system",
-            "content": (
-                "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)."
-            ),
-        }
-        memory_msg = {
-            "role": "system",
-            "content": "Relevant CSV rows:\n" + "\n".join(relevant_rows) + "\n<end of rows>",
-        }
-        user_msg = {"role": "user", "content": prompt}
-
-        # 5) Call the Chat API
+        # 1) First call: ask OpenAI if it wants to use a function
+        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."
+                ),
+            },
+            {"role": "user", "content": prompt}
+        ]
+
         chat_resp = requests.post(
             "https://api.openai.com/v1/chat/completions",
             headers=HEADERS,
             json={
-                "model": "gpt-3.5-turbo",
-                "messages": [system_msg, memory_msg, user_msg],
+                "model": "gpt-3.5-turbo-1106",  # or "gpt-4-1106-preview" if available
+                "messages": messages,
+                "functions": function_schema,
+                "function_call": "auto",
                 "temperature": 0,
-                "max_tokens": 1500,
+                "max_tokens": 1000,
             },
             timeout=60,
         )
         chat_resp.raise_for_status()
-        answer = chat_resp.json()["choices"][0]["message"]["content"]
-
-        # 6) Extract the JSON payload
-        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
-
-        # 7) Display results
-        if parsed is not None:
-            st.subheader("✅ JSON Output")
-            st.json(parsed)
+        response_json = chat_resp.json()
+        msg = response_json["choices"][0]["message"]
+
+        # 2) Check if function call is requested
+        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"])
+
+            # 3) Pass function result back to OpenAI for final answer
+            followup_messages = [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are an AI agent helping users analyze a CSV file."
+                    ),
+                },
+                {"role": "user", "content": prompt},
+                {
+                    "role": "function",
+                    "name": func_name,
+                    "content": json.dumps(search_result),
+                }
+            ]
+
+            final_resp = requests.post(
+                "https://api.openai.com/v1/chat/completions",
+                headers=HEADERS,
+                json={
+                    "model": "gpt-3.5-turbo-1106",  # or "gpt-4-1106-preview"
+                    "messages": followup_messages,
+                    "temperature": 0,
+                    "max_tokens": 1500,
+                },
+                timeout=60,
+            )
+            final_resp.raise_for_status()
+            answer = final_resp.json()["choices"][0]["message"]["content"]
+
+            st.subheader("✅ Agent Answer")
+            st.markdown(answer)
+            st.subheader("📊 Filtered CSV Data")
+            st.json(search_result)
         else:
-            st.subheader("🔍 Raw Output")
-            st.text(answer)
+            # No function call: model answered directly
+            st.subheader("✅ Agent Answer")
+            st.markdown(msg["content"])