ASUllmAPI 2.0.3__tar.gz → 2.0.5__tar.gz

{asullmapi-2.0.3 → asullmapi-2.0.5}/ASUllmAPI/utils.py

@@ -3,6 +3,7 @@ from concurrent.futures import ThreadPoolExecutor
 import time
 import json
 import asyncio
+import sys
 
 
 def begin_task_execution(async_func):
@@ -13,17 +14,31 @@ def begin_task_execution(async_func):
     If an event loop doesn't exist, it reverts back to existing asyncio logic.
     """
     def wrap(*args, **kwargs):
-        try:
-            asyncio.get_running_loop()  # Triggers RuntimeError if no running event loop
-            # Create a separate thread so we can block before returning
+        # It is safer to use a if/else statement than try/except since
+        # the underlying code we wrap our function around can also raise RuntimeErrors.
+        # When this happens, multiple asyncio.run() executions occur, which is dangerous.
+        if is_jupyter():
             with ThreadPoolExecutor(1) as pool:
                 result = pool.submit(lambda: asyncio.run(async_func(*args, **kwargs))).result()
-        except RuntimeError:
+        else:
             result = asyncio.run(async_func(*args, **kwargs))
         return result
     return wrap
 
 
+def is_jupyter():
+    try:
+        # Check if 'IPython' is in sys.modules
+        if 'IPython' in sys.modules:
+            from IPython import get_ipython
+            # Check if we're in an IPython environment
+            if get_ipython() is not None:
+                return True
+    except ImportError:
+        pass
+    return False
+
+
 def load_json_buffer(string):
     try:
         return json.loads(string)

asullmapi-2.0.5/ASUllmAPI/web_socket.py

@@ -0,0 +1,166 @@
+import asyncio
+import json
+import logging
+import traceback
+from typing import Dict, Any, Union
+import ssl
+import time
+
+import websockets
+import certifi
+
+from .model_config import ModelConfig
+from .utils import load_json_buffer, begin_task_execution
+
+SSL_CONTEXT = ssl.create_default_context(cafile=certifi.where())
+END_OF_STREAM = '<EOS>'
+DEFAULT_RESPONSE = {"response": "", "success": 0}
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
+
+
+class ResponseDataError(Exception):
+    def __init__(self, message="An error message has been sent by the endpoint."):
+        super().__init__(message)
+        self.message = message
+
+
+async def interact_with_websocket(uri: str, queue: asyncio.Queue,
+                                  response_payloads: Dict[Union[str, int], Dict[str, Any]],
+                                  ws_timeout_min: int = 8,
+                                  error_threshold: int = 100,
+                                  reconnect_timeout_secs: int = 2):
+    """
+    :param uri: WebSocket URI
+    :param queue: An instantiated asyncio Queue with your dataset fully inserted.
+    :param response_payloads: A dictionary with base responses set up for processing.
+    :param ws_timeout_min: The number of minutes before a connection to the WebSocket is terminated on the async task.
+    :param error_threshold: The number of errors permissible for a query in the WebSocket before the program
+    moves to the next question.
+    :param reconnect_timeout_secs: The number of seconds before another connection is instantiated.
+    :return: the `response_payloads` object initially passed into the function
+    """
+    error_ct = 0
+    ws_timeout = ws_timeout_min * 60
+    qid = None
+    tmp_input_payload = None
+
+    # START - WEBSOCKET LOOP
+    while not (queue.empty() and error_ct == 0):
+        connection_start_time = time.time()
+        try:
+            async with websockets.connect(uri, ssl=SSL_CONTEXT) as ws:
+                # START - QUERY QUEUE LOOP
+                while not (queue.empty() and error_ct == 0):
+                    if error_ct == 0:
+                        qid, tmp_input_payload = await queue.get()
+
+                    # To prevent unnecessary timeouts DURING a query, we can
+                    # cautiously opt out of the connection at a user specified threshold.
+                    query_time = time.time()
+                    connection_time = query_time - connection_start_time
+                    if connection_time >= ws_timeout:
+                        raise asyncio.TimeoutError()
+
+                    if tmp_input_payload["query"] != "":
+                        await ws.send(json.dumps(tmp_input_payload))
+                        # START - QUERY CHUNK LOOP
+                        while True:
+                            # First chunk will typically take 29 seconds of time for every model
+                            # Remaining chunks can take up to 9 minutes to be received.
+                            # First chunk will take more time with GeminiPro - long time to load in comparison
+                            # to other models
+                            response = await asyncio.wait_for(ws.recv(), ws_timeout)
+
+                            parsed_response = load_json_buffer(response)
+                            if isinstance(parsed_response, dict):
+                                if 'response' in parsed_response.keys():
+                                    cleaned = parsed_response["response"].replace(END_OF_STREAM, "")
+                                    response_payloads[qid]["response"] += cleaned
+                                    if 'metadata' in parsed_response.keys() \
+                                            or END_OF_STREAM in parsed_response['response']:
+                                        response_payloads[qid]["metadata"] = parsed_response["metadata"]
+                                        response_payloads[qid]["success"] = 1
+                                        error_ct = 0
+                                        break
+                                else:
+                                    # Connection ID expires - first chunk took more than 29 seconds.
+                                    response_payloads[qid].update(parsed_response)
+                                    raise ResponseDataError()
+                            else:
+                                response_payloads[qid]["response"] += response.replace(END_OF_STREAM, "")
+                                if END_OF_STREAM in response:
+                                    response_payloads[qid]["success"] = 1
+                                    error_ct = 0
+                                    break
+                        # END - QUERY CHUNK LOOP
+                        if error_ct == 0:
+                            logging.info(f"Query ID {qid} completed... Message: {response_payloads[qid]['response']}")
+                    queue.task_done()
+                # END - QUERY QUEUE LOOP
+        except (asyncio.TimeoutError, websockets.ConnectionClosed, Exception) as exc:
+            if isinstance(exc, asyncio.TimeoutError):
+                logging.error(f"Error {error_ct} on {qid} stream timeout: resetting connection...")
+            elif isinstance(exc, ResponseDataError):
+                logging.error(f"Error {error_ct} on {qid}: invalid response from endpoint.\n"
+                              f"{response_payloads[qid]}")
+            elif isinstance(exc, websockets.ConnectionClosed):
+                logging.error(f"Error {error_ct} on {qid}: WebSocket connection closed on "
+                              f"query ID {qid}. Reopening...")
+            else:
+                logging.error(f"Error {error_ct} on {qid}: {traceback.format_exc()}")
+
+            time.sleep(reconnect_timeout_secs)
+            # If the query is already complete, we don't want to increment the error count
+            if response_payloads[qid]["success"] == 0:
+                error_ct += 1
+                # Reset buffer stream so that you don't get messed by pre-existing data.
+                response_payloads[qid]["response"] = ""
+            else:
+                logging.info(f"Query ID {qid} completed...")
+                queue.task_done()
+
+            # prevent any further retries if at error limit.
+            if error_ct == error_threshold:
+                error_ct = 0
+    # END - WEBSOCKET LOOP
+    logging.info("WebSocket connection closed. Queue appears to be empty...")
+
+
+@begin_task_execution
+async def batch_query_llm_socket(model: ModelConfig, queries: Dict[Union[str, int], str],
+                                 max_concurrent_tasks: int = 3,
+                                 ws_timeout_min: int = 8,
+                                 error_threshold: int = 100,
+                                 reconnect_timeout_secs: int = 2) -> Dict[Union[str, int], Dict[str, Any]]:
+    payloads = {}
+    for qid, message in queries.items():
+        payloads[qid] = model.compute_payload(message)
+
+    response_payloads = {qid: DEFAULT_RESPONSE.copy() for qid in queries}
+
+    if len(response_payloads) > 0:
+        queue = asyncio.Queue()
+        for qid, payload in payloads.items():
+            await queue.put((qid, payload))
+        tasks = [interact_with_websocket(model.api_url, queue, response_payloads,
+                                         ws_timeout_min, error_threshold, reconnect_timeout_secs)
+                 for _ in range(max_concurrent_tasks)]
+        await asyncio.gather(*tasks)
+
+    return response_payloads
+
+
+@begin_task_execution
+async def query_llm_socket(model: ModelConfig, query: str,
+                           ws_timeout_min: int = 8,
+                           error_threshold: int = 100,
+                           reconnect_timeout_secs: int = 2) -> Dict[str, Any]:
+    response_payloads = {0: DEFAULT_RESPONSE.copy()}
+    tmp_payload = model.compute_payload(query=query)
+    if query != "":
+        queue = asyncio.Queue()
+        await queue.put((0, tmp_payload))
+
+        await interact_with_websocket(model.api_url, queue, response_payloads, ws_timeout_min,
+                                      error_threshold, reconnect_timeout_secs)
+    return response_payloads[0]

{asullmapi-2.0.3 → asullmapi-2.0.5}/ASUllmAPI.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ASUllmAPI
-Version: 2.0.3
+Version: 2.0.5
 Summary: A simple python package to facilitate connection to ASU LLM API
 Author-email: Stella Wenxing Liu <stellawenxingliu@gmail.com>, Varun Shourie <svarun195@gmail.com>
 Project-URL: Homepage, https://github.com/ASU/aiml-ssmdv-student-support-ml-data-visualization

{asullmapi-2.0.3 → asullmapi-2.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ASUllmAPI
-Version: 2.0.3
+Version: 2.0.5
 Summary: A simple python package to facilitate connection to ASU LLM API
 Author-email: Stella Wenxing Liu <stellawenxingliu@gmail.com>, Varun Shourie <svarun195@gmail.com>
 Project-URL: Homepage, https://github.com/ASU/aiml-ssmdv-student-support-ml-data-visualization

{asullmapi-2.0.3 → asullmapi-2.0.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ASUllmAPI"
-version = "2.0.3"
+version = "2.0.5"
 authors = [
     { name="Stella Wenxing Liu", email="stellawenxingliu@gmail.com" },
     { name="Varun Shourie", email="svarun195@gmail.com" }

asullmapi-2.0.3/ASUllmAPI/web_socket.py

@@ -1,105 +0,0 @@
-import asyncio
-import json
-import logging
-from typing import Dict, Any, Union
-import ssl
-
-import websockets
-import certifi
-
-from .model_config import ModelConfig
-from .utils import load_json_buffer, begin_task_execution, split_dict_into_chunks
-
-SSL_CONTEXT = ssl.create_default_context(cafile=certifi.where())
-END_OF_STREAM = '<EOS>'
-DEFAULT_RESPONSE = {"response": ""}
-logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
-
-
-async def interact_with_websocket(uri: str, query_payloads: Dict[Union[str, int], Dict[str, Any]]):
-    response_payloads = {qid: {"response": ""} for qid in query_payloads}
-    payload_iterator = iter(query_payloads.items())
-    tmp_qid, tmp_input_payload = next(payload_iterator, (None, None))
-    retry = False
-    ws_ready_for_closure = False
-    error_ct = 0
-
-    # WEBSOCKET LOOP: Indefinitely opens a new connection when a new one is closed in the SAME asyncio task.
-    while not ws_ready_for_closure:
-        async with websockets.connect(uri, ssl=SSL_CONTEXT) as ws:
-            try:
-                if ws_ready_for_closure or tmp_qid is None:
-                    break
-                # QUERY LOOP: While there are queries remaining to send to the server.
-                while tmp_qid is not None:
-                    response_payloads[tmp_qid]["success"] = 0
-                    if not retry:
-                        error_ct = 0
-                    await ws.send(json.dumps(tmp_input_payload))
-                    # QUERY CHUNK LOOP: While chunks are being received by the WebSocket Client.
-                    while True:
-                        # Wait for the duration of the WebSocket connection for ANY response
-                        response = await ws.recv()
-                        parsed_response = load_json_buffer(response)
-                        # Parse response stream from WebSocket server based on every use case
-                        if isinstance(parsed_response, dict):
-                            # Intended use case: user receives JSON response.
-                            if 'response' in parsed_response.keys():
-                                cleaned = parsed_response["response"].replace(END_OF_STREAM, "")
-                                response_payloads[tmp_qid]["response"] += cleaned
-                                if 'metadata' in parsed_response.keys() or END_OF_STREAM in parsed_response['response']:
-                                    response_payloads[tmp_qid]["metadata"] = parsed_response["metadata"]
-                                    response_payloads[tmp_qid]["success"] = 1
-                                    retry = False
-                                    break
-                            # Error: user denied access to the response for some reason.
-                            else:
-                                logging.log(logging.ERROR, f"Data error: query ID {tmp_qid}: {parsed_response}")
-                                response_payloads[tmp_qid].update(parsed_response)
-                                retry = False
-                                break
-                        # Intended use case: user receives TEXT response from websocket
-                        # Data format is streamlined regardless of which response type the user selects.
-                        else:
-                            response_payloads[tmp_qid]["response"] += response.replace(END_OF_STREAM, "")
-                            if END_OF_STREAM in response:
-                                retry = False
-                                response_payloads[tmp_qid]["success"] = 1
-                                break
-                    # END - QUERY CHUNK LOOP
-                    if not retry:
-                        logging.log(level=logging.INFO, msg=f"Query ID {tmp_qid} completed...")
-                        tmp_qid, tmp_input_payload = next(payload_iterator, (None, None))
-                # END - QUERY LOOP
-                ws_ready_for_closure = True
-            except websockets.ConnectionClosed:
-                error_ct += 1
-                logging.log(level=logging.INFO, msg=f"Error {error_ct}: WebSocket connection closed on query ID "
-                                                    f"{tmp_qid}. Reopening...")
-                retry = True if error_ct <= 2 else False
-    # END - WEBSOCKET LOOP
-    logging.log(logging.INFO, msg="WebSocket connection closed.")
-    return response_payloads
-
-
-@begin_task_execution
-async def batch_query_llm_socket(model: ModelConfig, queries: Dict[Union[str, int], str],
-                                 max_concurrent_tasks: int = 3) -> Dict[Union[str, int], Dict[str, Any]]:
-    payloads = {}
-    for qid, message in queries.items():
-        payloads[qid] = model.compute_payload(message)
-    batches = split_dict_into_chunks(payloads, max_concurrent_tasks)
-    tasks = [interact_with_websocket(model.api_url, batch) for batch in batches]
-
-    results = await asyncio.gather(*tasks)
-    final_results = {}
-    for result_dict in results:
-        final_results.update(result_dict)
-    return final_results
-
-
-@begin_task_execution
-async def query_llm_socket(model: ModelConfig, query: str) -> Dict[str, Any]:
-    tmp_payload = model.compute_payload(query=query)
-    result = await interact_with_websocket(model.api_url, {0: tmp_payload})
-    return result[0]