[OPENAI]: Fix chat completion response and memory leak

Jaseci-Labs · Dec 6, 2023 · 18813cc · 18813cc
1 parent f094e0b
commit 18813cc
Showing 1 changed file with 31 additions and 14 deletions.
diff --git a/jaseci_ai_kit/jac_misc/jac_misc/openai/main.py b/jaseci_ai_kit/jac_misc/jac_misc/openai/main.py
@@ -8,6 +8,7 @@
 
 Requires OpenAI API Key
 """
+from enum import Enum
 from jaseci.jsorc.live_actions import jaseci_action
 from jaseci.utils.utils import logger
 from openai import OpenAI
@@ -17,6 +18,12 @@
 openai_client = None
 
 
+class SetupResponse(int, Enum):
+    SUCCESS = 1
+    FAILED = 0
+    IGNORED = 2
+
+
 def client() -> OpenAI:
     if openai_client:
         return openai_client
@@ -51,21 +58,31 @@ def setup(
     """
     try:
         global openai_client
-        openai_client = OpenAI(
-            api_key=api_key,
-            organization=organization,
-            base_url=base_url,
-            timeout=timeout,
-            max_retries=max_retries,
-            default_headers=default_headers,
-            default_query=default_query,
-            http_client=http_client,
-            _strict_response_validation=_strict_response_validation,
-        )
-        return True
+        if not openai_client:
+            openai_client = OpenAI(
+                api_key=api_key,
+                organization=organization,
+                base_url=base_url,
+                timeout=timeout,
+                max_retries=max_retries,
+                default_headers=default_headers,
+                default_query=default_query,
+                http_client=http_client,
+                _strict_response_validation=_strict_response_validation,
+            )
+            return SetupResponse.SUCCESS
+        return SetupResponse.IGNORED
     except Exception:
         logger.error("Error occured during initialization!")
-        return False
+        return SetupResponse.FAILED
+
+
+@jaseci_action(act_group=["openai"], allow_remote=True)
+def close():
+    global openai_client
+    if openai_client:
+        openai_client.close()
+        openai_client = None
 
 
 @jaseci_action(act_group=["openai"], allow_remote=True)
@@ -192,7 +209,7 @@ def chat(
         frequency_penalty=frequency_penalty,
         **kwargs,
     )
-    response = [x.message for x in response.choices]
+    response = [x.message.dict() for x in response.choices]
     return response