From 4ecb7a80a7c3647f0f7969c52fbfdb7bf3ac33b1 Mon Sep 17 00:00:00 2001
From: Parish Wolfe <parishwolfe@gmail.com>
Date: Mon, 23 Sep 2024 18:48:52 -0400
Subject: [PATCH] update ai util

---
 util/ai_util.py | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/util/ai_util.py b/util/ai_util.py
index 4a19561..9173f45 100644
--- a/util/ai_util.py
+++ b/util/ai_util.py
@@ -2,23 +2,39 @@
 from os import getenv
 
 class ai_util:
-    def __init__(self, api_key=None):
-        self.api_key = api_key
-        if not self.api_key:
+    def __init__(
+        self, 
+        api_key=None,
+        model="gpt-4o",
+        temperature=0.7,
+        max_response_len=None,
+        frequency_penalty=0
+    ):
+        if not api_key:
             self.api_key = getenv("OPENAI_API_KEY")
+        self.model = model
+        self.temperature = temperature
+        self.max_response_len = max_response_len
         self.client = OpenAI()
+        self.frequency_penalty = frequency_penalty
 
     def chat(self, messages: list):
         completion = self.client.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role": "user", "content": "write a haiku about ai"}
-            ]
+            model=self.model,
+            messages=messages,
+            temperature=self.temperature,
+            max_tokens=self.max_response_len,
+            frequency_penalty=self.frequency_penalty
         )
         return completion.choices[0].message.content
 
 
 # Example usage:
 ai = ai_util()
-response = ai.chat(["What is the meaning of life?"])
+response = ai.chat(
+    [
+        { "role": "system", "content": "You are a helpful assistant.", 
+        "role": "user", "content": "Tell me about yourself" }
+    ]
+)
 print(response)
\ No newline at end of file