ASUllmAPI 2.0.5__tar.gz → 2.0.7__tar.gz

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

@@ -6,19 +6,26 @@ __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,
-                 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,
+    def __init__(self, access_token: str = None, action: str = None, api_url: str = "",
+                 enhance_prompt_timezone: str = None, enhance_prompt_time: bool = None,
+                 enhance_prompt_date: bool = None, enhance_prompt_verbosity: str = None,
+                 enable_history: bool = False, 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 = "opensearch",
+                 search_expr: str = None, search_output_fields: List[str] = None, search_reranker: bool = False,
                  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_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
         self.enable_history = enable_history
         self.enable_search = enable_search
+        self.enhance_prompt_timezone = enhance_prompt_timezone
+        self.enhance_prompt_time = enhance_prompt_time
+        self.enhance_prompt_date = enhance_prompt_date
+        self.enhance_prompt_verbosity = enhance_prompt_verbosity
         self.history = history
         self.model_max_tokens = model_max_tokens
         self.model_temperature = model_temperature
@@ -30,11 +37,15 @@ class ModelConfig:
         self.response_format_type = response_format_type
         self.search_collection = search_collection
         self.search_db_type = search_db_type
+        self.search_expr = search_expr
         self.search_output_fields = search_output_fields
+        self.search_prompt = search_prompt
+        self.search_prompt_mode = search_prompt_mode
+        self.search_reranker = search_reranker
+        self.search_retrieval_type = search_retrieval_type
         self.search_source_name = search_source_name
         self.search_tags = search_tags
         self.search_top_k = search_top_k
-        self.search_retrieval_type = search_retrieval_type
         self.semantic_caching = semantic_caching
         self.session_id = session_id
         self.system_prompt = system_prompt
@@ -76,6 +87,8 @@ class ModelConfig:
             payload["model_provider"] = self.provider
         if self.project_id:
             payload["project_id"] = self.project_id
+        if self.prompt_enhancers:
+            payload["enhance_prompt"] = self.prompt_enhancers
         if self.response_format:
             payload["response_format"] = self.response_format
         if self.search_params:
@@ -102,6 +115,22 @@ class ModelConfig:
             model_params["system_prompt"] = self.system_prompt
         return model_params
 
+    @property
+    def prompt_enhancers(self):
+        self.__validate_access()
+
+        prompt_enhancer_params = {}
+        if self.enhance_prompt_date is not None:
+            prompt_enhancer_params["date"] = self.enhance_prompt_date
+        if self.enhance_prompt_time is not None:
+            prompt_enhancer_params["time"] = self.enhance_prompt_time
+        if self.enhance_prompt_timezone is not None:
+            prompt_enhancer_params["timezone"] = self.enhance_prompt_timezone
+        if self.enhance_prompt_verbosity is not None:
+            prompt_enhancer_params["verbosity"] = self.enhance_prompt_verbosity
+
+        return prompt_enhancer_params
+
     @property
     def search_params(self):
         self.__validate_access()
@@ -113,8 +142,16 @@ 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_expr is not None:
+                search_params["expr"] = self.search_expr
+            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_reranker is not None:
+                search_params["reranker"] = self.search_reranker
             if self.search_tags is not None:
                 search_params["tags"] = self.search_tags
             if self.search_output_fields is not None:

{asullmapi-2.0.5 → asullmapi-2.0.7}/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.7}/ASUllmAPI.egg-info/PKG-INFO

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

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

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