update ai util

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)