ASUllmAPI 2.0.2__tar.gz → 2.0.4__tar.gz

{asullmapi-2.0.2 → asullmapi-2.0.4}/ASUllmAPI/multithreading.py

@@ -1,5 +1,5 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Dict
+from typing import Dict, Union
 
 from tqdm.auto import tqdm
 
@@ -9,9 +9,13 @@ from .utils import time_api
 
 
 @time_api
-def batch_query_llm(model: ModelConfig, queries: Dict[int, str], max_threads: int,
-                    num_retry: int = 3, auto_increase_retry: bool = False,
-                    success_sleep: float = 0.0, fail_sleep: float = 1.0) -> Dict[int, dict]:
+def batch_query_llm(model: ModelConfig, 
+                    queries: Dict[Union[str, int], str], 
+                    max_threads: int,
+                    num_retry: int = 3, 
+                    auto_increase_retry: bool = False,
+                    success_sleep: float = 0.0, 
+                    fail_sleep: float = 1.0) -> Dict[Union[str, int], dict]:
     with ThreadPoolExecutor(max_workers=max_threads) as executor:
         # Submit tasks to the executor - order of return will be asynchronous.
         # If `auto_increase_retry` enabled, then scaling API backoff is implemented.

asullmapi-2.0.4/ASUllmAPI/utils.py

@@ -0,0 +1,84 @@
+from typing import Dict, List
+from concurrent.futures import ThreadPoolExecutor
+import time
+import json
+import asyncio
+
+
+def begin_task_execution(async_func):
+    """
+    See https://stackoverflow.com/a/75341431 for underlying rationale.
+    This code should allow any async function to be called synchronously by circumventing
+    Jupyter's existing event loop via the creation of a separate thread.
+    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
+            with ThreadPoolExecutor(1) as pool:
+                result = pool.submit(lambda: asyncio.run(async_func(*args, **kwargs))).result()
+        except RuntimeError:
+            result = asyncio.run(async_func(*args, **kwargs))
+        return result
+    return wrap
+
+
+def load_json_buffer(string):
+    try:
+        return json.loads(string)
+    except json.JSONDecodeError:
+        return None
+
+
+def split_dict_into_chunks(input_dict: Dict, n: int) -> List[Dict]:
+    """
+    Split a dictionary into `n` chunks.
+
+    Parameters:
+    input_dict (dict): The dictionary to split.
+    n (int): The number of chunks to split the dictionary into.
+
+    Returns:
+    List[Dict]: A list of dictionaries, where each dictionary is a chunk of the input dictionary.
+    """
+    if n <= 0:
+        raise ValueError("Number of chunks must be greater than 0")
+
+    if not input_dict:
+        return []
+
+    items = list(input_dict.items())
+    chunk_size = len(items) // n
+    remainder = len(items) % n
+
+    chunks = []
+    start = 0
+    for i in range(n):
+        end = start + chunk_size + (1 if i < remainder else 0)
+        chunk = dict(items[start:end])
+        chunks.append(chunk)
+        start = end
+
+    return chunks
+
+
+def time_api(func):
+    """
+    Decorator to measure the execution time of a function.
+    """
+
+    def time_wrapper(*args, **kwargs):
+        """
+        Passed function reference is utilized to run the function with its
+        original arguments while maintaining timing and logging functions.
+        """
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        print(f"{func.__name__} executed in {elapsed_time:.4f} seconds.")
+        return result
+
+    # We return the augmented function's reference.
+    return time_wrapper

asullmapi-2.0.4/ASUllmAPI/web_socket.py

@@ -0,0 +1,165 @@
+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:
+            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()
+            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()}")
+
+            # 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.2 → asullmapi-2.0.4}/ASUllmAPI.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ASUllmAPI
-Version: 2.0.2
+Version: 2.0.4
 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.2 → asullmapi-2.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ASUllmAPI
-Version: 2.0.2
+Version: 2.0.4
 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.2 → asullmapi-2.0.4}/pyproject.toml

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

asullmapi-2.0.2/ASUllmAPI/utils.py

@@ -1,38 +0,0 @@
-import time
-import json
-import asyncio
-
-
-def begin_task_execution(async_func):
-    def wrap(*args, **kwargs):
-        return asyncio.run(async_func(*args, **kwargs))
-
-    return wrap
-
-
-def load_json_buffer(string):
-    try:
-        return json.loads(string)
-    except json.JSONDecodeError:
-        return None
-
-
-def time_api(func):
-    """
-    Decorator to measure the execution time of a function.
-    """
-
-    def time_wrapper(*args, **kwargs):
-        """
-        Passed function reference is utilized to run the function with its
-        original arguments while maintaining timing and logging functions.
-        """
-        start_time = time.time()
-        result = func(*args, **kwargs)
-        end_time = time.time()
-        elapsed_time = end_time - start_time
-        print(f"{func.__name__} executed in {elapsed_time:.4f} seconds.")
-        return result
-
-    # We return the augmented function's reference.
-    return time_wrapper

asullmapi-2.0.2/ASUllmAPI/web_socket.py

@@ -1,89 +0,0 @@
-import asyncio
-import json
-from typing import Dict
-import ssl
-import warnings
-
-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())
-
-
-async def interact_with_websocket(uri: str, query_payload: Dict[str, str], qid: int):
-    end_of_stream = "<EOS>"
-    payload = {qid: {"response": ""}}
-
-    async with websockets.connect(uri, ssl=SSL_CONTEXT) as websocket:
-        # Send the user-provided message to the WebSocket server
-        await websocket.send(json.dumps(query_payload))
-
-        # Loop to receive messages until the server is finished
-        while True:
-            try:
-                response = await websocket.recv()
-                parsed_response = load_json_buffer(response)
-
-                if isinstance(parsed_response, dict):
-                    # Case: json/text response and user is denied entry
-                    if 'message' in parsed_response.keys() or 'error' in parsed_response.keys():
-                        # Embed entire Forbidden message payload and leave response blank.
-                        payload[qid].update(parsed_response)
-                        break
-                    # Case: json response and user is not denied entry
-                    elif 'response' in parsed_response.keys():
-                        payload[qid]["response"] += (parsed_response["response"].replace(end_of_stream, ""))
-                        if 'metadata' in parsed_response.keys() or end_of_stream in parsed_response['response']:
-                            payload[qid]["metadata"] = parsed_response["metadata"]
-                            break
-                    # Unknown edge case: json getting returned that is not a `message`, `response`, `error`
-                    # (not observed yet)
-                    else:
-                        warnings.warn(f"Unknown ASU LLM endpoint edge case detected: JSON parsed but data does not "
-                                      f"contain a message or response field.", RuntimeWarning)
-                        payload[qid].update(parsed_response)
-                        break
-                # Intended case: user asked for a response type of `text` and was not denied entry.
-                else:
-                    payload[qid]["response"] += response.replace(end_of_stream, "")
-                    if end_of_stream in response:
-                        break
-
-            except websockets.ConnectionClosed:
-                print(f"Question {qid} connection closed by the server")
-                break
-    print(f"......Question {qid} response sent by the WebSocket server.")
-    return payload
-
-
-async def limited_task(semaphore: asyncio.Semaphore, task):
-    async with semaphore:
-        return await task
-
-
-@begin_task_execution
-async def batch_query_llm_socket(model: ModelConfig, queries: Dict[int, str], max_concurrent_tasks: int = 3):
-    tasks = []
-    for qid, message in queries.items():
-        tmp_payload = model.compute_payload(message)
-        tasks.append(interact_with_websocket(model.api_url, tmp_payload, qid))
-    semaphore = asyncio.Semaphore(max_concurrent_tasks)
-    limited_tasks = [limited_task(semaphore, task) for task in tasks]
-
-    # Gather all tasks to run them concurrently and collect results
-    results = await asyncio.gather(*limited_tasks)
-    final_results = {}
-    # Desired format: {1: {"response": ..., "metadata": ...}, 2: {...}}
-    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):
-    tmp_payload = model.compute_payload(query=query)
-    result = await interact_with_websocket(model.api_url, tmp_payload, qid=0)
-    return result