Create openai_agent.py

openai_agent.py

@@ -0,0 +1,23 @@
+import openai
+
+class OpenAIAgent:
+    def __init__(self):
+        # Make sure to set your OpenAI API key in your environment variables
+        openai.api_key = "c496d9094ba54ddb9d66eeeb35a6196f"
+
+    def get_insights(self, df, prompt):
+        # Prepare the data for the API call
+        data_sample = df.head(5).to_string()
+        
+        messages = [
+            {"role": "system", "content": "You are a data analyst assistant."},
+            {"role": "user", "content": f"Here's a sample of my data:\n{data_sample}\n\nBased on this data, {prompt}"}
+        ]
+        
+        response = openai.ChatCompletion.create(
+            model="gpt 4o",
+            messages=messages,
+            max_tokens=150
+        )
+        
+        return response.choices[0].message['content']