abacusai 1.2.0__tar.gz → 1.2.1__tar.gz

{abacusai-1.2.0/abacusai.egg-info → abacusai-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: abacusai
-Version: 1.2.0
+Version: 1.2.1
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/__init__.py

@@ -4,4 +4,4 @@ from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/api_class/enums.py

@@ -184,6 +184,7 @@ class FileFormat(ApiEnum):
     PPTX = 'PPTX'
     PPT = 'PPT'
     HTML = 'HTML'
+    TXT = 'txt'
 
 
 class ExperimentationMode(ApiEnum):
@@ -411,6 +412,7 @@ class LLMName(ApiEnum):
     CLAUDE_V2_1 = 'CLAUDE_V2_1'
     CLAUDE_V3_OPUS = 'CLAUDE_V3_OPUS'
     CLAUDE_V3_SONNET = 'CLAUDE_V3_SONNET'
+    CLAUDE_V3_HAIKU = 'CLAUDE_V3_HAIKU'
     ABACUS_GIRAFFE = 'ABACUS_GIRAFFE'
     ABACUS_GIRAFFE_LARGE = 'ABACUS_GIRAFFE_LARGE'
     LLAMA2_CHAT = 'LLAMA2_CHAT'
@@ -419,6 +421,7 @@ class LLMName(ApiEnum):
     GEMINI_PRO = 'GEMINI_PRO'
     MIXTRAL_CHAT = 'MIXTRAL_CHAT'
     MISTRAL_MEDIUM = 'MISTRAL_MEDIUM'
+    ABACUS_SMAUG2 = 'ABACUS_SMAUG2'
 
 
 class MonitorAlertType(ApiEnum):

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/api_class/model.py

@@ -732,12 +732,14 @@ class CustomTrainedModelTrainingConfig(TrainingConfig):
         index_output_path (str): Fully qualified cloud location (GCS, S3, etc) to export snapshots of the embedding to.
         docker_image_uri (str): Docker image URI.
         service_port (int): Service port.
+        streaming_embeddings (bool): Flag to enable streaming embeddings.
     """
     max_catalog_size: int = dataclasses.field(default=None)
     max_dimension: int = dataclasses.field(default=None)
     index_output_path: str = dataclasses.field(default=None)
     docker_image_uri: str = dataclasses.field(default=None)
     service_port: int = dataclasses.field(default=None)
+    streaming_embeddings: bool = dataclasses.field(default=None)
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.CUSTOM_TRAINED_MODEL

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/batch_prediction.py

@@ -81,7 +81,7 @@ class BatchPrediction(AbstractApiClass):
             BatchPredictionArgs, globalPredictionArgs)
         self.batch_prediction_args = client._build_class(
             BatchPredictionArgs, batchPredictionArgs)
-        self.deprecated_keys = {'global_prediction_args', 'explanations'}
+        self.deprecated_keys = {'explanations', 'global_prediction_args'}
 
     def __repr__(self):
         repr_dict = {f'batch_prediction_id': repr(self.batch_prediction_id), f'created_at': repr(self.created_at), f'name': repr(self.name), f'deployment_id': repr(self.deployment_id), f'file_connector_output_location': repr(self.file_connector_output_location), f'database_connector_id': repr(self.database_connector_id), f'database_output_configuration': repr(self.database_output_configuration), f'explanations': repr(self.explanations), f'file_output_format': repr(self.file_output_format), f'connector_type': repr(self.connector_type), f'legacy_input_location': repr(self.legacy_input_location), f'output_feature_group_id': repr(self.output_feature_group_id), f'feature_group_table_name': repr(self.feature_group_table_name), f'output_feature_group_table_name': repr(self.output_feature_group_table_name), f'summary_feature_group_table_name': repr(

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/batch_prediction_version.py

@@ -95,7 +95,7 @@ class BatchPredictionVersion(AbstractApiClass):
             BatchPredictionArgs, globalPredictionArgs)
         self.batch_prediction_args = client._build_class(
             BatchPredictionArgs, batchPredictionArgs)
-        self.deprecated_keys = {'global_prediction_args', 'explanations'}
+        self.deprecated_keys = {'explanations', 'global_prediction_args'}
 
     def __repr__(self):
         repr_dict = {f'batch_prediction_version': repr(self.batch_prediction_version), f'batch_prediction_id': repr(self.batch_prediction_id), f'status': repr(self.status), f'drift_monitor_status': repr(self.drift_monitor_status), f'deployment_id': repr(self.deployment_id), f'model_id': repr(self.model_id), f'model_version': repr(self.model_version), f'predictions_started_at': repr(self.predictions_started_at), f'predictions_completed_at': repr(self.predictions_completed_at), f'database_output_error': repr(self.database_output_error), f'total_predictions': repr(self.total_predictions), f'failed_predictions': repr(self.failed_predictions), f'database_connector_id': repr(self.database_connector_id), f'database_output_configuration': repr(self.database_output_configuration), f'explanations': repr(self.explanations), f'file_connector_output_location': repr(self.file_connector_output_location), f'file_output_format': repr(self.file_output_format), f'connector_type': repr(self.connector_type), f'legacy_input_location': repr(self.legacy_input_location), f'error': repr(self.error), f'drift_monitor_error': repr(self.drift_monitor_error), f'monitor_warnings': repr(

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/client.py

@@ -585,7 +585,7 @@ class BaseApiClient:
         client_options (ClientOptions): Optional API client configurations
         skip_version_check (bool): If true, will skip checking the server's current API version on initializing the client
     """
-    client_version = '1.2.0'
+    client_version = '1.2.1'
 
     def __init__(self, api_key: str = None, server: str = None, client_options: ClientOptions = None, skip_version_check: bool = False):
         self.api_key = api_key
@@ -6365,7 +6365,7 @@ Creates a new feature group defined as the union of other feature group versions
             deployment_id (str): The unique identifier to a deployment created under the project.
             vector (list): Input vector to perform the k nearest neighbors with.
             k (int): Overrideable number of items to return.
-            distance (str): Specify the distance function to use when finding nearest neighbors.
+            distance (str): Specify the distance function to use. Options include “dot“, “cosine“, “euclidean“, and “manhattan“. Default = “dot“
             include_score (bool): If True, will return the score alongside the resulting embedding value.
             catalog_id (str): An optional parameter honored only for embeddings that provide a catalog id"""
         prediction_url = self._get_prediction_endpoint(
@@ -6383,16 +6383,17 @@ Creates a new feature group defined as the union of other feature group versions
             deployment_id, deployment_token) if deployment_token else None
         return self._call_api('getMultipleKNearest', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'queries': queries}, server_override=prediction_url)
 
-    def get_labels(self, deployment_token: str, deployment_id: str, query_data: dict) -> Dict:
+    def get_labels(self, deployment_token: str, deployment_id: str, query_data: dict, return_extracted_entities: bool = False) -> Dict:
         """Returns a list of scored labels for a document.
 
         Args:
             deployment_token (str): The deployment token to authenticate access to created deployments. This token is only authorized to predict on deployments in this project, so it is safe to embed this model inside of an application or website.
             deployment_id (str): The unique identifier to a deployment created under the project.
-            query_data (dict): Dictionary where key is "Content" and value is the text from which entities are to be extracted."""
+            query_data (dict): Dictionary where key is "Content" and value is the text from which entities are to be extracted.
+            return_extracted_entities (bool): (Optional) If True, will return the extracted entities in simpler format"""
         prediction_url = self._get_prediction_endpoint(
             deployment_id, deployment_token) if deployment_token else None
-        return self._call_api('getLabels', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'queryData': query_data}, server_override=prediction_url)
+        return self._call_api('getLabels', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'queryData': query_data, 'returnExtractedEntities': return_extracted_entities}, server_override=prediction_url)
 
     def get_entities_from_pdf(self, deployment_token: str, deployment_id: str, pdf: io.TextIOBase = None, doc_id: str = None, return_extracted_features: bool = False, verbose: bool = False, save_extracted_features: bool = None) -> Dict:
         """Extracts text from the provided PDF and returns a list of recognized labels and their scores.
@@ -6993,72 +6994,6 @@ Creates a new feature group defined as the union of other feature group versions
             batch_prediction_id (str): Unique string identifier of the batch prediction."""
         return self._call_api('deleteBatchPrediction', 'DELETE', query_params={'batchPredictionId': batch_prediction_id})
 
-    def add_user_item_interaction(self, streaming_token: str, dataset_id: str, timestamp: int, user_id: str, item_id: list, event_type: str, additional_attributes: dict):
-        """Adds a user-item interaction record (data row) to a streaming dataset.
-
-        Args:
-            streaming_token (str): The streaming token for authenticating requests to the dataset.
-            dataset_id (str): The unique string identifier for the streaming dataset to record data to.
-            timestamp (int): The Unix timestamp of the event.
-            user_id (str): The unique identifier for the user.
-            item_id (list): The unique identifier for the items.
-            event_type (str): The event type.
-            additional_attributes (dict): Attributes of the user interaction."""
-        return self._call_api('addUserItemInteraction', 'POST', query_params={'streamingToken': streaming_token, 'datasetId': dataset_id}, body={'timestamp': timestamp, 'userId': user_id, 'itemId': item_id, 'eventType': event_type, 'additionalAttributes': additional_attributes})
-
-    def upsert_user_attributes(self, streaming_token: str, dataset_id: str, user_id: str, user_attributes: dict):
-        """Adds a user attribute record (data row) to a streaming dataset.
-
-        Either the streaming dataset ID or the project ID is required.
-
-
-        Args:
-            streaming_token (str): The streaming token for authenticating requests to the dataset.
-            dataset_id (str): The unique string identifier for the streaming dataset to record data to.
-            user_id (str): The unique identifier for the user.
-            user_attributes (dict): Attributes of the user interaction."""
-        return self._call_api('upsertUserAttributes', 'POST', query_params={'streamingToken': streaming_token, 'datasetId': dataset_id}, body={'userId': user_id, 'userAttributes': user_attributes})
-
-    def upsert_item_attributes(self, streaming_token: str, dataset_id: str, item_id: str, item_attributes: dict):
-        """Adds an item attributes record (data row) to a streaming dataset.
-
-        Either the streaming dataset ID or the project ID is required.
-
-
-        Args:
-            streaming_token (str): The streaming token for authenticating requests to the dataset.
-            dataset_id (str): The unique string identifier for the streaming dataset to record data to.
-            item_id (str): The unique identifier for the item.
-            item_attributes (dict): Attributes of the item interaction."""
-        return self._call_api('upsertItemAttributes', 'POST', query_params={'streamingToken': streaming_token, 'datasetId': dataset_id}, body={'itemId': item_id, 'itemAttributes': item_attributes})
-
-    def add_multiple_user_item_interactions(self, streaming_token: str, dataset_id: str, interactions: list):
-        """Adds a user-item interaction record (data row) to a streaming dataset.
-
-        Args:
-            streaming_token (str): The streaming token for authenticating requests to the dataset.
-            dataset_id (str): The unique string identifier of the streaming dataset to record data to.
-            interactions (list): List of interactions, each interaction of format {'userId': userId, 'timestamp': timestamp, 'itemId': itemId, 'eventType': eventType, 'additionalAttributes': {'attribute1': 'abc', 'attribute2': 123}}."""
-        return self._call_api('addMultipleUserItemInteractions', 'POST', query_params={'streamingToken': streaming_token, 'datasetId': dataset_id}, body={'interactions': interactions})
-
-    def upsert_multiple_user_attributes(self, streaming_token: str, dataset_id: str, upserts: list):
-        """Adds multiple user attributes records (data rows) to a streaming dataset.
-
-        Args:
-            streaming_token (str): The streaming token for authenticating requests to the dataset.
-            dataset_id (str): A unique string identifier for the streaming dataset to record data to.
-            upserts (list): List of upserts, each upsert of format {'userId': userId, 'userAttributes': {'attribute1': 'abc', 'attribute2': 123}}."""
-        return self._call_api('upsertMultipleUserAttributes', 'POST', query_params={'streamingToken': streaming_token, 'datasetId': dataset_id}, body={'upserts': upserts})
-
-    def upsert_multiple_item_attributes(self, streaming_token: str, dataset_id: str, upserts: list):
-        """Adds multiple item attributes records (data rows) to a streaming dataset.
-
-        Args:
-            streaming_token (str): The streaming token for authenticating requests to the dataset.
-            dataset_id (str): A unique string identifier for the streaming dataset to record data to.
-            upserts (list): A list of upserts, each upsert of format {'itemId': itemId, 'itemAttributes': {'attribute1': 'abc', 'attribute2': 123}}."""
-        return self._call_api('upsertMultipleItemAttributes', 'POST', query_params={'streamingToken': streaming_token, 'datasetId': dataset_id}, body={'upserts': upserts})
-
     def upsert_item_embeddings(self, streaming_token: str, model_id: str, item_id: str, vector: list, catalog_id: str = None):
         """Upserts an embedding vector for an item id for a model_id.
 
@@ -8101,13 +8036,17 @@ Creates a new feature group defined as the union of other feature group versions
             document_retriever_id (str): A unique string identifier associated with the document retriever."""
         return self._call_api('restartDocumentRetriever', 'POST', query_params={}, body={'documentRetrieverId': document_retriever_id})
 
-    def get_relevant_snippets(self, doc_ids: list = None, blobs: io.TextIOBase = None, query: str = None, document_retriever_config: Union[dict, DocumentRetrieverConfig] = None, honor_sentence_boundary: bool = True, num_retrieval_margin_words: int = None, max_words_per_snippet: int = None, max_snippets_per_document: int = None, start_word_index: int = None, end_word_index: int = None, including_bounding_boxes: bool = False) -> List[DocumentRetrieverLookupResult]:
-        """Get relevant snippets from documents with respect to the query. Document retrievers may be created on-the-fly to perform lookup.
+    def get_relevant_snippets(self, doc_ids: list = None, blobs: io.TextIOBase = None, query: str = None, document_retriever_config: Union[dict, DocumentRetrieverConfig] = None, honor_sentence_boundary: bool = True, num_retrieval_margin_words: int = None, max_words_per_snippet: int = None, max_snippets_per_document: int = None, start_word_index: int = None, end_word_index: int = None, including_bounding_boxes: bool = False, text: str = None) -> List[DocumentRetrieverLookupResult]:
+        """Retrieves snippets relevant to a given query from specified documents. This function supports flexible input options,
+
+        allowing for retrieval from a variety of data sources including document IDs, blob data, and plain text. When multiple data
+        sources are provided, all are considered in the retrieval process. Document retrievers may be created on-the-fly to perform lookup.
+
 
         Args:
             doc_ids (list): A list of document store IDs to retrieve the snippets from.
             blobs (io.TextIOBase): A dictionary mapping document names to the blob data.
-            query (str): The query that the documents are relevant to.
+            query (str): Query string to find relevant snippets in the documents.
             document_retriever_config (DocumentRetrieverConfig): If provided, used to configure the retrieval steps like chunking for embeddings.
             num_retrieval_margin_words (int): If provided, will add this number of words from left and right of the returned snippets.
             max_words_per_snippet (int): If provided, will limit the number of words in each snippet to the value specified.
@@ -8115,7 +8054,8 @@ Creates a new feature group defined as the union of other feature group versions
             start_word_index (int): If provided, will start the snippet at the index (of words in the document) specified.
             end_word_index (int): If provided, will end the snippet at the index of (of words in the document) specified.
             including_bounding_boxes (bool): If true, will include the bounding boxes of the snippets if they are available.
+            text (str): Plain text from which to retrieve snippets.
 
         Returns:
             list[DocumentRetrieverLookupResult]: The snippets found from the documents."""
-        return self._proxy_request('GetRelevantSnippets', 'POST', query_params={}, body={'docIds': doc_ids, 'query': query, 'documentRetrieverConfig': json.dumps(document_retriever_config), 'honorSentenceBoundary': honor_sentence_boundary, 'numRetrievalMarginWords': num_retrieval_margin_words, 'maxWordsPerSnippet': max_words_per_snippet, 'maxSnippetsPerDocument': max_snippets_per_document, 'startWordIndex': start_word_index, 'endWordIndex': end_word_index, 'includingBoundingBoxes': including_bounding_boxes}, files=blobs, parse_type=DocumentRetrieverLookupResult)
+        return self._proxy_request('GetRelevantSnippets', 'POST', query_params={}, body={'docIds': doc_ids, 'query': query, 'documentRetrieverConfig': json.dumps(document_retriever_config), 'honorSentenceBoundary': honor_sentence_boundary, 'numRetrievalMarginWords': num_retrieval_margin_words, 'maxWordsPerSnippet': max_words_per_snippet, 'maxSnippetsPerDocument': max_snippets_per_document, 'startWordIndex': start_word_index, 'endWordIndex': end_word_index, 'includingBoundingBoxes': including_bounding_boxes, 'text': text}, files=blobs, parse_type=DocumentRetrieverLookupResult)

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/document_retriever.py

@@ -150,25 +150,62 @@ class DocumentRetriever(AbstractApiClass):
         # we set deployment to be off by default for some cases.
         # So wait for indexing is done and then call restart.
         self.client._poll(self, {'PENDING', 'INDEXING'}, timeout=timeout / 2)
-        if self.get_status() == 'STOPPED':
+        if self.get_deployment_status() == 'STOPPED':
             self.client.restart_document_retriever(self.document_retriever_id)
         version = self.describe().latest_document_retriever_version
         if not version:
             from .client import ApiException
             raise ApiException(
                 409, 'This vector store does not have any versions')
-        version.wait_until_ready(timeout=timeout)
-        return self
+        else:
+            self.wait_until_deployment_ready(timeout=timeout / 2)
+        return self.refresh()
+
+    def wait_until_deployment_ready(self, timeout: int = 3600):
+        """
+        A waiting call until the document retriever deployment is ready to serve.
+
+        Args:
+            timeout (int, optional): The waiting time given to the call to finish, if it doesn't finish by the allocated time, the call is said to be timed out. Default value given is 3600 seconds.
+        """
+        import time
+
+        from .client import ApiException
+        start_time = time.time()
+        while (True):
+            deployment_status = self.get_deployment_status()
+            if deployment_status in {'PENDING', 'DEPLOYING'}:
+                if timeout and time.time() - start_time > timeout:
+                    raise TimeoutError(
+                        f'Maximum wait time of {timeout}s exceeded')
+            elif deployment_status in {'STOPPED'}:
+                raise ApiException(
+                    409, f'Document retriever deployment is stopped, please restart it.')
+            elif deployment_status in {'FAILED'}:
+                raise ApiException(
+                    409, f'Document retriever deployment failed, please retry deploying it.')
+            else:
+                return self.refresh()
+            time.sleep(15)
 
     def get_status(self):
         """
-        Gets the status of the document retriever.
+        Gets the indexing status of the document retriever.
 
         Returns:
             str: A string describing the status of a document retriever (pending, complete, etc.).
         """
         return self.describe().latest_document_retriever_version.status
 
+    def get_deployment_status(self):
+        """
+        Gets the deployment status of the document retriever.
+
+        Returns:
+            str: A string describing the deployment status of document retriever (pending, deploying, active, etc.).
+        """
+        return self.describe().latest_document_retriever_version.deployment_status
+
     def get_matching_documents(self, query: str, filters: dict = None, limit: int = None, result_columns: list = None, max_words: int = None, num_retrieval_margin_words: int = None, max_words_per_chunk: int = None, score_multiplier_column: str = None):
         """
         Lookup document retrievers and return the matching documents from the document retriever deployed with given query.

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/document_retriever_config.py

@@ -12,19 +12,21 @@ class DocumentRetrieverConfig(AbstractApiClass):
             chunkOverlapFraction (float): The fraction of overlap between two consecutive chunks.
             textEncoder (str): The text encoder used to encode texts in the vector store.
             scoreMultiplierColumn (str): The values in this metadata column are used to modify the relevance scores of returned chunks.
+            pruneVectors (bool): Corpus specific transformation of vectors that applies dimensional reduction techniques to strip common components from the vectors.
     """
 
-    def __init__(self, client, chunkSize=None, chunkOverlapFraction=None, textEncoder=None, scoreMultiplierColumn=None):
+    def __init__(self, client, chunkSize=None, chunkOverlapFraction=None, textEncoder=None, scoreMultiplierColumn=None, pruneVectors=None):
         super().__init__(client, None)
         self.chunk_size = chunkSize
         self.chunk_overlap_fraction = chunkOverlapFraction
         self.text_encoder = textEncoder
         self.score_multiplier_column = scoreMultiplierColumn
+        self.prune_vectors = pruneVectors
         self.deprecated_keys = {}
 
     def __repr__(self):
         repr_dict = {f'chunk_size': repr(self.chunk_size), f'chunk_overlap_fraction': repr(self.chunk_overlap_fraction), f'text_encoder': repr(
-            self.text_encoder), f'score_multiplier_column': repr(self.score_multiplier_column)}
+            self.text_encoder), f'score_multiplier_column': repr(self.score_multiplier_column), f'prune_vectors': repr(self.prune_vectors)}
         class_name = "DocumentRetrieverConfig"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None and key not in self.deprecated_keys])
@@ -37,6 +39,6 @@ class DocumentRetrieverConfig(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'chunk_size': self.chunk_size, 'chunk_overlap_fraction': self.chunk_overlap_fraction,
-                'text_encoder': self.text_encoder, 'score_multiplier_column': self.score_multiplier_column}
+        resp = {'chunk_size': self.chunk_size, 'chunk_overlap_fraction': self.chunk_overlap_fraction, 'text_encoder': self.text_encoder,
+                'score_multiplier_column': self.score_multiplier_column, 'prune_vectors': self.prune_vectors}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/document_retriever_version.py

@@ -13,6 +13,7 @@ class DocumentRetrieverVersion(AbstractApiClass):
             documentRetrieverVersion (str): The unique identifier of the vector store version.
             createdAt (str): When the vector store was created.
             status (str): The status of creating vector store version.
+            deploymentStatus (str): The status of deploying the vector store version.
             featureGroupId (str): The feature group id associated with the document retriever.
             featureGroupVersion (str): The unique identifier of the feature group version at which the vector store version is created.
             error (str): The error message when it failed to create the document retriever version.
@@ -22,12 +23,13 @@ class DocumentRetrieverVersion(AbstractApiClass):
             resolvedConfig (DocumentRetrieverConfig): The resolved configurations, such as default settings, for indexing documents.
     """
 
-    def __init__(self, client, documentRetrieverId=None, documentRetrieverVersion=None, createdAt=None, status=None, featureGroupId=None, featureGroupVersion=None, error=None, numberOfChunks=None, embeddingFileSize=None, warnings=None, resolvedConfig={}):
+    def __init__(self, client, documentRetrieverId=None, documentRetrieverVersion=None, createdAt=None, status=None, deploymentStatus=None, featureGroupId=None, featureGroupVersion=None, error=None, numberOfChunks=None, embeddingFileSize=None, warnings=None, resolvedConfig={}):
         super().__init__(client, documentRetrieverVersion)
         self.document_retriever_id = documentRetrieverId
         self.document_retriever_version = documentRetrieverVersion
         self.created_at = createdAt
         self.status = status
+        self.deployment_status = deploymentStatus
         self.feature_group_id = featureGroupId
         self.feature_group_version = featureGroupVersion
         self.error = error
@@ -39,8 +41,8 @@ class DocumentRetrieverVersion(AbstractApiClass):
         self.deprecated_keys = {}
 
     def __repr__(self):
-        repr_dict = {f'document_retriever_id': repr(self.document_retriever_id), f'document_retriever_version': repr(self.document_retriever_version), f'created_at': repr(self.created_at), f'status': repr(self.status), f'feature_group_id': repr(self.feature_group_id), f'feature_group_version': repr(
-            self.feature_group_version), f'error': repr(self.error), f'number_of_chunks': repr(self.number_of_chunks), f'embedding_file_size': repr(self.embedding_file_size), f'warnings': repr(self.warnings), f'resolved_config': repr(self.resolved_config)}
+        repr_dict = {f'document_retriever_id': repr(self.document_retriever_id), f'document_retriever_version': repr(self.document_retriever_version), f'created_at': repr(self.created_at), f'status': repr(self.status), f'deployment_status': repr(self.deployment_status), f'feature_group_id': repr(
+            self.feature_group_id), f'feature_group_version': repr(self.feature_group_version), f'error': repr(self.error), f'number_of_chunks': repr(self.number_of_chunks), f'embedding_file_size': repr(self.embedding_file_size), f'warnings': repr(self.warnings), f'resolved_config': repr(self.resolved_config)}
         class_name = "DocumentRetrieverVersion"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None and key not in self.deprecated_keys])
@@ -53,7 +55,7 @@ class DocumentRetrieverVersion(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'document_retriever_id': self.document_retriever_id, 'document_retriever_version': self.document_retriever_version, 'created_at': self.created_at, 'status': self.status, 'feature_group_id': self.feature_group_id,
+        resp = {'document_retriever_id': self.document_retriever_id, 'document_retriever_version': self.document_retriever_version, 'created_at': self.created_at, 'status': self.status, 'deployment_status': self.deployment_status, 'feature_group_id': self.feature_group_id,
                 'feature_group_version': self.feature_group_version, 'error': self.error, 'number_of_chunks': self.number_of_chunks, 'embedding_file_size': self.embedding_file_size, 'warnings': self.warnings, 'resolved_config': self._get_attribute_as_dict(self.resolved_config)}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}
 
@@ -86,7 +88,11 @@ class DocumentRetrieverVersion(AbstractApiClass):
         Args:
             timeout (int, optional): The waiting time given to the call to finish, if it doesn't finish by the allocated time, the call is said to be timed out.
         """
-        return self.client._poll(self, {'PENDING', 'INDEXING', 'DEPLOYING'}, timeout=timeout)
+        self.client._poll(self, {'PENDING', 'INDEXING'}, timeout=timeout / 2)
+        if self.get_deployment_status() == 'STOPPED':
+            self.client.restart_document_retriever(self.document_retriever_id)
+        self.wait_until_deployment_ready(timeout=timeout / 2)
+        return self.refresh()
 
     def wait_until_ready(self, timeout=3600):
         """
@@ -97,6 +103,33 @@ class DocumentRetrieverVersion(AbstractApiClass):
         """
         return self.wait_for_results(timeout)
 
+    def wait_until_deployment_ready(self, timeout: int = 3600):
+        """
+        A waiting call until the document retriever deployment is ready to serve.
+
+        Args:
+            timeout (int, optional): The waiting time given to the call to finish, if it doesn't finish by the allocated time, the call is said to be timed out. Default value given is 3600 seconds.
+        """
+        import time
+
+        from .client import ApiException
+        start_time = time.time()
+        while (True):
+            deployment_status = self.get_deployment_status()
+            if deployment_status in {'PENDING', 'DEPLOYING'}:
+                if timeout and time.time() - start_time > timeout:
+                    raise TimeoutError(
+                        f'Maximum wait time of {timeout}s exceeded')
+            elif deployment_status in {'STOPPED'}:
+                raise ApiException(
+                    409, f'Document retriever deployment is stopped, please restart it.')
+            elif deployment_status in {'FAILED'}:
+                raise ApiException(
+                    409, f'Document retriever deployment failed, please retry deploying it.')
+            else:
+                return self.refresh()
+            time.sleep(15)
+
     def get_status(self):
         """
         Gets the status of the document retriever version.
@@ -105,3 +138,12 @@ class DocumentRetrieverVersion(AbstractApiClass):
             str: A string describing the status of a document retriever version (pending, complete, etc.).
         """
         return self.describe().status
+
+    def get_deployment_status(self):
+        """
+        Gets the status of the document retriever version.
+
+        Returns:
+            str: A string describing the deployment status of a document retriever version (pending, deploying, etc.).
+        """
+        return self.describe().deployment_status

abacusai-1.2.1/abacusai/global_context.py

@@ -0,0 +1,13 @@
+import threading
+
+
+lock = threading.Lock()
+_context = {}
+
+
+def get(key):
+    return _context.get(key)
+
+
+def set(key, val):
+    _context[key] = val

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai/prediction_client.py

@@ -233,7 +233,7 @@ class PredictionClient(BaseApiClient):
             deployment_id (str): The unique identifier to a deployment created under the project.
             vector (list): Input vector to perform the k nearest neighbors with.
             k (int): Overrideable number of items to return.
-            distance (str): Specify the distance function to use when finding nearest neighbors.
+            distance (str): Specify the distance function to use. Options include “dot“, “cosine“, “euclidean“, and “manhattan“. Default = “dot“
             include_score (bool): If True, will return the score alongside the resulting embedding value.
             catalog_id (str): An optional parameter honored only for embeddings that provide a catalog id"""
         prediction_url = self._get_prediction_endpoint(
@@ -251,16 +251,17 @@ class PredictionClient(BaseApiClient):
             deployment_id, deployment_token) if deployment_token else None
         return self._call_api('getMultipleKNearest', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'queries': queries}, server_override=prediction_url)
 
-    def get_labels(self, deployment_token: str, deployment_id: str, query_data: dict) -> Dict:
+    def get_labels(self, deployment_token: str, deployment_id: str, query_data: dict, return_extracted_entities: bool = False) -> Dict:
         """Returns a list of scored labels for a document.
 
         Args:
             deployment_token (str): The deployment token to authenticate access to created deployments. This token is only authorized to predict on deployments in this project, so it is safe to embed this model inside of an application or website.
             deployment_id (str): The unique identifier to a deployment created under the project.
-            query_data (dict): Dictionary where key is "Content" and value is the text from which entities are to be extracted."""
+            query_data (dict): Dictionary where key is "Content" and value is the text from which entities are to be extracted.
+            return_extracted_entities (bool): (Optional) If True, will return the extracted entities in simpler format"""
         prediction_url = self._get_prediction_endpoint(
             deployment_id, deployment_token) if deployment_token else None
-        return self._call_api('getLabels', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'queryData': query_data}, server_override=prediction_url)
+        return self._call_api('getLabels', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'queryData': query_data, 'returnExtractedEntities': return_extracted_entities}, server_override=prediction_url)
 
     def get_entities_from_pdf(self, deployment_token: str, deployment_id: str, pdf: io.TextIOBase = None, doc_id: str = None, return_extracted_features: bool = False, verbose: bool = False, save_extracted_features: bool = None) -> Dict:
         """Extracts text from the provided PDF and returns a list of recognized labels and their scores.

abacusai-1.2.1/abacusai/public.pem

@@ -0,0 +1,5 @@
+-----BEGIN RSA PUBLIC KEY-----
+MIGJAoGBAJHbcVIwmETUUeR+tYt8xb015XMGV9eOhPsuyMCWjzdtCX1VZ83iPWPs
+JmygZpAqiV9ANGo562gFh8+xW10gJuQ698b+fuAuM8nMzxtKf3opbjBX1SMJxFms
+S0hU5MnMUJ4ZBVmbLB+EwbTECiiMBiu5EaOoZcr5918nw9xXjRP1AgMBAAE=
+-----END RSA PUBLIC KEY-----

{abacusai-1.2.0 → abacusai-1.2.1/abacusai.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: abacusai
-Version: 1.2.0
+Version: 1.2.1
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI

{abacusai-1.2.0 → abacusai-1.2.1}/abacusai.egg-info/SOURCES.txt

@@ -110,6 +110,7 @@ abacusai/forecasting_monitor_item_analysis.py
 abacusai/forecasting_monitor_summary.py
 abacusai/function_logs.py
 abacusai/generated_pit_feature_config_option.py
+abacusai/global_context.py
 abacusai/graph_dashboard.py
 abacusai/holdout_analysis.py
 abacusai/holdout_analysis_version.py

{abacusai-1.2.0 → abacusai-1.2.1}/docs/_sources/autoapi/abacusai/api_class/batch_prediction/index.rst.txt

@@ -13,7 +13,6 @@ Classes
 .. autoapisummary::
 
    abacusai.api_class.batch_prediction.BatchPredictionArgs
-   abacusai.api_class.batch_prediction.AnomalyOutliersBatchPredictionArgs
    abacusai.api_class.batch_prediction.ForecastingBatchPredictionArgs
    abacusai.api_class.batch_prediction.NamedEntityExtractionBatchPredictionArgs
    abacusai.api_class.batch_prediction.PersonalizationBatchPredictionArgs
@@ -55,39 +54,6 @@ Classes
 
 
 
-.. py:class:: AnomalyOutliersBatchPredictionArgs
-
-
-   Bases: :py:obj:`BatchPredictionArgs`
-
-   Batch Prediction Config for the ANOMALY_OUTLIERS problem type
-
-   :param for_eval: If True, the test fold which was created during training and used for metrics calculation will be used as input data. These predictions are hence, used for model evaluation.
-   :type for_eval: bool
-   :param threshold: The threshold for detecting an anomaly. Range: [0.8, 0.99]
-   :type threshold: float
-   :param explain_predictions: If True, calculates explanations for the predicted values along with predictions.
-   :type explain_predictions: bool
-
-   .. py:attribute:: for_eval
-      :type: bool
-
-      
-
-   .. py:attribute:: threshold
-      :type: float
-
-      
-
-   .. py:attribute:: explain_predictions
-      :type: bool
-
-      
-
-   .. py:method:: __post_init__()
-
-
-
 .. py:class:: ForecastingBatchPredictionArgs
 
 

{abacusai-1.2.0 → abacusai-1.2.1}/docs/_sources/autoapi/abacusai/api_class/deployment/index.rst.txt

@@ -15,7 +15,6 @@ Classes
    abacusai.api_class.deployment.PredictionArguments
    abacusai.api_class.deployment.OptimizationPredictionArguments
    abacusai.api_class.deployment.TimeseriesAnomalyPredictionArguments
-   abacusai.api_class.deployment.AnomalyOutliersPredictionArguments
    abacusai.api_class.deployment.ChatLLMPredictionArguments
    abacusai.api_class.deployment.RegressionPredictionArguments
    abacusai.api_class.deployment.ForecastingPredictionArguments
@@ -120,32 +119,6 @@ Classes
 
 
 
-.. py:class:: AnomalyOutliersPredictionArguments
-
-
-   Bases: :py:obj:`PredictionArguments`
-
-   Prediction arguments for the ANOMALY_OUTLIERS problem type
-
-   :param threshold: The threshold score of what is an anomaly. Valid values are between 0.8 and 0.99
-   :type threshold: float
-   :param histogram: If True, will return a histogram of the distribution of all points
-   :type histogram: bool
-
-   .. py:attribute:: threshold
-      :type: float
-
-      
-
-   .. py:attribute:: histogram
-      :type: str
-
-      
-
-   .. py:method:: __post_init__()
-
-
-
 .. py:class:: ChatLLMPredictionArguments
 
 

{abacusai-1.2.0 → abacusai-1.2.1}/docs/_sources/autoapi/abacusai/api_class/enums/index.rst.txt

@@ -102,11 +102,6 @@ Classes
 
       
 
-   .. py:attribute:: ANOMALY_OUTLIERS
-      :value: 'anomaly'
-
-      
-
    .. py:attribute:: EVENT_ANOMALY
       :value: 'event_anomaly'
 
@@ -1827,11 +1822,6 @@ Classes
 
       
 
-   .. py:attribute:: ABACUS_LLAMA2_QA
-      :value: 'ABACUS_LLAMA2_QA'
-
-      
-
    .. py:attribute:: LLAMA2_CHAT
       :value: 'LLAMA2_CHAT'