ASUllmAPI 2.0.5__tar.gz → 2.0.6__tar.gz

{asullmapi-2.0.5 → asullmapi-2.0.6}/ASUllmAPI/model_config.py

@@ -7,13 +7,14 @@ __author__ = ['swliu', 'vshourie']
 class ModelConfig:
 
     def __init__(self, access_token: str = None, action: str = None, api_url: str = "", enable_history: bool = False,
-                 enable_search: bool = True, history: List[Dict[str, str]] = None, model_max_tokens: int = None,
+                 enable_search: bool = True,  history: List[Dict[str, str]] = None, model_max_tokens: int = None,
                  model_temperature: float = None, model_top_k: int = None, model_top_p: int = None, name: str = "",
                  project_id: str = None, provider: str = "", response_format_type: str = "",
-                 search_collection: str = "asu", search_db_type: str = "milvus", search_output_fields: List[str] = None,
-                 search_retrieval_type: str = "chunk", search_source_name: List[str] = None, search_tags: list = None,
-                 search_top_k: int = 3, semantic_caching: bool = False, session_id: str = None,
-                 system_prompt: str = None):
+                 search_collection: str = "asu", search_db_type: str = "opensearch",
+                 search_output_fields: List[str] = None, search_retrieval_type: str = "chunk",
+                 search_source_name: List[str] = None, search_tags: list = None, search_top_k: int = 3,
+                 search_prompt_mode: str = "unrestricted", search_prompt: str = None, semantic_caching: bool = False,
+                 session_id: str = None, system_prompt: str = None):
         self.access_token = access_token
         self.action = action
         self.api_url = api_url
@@ -35,6 +36,8 @@ class ModelConfig:
         self.search_tags = search_tags
         self.search_top_k = search_top_k
         self.search_retrieval_type = search_retrieval_type
+        self.search_prompt_mode = search_prompt_mode
+        self.search_prompt = search_prompt
         self.semantic_caching = semantic_caching
         self.session_id = session_id
         self.system_prompt = system_prompt
@@ -113,6 +116,10 @@ class ModelConfig:
             search_params["top_k"] = self.search_top_k
             search_params["retrieval_type"] = self.search_retrieval_type
 
+            if self.search_prompt_mode is not None:
+                search_params["prompt_mode"] = self.search_prompt_mode
+            if self.search_prompt is not None:
+                search_params["search_prompt"] = self.search_prompt
             if self.search_source_name is not None:
                 search_params["source_name"] = self.search_source_name
             if self.search_tags is not None:

{asullmapi-2.0.5 → asullmapi-2.0.6}/ASUllmAPI/web_socket.py

@@ -99,29 +99,36 @@ async def interact_with_websocket(uri: str, queue: asyncio.Queue,
                 # 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...")
+                logging.error(f"Error {error_ct} on Question ID {qid} stream timeout: resetting connection...")
             elif isinstance(exc, ResponseDataError):
-                logging.error(f"Error {error_ct} on {qid}: invalid response from endpoint.\n"
+                logging.error(f"Error {error_ct} on Question ID {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 "
+                logging.error(f"Error {error_ct} on Question ID {qid}: WebSocket connection closed on "
                               f"query ID {qid}. Reopening...")
+            elif isinstance(exc, websockets.exceptions.InvalidStatusCode):
+                logging.error(f"Error {error_ct} on Question ID {qid}: Server rejected the connection. "
+                              f"Check URI again.")
             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
+                logging.error(f"Error {error_ct} on Question ID {qid}: {traceback.format_exc()}")
+
+            try:
+                # 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()
+            except KeyError:
+                logging.error(f"Question ID {qid} does not exist in the queue. Exiting...")
+                return
+            finally:
+                # prevent any further retries if at error limit.
+                if error_ct == error_threshold:
+                    error_ct = 0
+                time.sleep(reconnect_timeout_secs)
     # END - WEBSOCKET LOOP
     logging.info("WebSocket connection closed. Queue appears to be empty...")
 

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

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

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

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