Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,9 +6,10 @@ os.environ["STREAMLIT_HOME"] = "/tmp"
 import streamlit as st
 import pandas as pd
 import json
-import openai
+from openai import OpenAI
 
-openai.api_key = os.getenv("OPENAI_API_KEY")
+# Initialize OpenAI client
+client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
 PROMPT_INSTRUCTIONS_TEXT = """
 You are a forensic auditor AI with deep domain expertise and a sharp eye for irregularities. Your job is to identify **anomalies** in financial transaction data.
@@ -45,19 +46,21 @@ Output your findings in this JSON format:
 """
 
 def query_openai(prompt: str) -> dict:
-    response = openai.ChatCompletion.create(
-        model="gpt-4-turbo",
-        messages=[
-            {"role": "system", "content": "You analyze financial transactions for anomalies."},
-            {"role": "user", "content": prompt}
-        ],
-        temperature=0.2,
-        max_tokens=2048
-    )
     try:
-        return json.loads(response.choices[0].message["content"])
+        response = client.chat.completions.create(
+            model="gpt-4-turbo",
+            messages=[
+                {"role": "system", "content": "You analyze financial transactions for anomalies."},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.2,
+            max_tokens=2048
+        )
+        return json.loads(response.choices[0].message.content)
     except json.JSONDecodeError:
         return {"error": "Failed to parse JSON from LLM response."}
+    except Exception as e:
+        return {"error": str(e)}
 
 st.set_page_config(page_title="LLM Financial Anomaly Detector", layout="wide")
 st.title("LLM-Powered Financial Anomaly Detector")
@@ -80,8 +83,8 @@ if uploaded_file:
                 data_subset = df.to_dict(orient="records")
                 focused_values = df[selected_column].dropna().unique().tolist()
                 
-                # Add a section to prompt highlighting selected column
-                column_focus_note = f"\n\nFocus especially on anomalies involving the column `{selected_column}`. The unique values in this column are:\n{json.dumps(focused_values[:30])[:1000]}\n\n"
+                # Highlight selected column in the prompt
+                column_focus_note = f"\n\nFocus especially on anomalies involving the column `{selected_column}`. Here are some example values:\n{json.dumps(focused_values[:30])[:1000]}\n\n"
 
                 prompt = PROMPT_INSTRUCTIONS_TEXT + column_focus_note + "DATA:\n" + json.dumps(data_subset[:100], indent=2)
                 result = query_openai(prompt)
@@ -91,7 +94,8 @@ if uploaded_file:
                     st.dataframe(pd.json_normalize(result["anomalies"]), use_container_width=True)
                 else:
                     st.warning("No anomalies found or the model response was invalid.")
+                    st.json(result)
     except Exception as e:
         st.error(f"Could not read CSV. Error: {e}")
 else:
-    st.info("Please upload a CSV file to begin.")
+    st.info("Please upload a CSV file to begin.")