diff --git a/jaseci_ai_kit/jac_misc/jac_misc/openai/main.py b/jaseci_ai_kit/jac_misc/jac_misc/openai/main.py index c31451718f..c7a014e862 100644 --- 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