Merge remote-tracking branch 'upstream/hotfixes' into release

process-intelligence-solutions · Apr 22, 2024 · f53d01d · f53d01d
2 parents 68dd348 + 928ecf5
commit f53d01d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 17 deletions.
diff --git a/pm4py/algo/querying/llm/connectors/openai.py b/pm4py/algo/querying/llm/connectors/openai.py
@@ -23,6 +23,7 @@
 
 
 class Parameters(Enum):
+    API_URL = "api_url"
     API_KEY = "api_key"
     OPENAI_MODEL = "openai_model"
     IMAGE_PATH = "image_path"
@@ -39,8 +40,17 @@ def apply(prompt: str, parameters: Optional[Dict[Any, Any]] = None) -> str:
     if parameters is None:
         parameters = {}
 
-    api_key = exec_utils.get_param_value(Parameters.API_KEY, parameters, constants.OPENAI_API_KEY)
     image_path = exec_utils.get_param_value(Parameters.IMAGE_PATH, parameters, None)
+    api_key = exec_utils.get_param_value(Parameters.API_KEY, parameters, constants.OPENAI_API_KEY)
+    api_url = exec_utils.get_param_value(Parameters.API_URL, parameters, None)
+    simple_content_specification = image_path is None
+
+    if api_url is None:
+        api_url = constants.OPENAI_API_URL
+    else:
+        if not api_url.endswith("/"):
+            api_url += "/"
+
     model = exec_utils.get_param_value(Parameters.OPENAI_MODEL, parameters,
                                        constants.OPENAI_DEFAULT_MODEL if image_path is None else constants.OPENAI_DEFAULT_VISION_MODEL)
 
@@ -49,28 +59,24 @@ def apply(prompt: str, parameters: Optional[Dict[Any, Any]] = None) -> str:
         "Authorization": f"Bearer {api_key}"
     }
 
+    messages = []
+    if simple_content_specification:
+        messages.append({"role": "user", "content": prompt})
+    else:
+        messages.append({"role": "user", "content": [{"type": "text", "text": prompt}]})
+
     payload = {
-        "model": model,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": prompt
-                    }
-                ]
-            }
-        ]
+        "model": model
     }
 
     if image_path is not None:
         base64_image = encode_image(image_path)
-        payload["messages"][0]["content"].append(
-            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image} "}})
+        messages[0]["content"].append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image} "}})
         payload["max_tokens"] = 4096
 
-    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload).json()
+    payload["messages"] = messages
+
+    response = requests.post(api_url+"chat/completions", headers=headers, json=payload).json()
 
     if "error" in response:
         # raise an exception when the request fails, with the provided message

diff --git a/pm4py/llm.py b/pm4py/llm.py
@@ -27,13 +27,14 @@
 from pm4py.objects.petri_net.obj import PetriNet, Marking
 
 
-def openai_query(prompt: str, api_key: Optional[str] = None, openai_model: Optional[str] = None, **kwargs) -> str:
+def openai_query(prompt: str, api_key: Optional[str] = None, openai_model: Optional[str] = None, api_url: Optional[str] = None, **kwargs) -> str:
     """
     Executes the provided prompt, obtaining the answer from the OpenAI APIs.
 
     :param prompt: prompt that should be executed
     :param api_key: OpenAI API key
     :param openai_model: OpenAI model to be used (default: gpt-3.5-turbo)
+    :param api_url: OpenAI API URL
     :rtype: ``str``
 
     .. code-block:: python3
@@ -44,6 +45,8 @@ def openai_query(prompt: str, api_key: Optional[str] = None, openai_model: Optio
         print(resp)
     """
     parameters = copy(kwargs) if kwargs is not None else {}
+    if api_url is not None:
+        parameters["api_url"] = api_url
     if api_key is not None:
         parameters["api_key"] = api_key
     if openai_model is not None:

diff --git a/pm4py/util/constants.py b/pm4py/util/constants.py
@@ -128,6 +128,7 @@ def get_default_is_aware_enabled():
 
 OPENAI_MAX_LEN = int(get_param_from_env("PM4PY_OPENAI_MAX_LEN", "10000"))
 OPENAI_API_KEY = get_param_from_env("PM4PY_OPENAI_API_KEY", None)
+OPENAI_API_URL = get_param_from_env("PM4PY_OPENAI_API_URL", "https://api.openai.com/v1/")
 OPENAI_DEFAULT_MODEL = get_param_from_env("PM4PY_OPENAI_DEFAULT_MODEL", "gpt-4-turbo-preview")
 OPENAI_DEFAULT_VISION_MODEL = get_param_from_env("PM4PY_OPENAI_DEFAULT_VISION_MODEL", "gpt-4-vision-preview")
 OPENAI_DEFAULT_STT_MODEL = get_param_from_env("PM4PY_OPENAI_DEFAULT_STT_MODEL", "whisper-1")