ASUllmAPI 2.0.1__tar.gz → 2.0.3__tar.gz

{asullmapi-2.0.1 → asullmapi-2.0.3}/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.3/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.3/ASUllmAPI/web_socket.py

@@ -0,0 +1,105 @@
+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]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ASUllmAPI
-Version: 2.0.1
+Version: 2.0.3
 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
@@ -13,6 +13,8 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: tqdm
+Requires-Dist: websockets
+Requires-Dist: certifi
 
 # ASU LLM API
 This package allows individuals at Arizona State University to access ASU GPT through API. You will need API access token, API endpoints in order to use this package.

asullmapi-2.0.3/ASUllmAPI.egg-info/requires.txt

@@ -0,0 +1,4 @@
+requests
+tqdm
+websockets
+certifi

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ASUllmAPI
-Version: 2.0.1
+Version: 2.0.3
 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
@@ -13,6 +13,8 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: tqdm
+Requires-Dist: websockets
+Requires-Dist: certifi
 
 # ASU LLM API
 This package allows individuals at Arizona State University to access ASU GPT through API. You will need API access token, API endpoints in order to use this package.

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

@@ -1,6 +1,6 @@
 [project]
 name = "ASUllmAPI"
-version = "2.0.1"
+version = "2.0.3"
 authors = [
     { name="Stella Wenxing Liu", email="stellawenxingliu@gmail.com" },
     { name="Varun Shourie", email="svarun195@gmail.com" }
@@ -15,7 +15,9 @@ classifiers = [
 ]
 dependencies = [
     "requests",
-    "tqdm"
+    "tqdm",
+    "websockets",
+    "certifi"
 ]
 
 [project.urls]

asullmapi-2.0.1/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.1/ASUllmAPI/web_socket.py

@@ -1,91 +0,0 @@
-import asyncio
-import json
-from typing import Dict
-
-import websockets
-
-from .model_config import ModelConfig
-from .utils import load_json_buffer, begin_task_execution
-
-
-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) 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():
-                        if parsed_response['message'] == 'Forbidden':
-                            # Embed entire Forbidden message payload and leave response blank.
-                            payload[qid].update(parsed_response)
-                            break
-                        else:
-                            raise RuntimeWarning(f"Unknown ASU LLM endpoint edge case detected: `message` as key in "
-                                                 f"parsed response: {parsed_response}.")
-
-                    # 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
-                    # Edge case: user gets a message to contact the administrator after an error
-                    elif 'error' in parsed_response.keys():
-                        payload[qid]["error"] = parsed_response["error"]
-                    # Unknown edge case: json getting returned that is not a `message`, `response`, `error`
-                    # (not observed yet)
-                    else:
-                        raise RuntimeWarning(f"Unknown ASU LLM endpoint edge case detected: JSON parsed but data "
-                                             f"does not contain a message or response field.")
-
-                # 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

asullmapi-2.0.1/ASUllmAPI.egg-info/requires.txt

@@ -1,2 +0,0 @@
-requests
-tqdm