ASUllmAPI 2.0.7__tar.gz → 2.0.11__tar.gz

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

@@ -7,19 +7,24 @@ __author__ = ['swliu', 'vshourie']
 class ModelConfig:
 
     def __init__(self, access_token: str = None, action: str = None, api_url: str = "",
+                 chat_upload_images: List[str] = None, chat_upload_docs: List[str] = None,
                  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,
+                 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_expr: str = None, search_output_fields: List[str] = None, rerank: bool = False,
+                 search_reranker_model: str = None, search_reranker_provider: 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):
+                 semantic_caching: bool = False, session_id: str = None, system_prompt: str = None,
+                 query_id: str = None):
         self.access_token = access_token
         self.action = action
         self.api_url = api_url
+        self.chat_upload_images = chat_upload_images
+        self.chat_upload_docs = chat_upload_docs
         self.enable_history = enable_history
         self.enable_search = enable_search
         self.enhance_prompt_timezone = enhance_prompt_timezone
@@ -41,7 +46,9 @@ class ModelConfig:
         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.rerank = rerank
+        self.search_reranker_model = search_reranker_model
+        self.search_reranker_provider = search_reranker_provider
         self.search_retrieval_type = search_retrieval_type
         self.search_source_name = search_source_name
         self.search_tags = search_tags
@@ -49,6 +56,7 @@ class ModelConfig:
         self.semantic_caching = semantic_caching
         self.session_id = session_id
         self.system_prompt = system_prompt
+        self.query_id = query_id
 
         self.__validate_access()
 
@@ -73,6 +81,8 @@ class ModelConfig:
         payload = {"query": query}
         if self.action:
             payload["action"] = self.action
+        if self.chat_upload:
+            payload["chat_upload"] = self.chat_upload
         if self.enable_history is not None:
             payload["enable_history"] = self.enable_history
         if self.enable_search is not None:
@@ -95,9 +105,20 @@ class ModelConfig:
             payload["search_params"] = self.search_params
         if self.semantic_caching:
             payload["semantic_caching"] = self.semantic_caching
+        if self.query_id:
+            payload["query_id"] = self.query_id
 
         return payload
 
+    @property
+    def chat_upload(self):
+        chat_upload_payload = {}
+        if self.chat_upload_docs:
+            chat_upload_payload['docs'] = self.chat_upload_docs
+        if self.chat_upload_images:
+            chat_upload_payload['images'] = self.chat_upload_images
+        return chat_upload_payload
+
     @property
     def model_params(self):
         self.__validate_access()
@@ -150,8 +171,12 @@ class ModelConfig:
                 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.rerank is not None:
+                search_params["rerank"] = self.rerank
+            if self.search_reranker_model is not None:
+                search_params["reranker_model"] = self.search_reranker_model
+            if self.search_reranker_provider is not None:
+                search_params["reranker_provider"] = self.search_reranker_provider
             if self.search_tags is not None:
                 search_params["tags"] = self.search_tags
             if self.search_output_fields is not None:

{asullmapi-2.0.7 → asullmapi-2.0.11}/ASUllmAPI/web_socket.py

@@ -123,11 +123,13 @@ async def interact_with_websocket(uri: str, queue: asyncio.Queue,
                     queue.task_done()
             except KeyError:
                 logging.error(f"Question ID {qid} does not exist in the queue. Exiting...")
+                queue.task_done()
                 return
             finally:
                 # prevent any further retries if at error limit.
                 if error_ct == error_threshold:
                     error_ct = 0
+                    queue.task_done()
                 time.sleep(reconnect_timeout_secs)
     # END - WEBSOCKET LOOP
     logging.info("WebSocket connection closed. Queue appears to be empty...")
@@ -139,6 +141,14 @@ async def batch_query_llm_socket(model: ModelConfig, queries: Dict[Union[str, in
                                  ws_timeout_min: int = 8,
                                  error_threshold: int = 100,
                                  reconnect_timeout_secs: int = 2) -> Dict[Union[str, int], Dict[str, Any]]:
+    return await async_batch_query_llm_socket(model, queries, max_concurrent_tasks, ws_timeout_min,
+                                              error_threshold, reconnect_timeout_secs)
+
+
+async def async_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)
@@ -162,12 +172,18 @@ 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]:
+    return await async_query_llm_socket(model, query, ws_timeout_min, error_threshold, reconnect_timeout_secs)
+
+
+async def async_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.7 → asullmapi-2.0.11}/ASUllmAPI.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: ASUllmAPI
-Version: 2.0.7
+Version: 2.0.11
 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
@@ -15,6 +15,7 @@ Requires-Dist: requests
 Requires-Dist: tqdm
 Requires-Dist: websockets
 Requires-Dist: certifi
+Dynamic: license-file
 
 # 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.7 → asullmapi-2.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: ASUllmAPI
-Version: 2.0.7
+Version: 2.0.11
 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
@@ -15,6 +15,7 @@ Requires-Dist: requests
 Requires-Dist: tqdm
 Requires-Dist: websockets
 Requires-Dist: certifi
+Dynamic: license-file
 
 # 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.7 → asullmapi-2.0.11}/pyproject.toml

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