abacusai 1.3.0__tar.gz → 1.3.2__tar.gz

{abacusai-1.3.0/abacusai.egg-info → abacusai-1.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: abacusai
-Version: 1.3.0
+Version: 1.3.2
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI
@@ -21,3 +21,4 @@ Requires-Dist: packaging
 Requires-Dist: requests
 Requires-Dist: pandas
 Requires-Dist: fastavro
+Requires-Dist: typing_inspect; python_version < "3.8"

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/__init__.py

@@ -4,4 +4,4 @@ from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.2.5"
+__version__ = "1.3.2"

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/abstract.py

@@ -2,10 +2,17 @@ import dataclasses
 import datetime
 import inspect
 import re
+import sys
 from abc import ABC
 from copy import deepcopy
 from textwrap import dedent
-from typing import Any, Callable, get_origin, get_type_hints
+from typing import Any, Callable, get_type_hints
+
+
+if sys.version_info >= (3, 8):
+    from typing import get_origin
+else:
+    from typing_inspect import get_origin
 
 from .enums import ApiEnum
 

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/ai_agents.py

@@ -28,7 +28,7 @@ class WorkflowNodeInputSchema(ApiClass):
     A react-jsonschema-form conformant schema for workflow node input.
 
     Args:
-        json_schema (dict): The json schema for the input conformant to react-jsonschema-form specification. Must define keys like "title", "type" and "properties".
+        json_schema (dict): The json schema for the input conformant to react-jsonschema-form specification. Must define keys like "title", "type" and "properties". Supported elements - Checkbox, Radio Button, Dropdown, Textarea, Number, Date, File Upload. Not supported - Nested elements, arrays and other complex types.
         ui_schema (dict): The ui schema for the input conformant to react-jsonschema-form specification.
     """
     json_schema: dict

abacusai-1.3.2/abacusai/api_class/connectors.py

@@ -0,0 +1,43 @@
+import dataclasses
+
+from . import enums
+from .abstract import _ApiClassFactory
+from .dataset import DatasetConfig
+
+
+@dataclasses.dataclass
+class StreamingConnectorDatasetConfig(DatasetConfig):
+    """
+    An abstract class for dataset configs specific to streaming connectors.
+
+    Args:
+        streaming_connector_type (StreamingConnectorType): The type of streaming connector
+    """
+    streaming_connector_type: enums.StreamingConnectorType = dataclasses.field(default=None, repr=False, init=False)
+
+    @classmethod
+    def _get_builder(cls):
+        return _StreamingConnectorDatasetConfigFactory
+
+
+@dataclasses.dataclass
+class KafkaDatasetConfig(StreamingConnectorDatasetConfig):
+    """
+    Dataset config for Kafka Streaming Connector
+
+    Args:
+        topic (str): The kafka topic to consume
+    """
+    topic: str
+
+    def __post_init__(self):
+        self.streaming_connector_type = enums.StreamingConnectorType.KAFKA
+
+
+@dataclasses.dataclass
+class _StreamingConnectorDatasetConfigFactory(_ApiClassFactory):
+    config_abstract_class = StreamingConnectorDatasetConfig
+    config_class_key = 'streaming_connector_type'
+    config_class_map = {
+        enums.StreamingConnectorType.KAFKA: KafkaDatasetConfig,
+    }

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/dataset.py

@@ -4,6 +4,17 @@ from .abstract import ApiClass
 from .enums import OcrMode
 
 
+@dataclasses.dataclass
+class DatasetConfig(ApiClass):
+    """
+    An abstract class for dataset configs
+
+    Args:
+        is_documentset (bool): Whether the dataset is a document set
+    """
+    is_documentset: bool = dataclasses.field(default=None)
+
+
 @dataclasses.dataclass
 class ParsingConfig(ApiClass):
     """

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/dataset_application_connector.py

@@ -1,28 +1,27 @@
 import dataclasses
 
 from . import enums
-from .abstract import ApiClass, _ApiClassFactory
+from .abstract import _ApiClassFactory
+from .dataset import DatasetConfig
 
 
 @dataclasses.dataclass
-class DatasetConfig(ApiClass):
+class ApplicationConnectorDatasetConfig(DatasetConfig):
     """
     An abstract class for dataset configs specific to application connectors.
 
     Args:
         application_connector_type(enums.ApplicationConnectorType): The type of application connector
-        is_documentset (bool): Whether the dataset is a document set
     """
     application_connector_type: enums.ApplicationConnectorType = dataclasses.field(default=None, repr=False, init=False)
-    is_documentset: bool = dataclasses.field(default=None)
 
     @classmethod
     def _get_builder(cls):
-        return _DatasetConfigFactory
+        return _ApplicationConnectorDatasetConfigFactory
 
 
 @dataclasses.dataclass
-class ConfluenceDatasetConfig(DatasetConfig):
+class ConfluenceDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Confluence Application Connector
     Args:
@@ -42,7 +41,7 @@ class ConfluenceDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class GoogleAnalyticsDatasetConfig(DatasetConfig):
+class GoogleAnalyticsDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Google Analytics Application Connector
 
@@ -60,7 +59,7 @@ class GoogleAnalyticsDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class GoogleDriveDatasetConfig(DatasetConfig):
+class GoogleDriveDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Google Drive Application Connector
 
@@ -80,7 +79,7 @@ class GoogleDriveDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class JiraDatasetConfig(DatasetConfig):
+class JiraDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Jira Application Connector
 
@@ -100,7 +99,7 @@ class JiraDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class OneDriveDatasetConfig(DatasetConfig):
+class OneDriveDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for OneDrive Application Connector
 
@@ -120,13 +119,12 @@ class OneDriveDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class SharepointDatasetConfig(DatasetConfig):
+class SharepointDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Sharepoint Application Connector
 
     Args:
         location (str): The regex location of the files to fetch
-        is_documentset (bool): Whether the dataset is a document set
         csv_delimiter (str): If the file format is CSV, use a specific csv delimiter
         extract_bounding_boxes (bool): Signifies whether to extract bounding boxes out of the documents. Only valid if is_documentset if True
         merge_file_schemas (bool): Signifies if the merge file schema policy is enabled. Not applicable if is_documentset is True
@@ -141,7 +139,7 @@ class SharepointDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class ZendeskDatasetConfig(DatasetConfig):
+class ZendeskDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Zendesk Application Connector
 
@@ -155,7 +153,7 @@ class ZendeskDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class AbacusUsageMetricsDatasetConfig(DatasetConfig):
+class AbacusUsageMetricsDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Abacus Usage Metrics Application Connector
 
@@ -172,7 +170,7 @@ class AbacusUsageMetricsDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class FreshserviceDatasetConfig(DatasetConfig):
+class FreshserviceDatasetConfig(ApplicationConnectorDatasetConfig):
     """
     Dataset config for Freshservice Application Connector
     """
@@ -182,8 +180,8 @@ class FreshserviceDatasetConfig(DatasetConfig):
 
 
 @dataclasses.dataclass
-class _DatasetConfigFactory(_ApiClassFactory):
-    config_abstract_class = DatasetConfig
+class _ApplicationConnectorDatasetConfigFactory(_ApiClassFactory):
+    config_abstract_class = ApplicationConnectorDatasetConfig
     config_class_key = 'application_connector_type'
     config_class_map = {
         enums.ApplicationConnectorType.CONFLUENCE: ConfluenceDatasetConfig,

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/document_retriever.py

@@ -7,7 +7,7 @@ from .enums import VectorStoreTextEncoder
 @dataclasses.dataclass
 class VectorStoreConfig(ApiClass):
     """
-    Configs for vector store indexing.
+    Config for indexing options of a document retriever. Default values of optional arguments are heuristically selected by the Abacus.AI platform based on the underlying data.
 
     Args:
         chunk_size (int): The size of text chunks in the vector store.
@@ -25,8 +25,4 @@ class VectorStoreConfig(ApiClass):
     prune_vectors: bool = dataclasses.field(default=None)
 
 
-@dataclasses.dataclass
-class DocumentRetrieverConfig(VectorStoreConfig):
-    """
-    Configs for document retriever. If any configuration value is not explicitly provided, Abacus.AI will automatically infer default values based on the data.
-    """
+DocumentRetrieverConfig = VectorStoreConfig

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/enums.py

@@ -362,7 +362,11 @@ class ApplicationConnectorType(ApiEnum):
     ABACUSUSAGEMETRICS = 'ABACUSUSAGEMETRICS'
     MICROSOFTAUTH = 'MICROSOFTAUTH'
     FRESHSERVICE = 'FRESHSERVICE'
-    ZENDESKSUNSHINEMESSAGING = 'zendesksunshinemessaging'
+    ZENDESKSUNSHINEMESSAGING = 'ZENDESKSUNSHINEMESSAGING'
+
+
+class StreamingConnectorType(ApiEnum):
+    KAFKA = 'KAFKA'
 
 
 class PythonFunctionArgumentType(ApiEnum):
@@ -428,7 +432,7 @@ class LLMName(ApiEnum):
     GEMINI_1_5_PRO = 'GEMINI_1_5_PRO'
     MIXTRAL_CHAT = 'MIXTRAL_CHAT'
     MISTRAL_MEDIUM = 'MISTRAL_MEDIUM'
-    ABACUS_SMAUG2 = 'ABACUS_SMAUG2'
+    ABACUS_SMAUG3 = 'ABACUS_SMAUG3'
 
 
 class MonitorAlertType(ApiEnum):
@@ -532,13 +536,6 @@ class DataType(ApiEnum):
     STRUCT = 'struct'
     NULL = 'null'
 
-    @classmethod
-    def from_str(cls, value):
-        if not value:
-            return None
-        default_map = {val.value: val for val in DataType}
-        return default_map[value.lower()]
-
 
 class AgentInterface(ApiEnum):
     # Duplicated in reainternal.enums, both should be kept in sync

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/api_class/model.py

@@ -438,18 +438,19 @@ class ChatLLMTrainingConfig(TrainingConfig):
     Training config for the CHAT_LLM problem type
 
     Args:
-        document_retrievers (List[str]): List of document retriever names to use for the feature stores this model was trained with.
-        num_completion_tokens (int): Default for maximum number of tokens for chat answers. Reducing this will get faster responses which are more succinct
-        temperature (float): The generative LLM temperature
+        document_retrievers (List[str]): List of names of document retrievers to use as vector stores of information for RAG responses.
+        num_completion_tokens (int): Default for maximum number of tokens for chat answers. Reducing this will get faster responses which are more succinct.
+        temperature (float): The generative LLM temperature.
         retrieval_columns (list): Include the metadata column values in the retrieved search results.
         filter_columns (list): Allow users to filter the document retrievers on these metadata columns.
-        include_general_knowledge (bool): Allow the LLM to rely not just on search results, but to fall back on general knowledge.
+        include_general_knowledge (bool): Allow the LLM to rely not just on RAG search results, but to fall back on general knowledge. Disabled by default.
+        enable_web_search (bool) : Allow the LLM to use Web Search Engines to retrieve information for better results.
         behavior_instructions (str): Customize the overall role instructions for the LLM.
-        response_instructions (str): Customize instructions for what the LLM responses should look like.
-        enable_llm_rewrite (bool): Enable LLM rewrite for the ChatLLM. If None, LLM rewrite will happen automatically. Defaults to False.
+        response_instructions (str): Customized instructions for how the LLM should respond.
+        enable_llm_rewrite (bool): If enabled, an LLM will rewrite the RAG queries sent to document retriever. Disabled by default.
         column_filtering_instructions (str): Instructions for a LLM call to automatically generate filter expressions on document metadata to retrieve relevant documents for the conversation.
         keyword_requirement_instructions (str): Instructions for a LLM call to automatically generate keyword requirements to retrieve relevant documents for the conversation.
-        query_rewrite_instructions (str): Instructions for a LLM call to rewrite a search query.
+        query_rewrite_instructions (str): Special instructions for the LLM which rewrites the RAG query.
         max_search_results (int): Maximum number of search results in the retrieval augmentation step. If we know that the questions are likely to have snippets which are easily matched in the documents, then a lower number will help with accuracy.
         data_feature_group_ids: (List[str]): List of feature group IDs to use to possibly query for the ChatLLM. The created ChatLLM is commonly referred to as DataLLM.
         data_prompt_context (str): Prompt context for the data feature group IDs.
@@ -458,33 +459,32 @@ class ChatLLMTrainingConfig(TrainingConfig):
         search_score_cutoff (float): Minimum search score to consider a document as a valid search result.
         database_connector_id (str): Database connector ID to use for the ChatLLM.
         database_connector_tables (List[str]): List of tables to use from the database connector for the ChatLLM.
-        enable_code_execution (bool): Enable code execution in the ChatLLM.
-        metadata_columns (list): DEPRECATED. Include the metadata column values in the retrieved search results.
-        lookup_rewrite_instructions (str): DEPRECATED. Instructions for a LLM call to rewrite a search query.
-    """
-    document_retrievers: List[str] = None
-    num_completion_tokens: int = None
-    temperature: float = None
-    retrieval_columns: list = None
-    filter_columns: list = None
-    include_general_knowledge: bool = None
-    behavior_instructions: str = None
-    response_instructions: str = None
-    enable_llm_rewrite: bool = False
-    column_filtering_instructions: str = None
-    keyword_requirement_instructions: str = None
-    query_rewrite_instructions: str = None
-    max_search_results: int = None
-    data_feature_group_ids: List[str] = None
-    data_prompt_context: str = None
-    hide_generated_sql: bool = None
-    disable_data_summarization: bool = None
-    search_score_cutoff: float = None
-    database_connector_id: str = None
-    database_connector_tables: List[str] = None
-    enable_code_execution: bool = None
-    metadata_columns: list = None
-    lookup_rewrite_instructions: str = None
+        enable_code_execution (bool): Enable python code execution in the ChatLLM. This equips the LLM with a python kernel in which all its code is executed.
+    """
+    document_retrievers: List[str] = dataclasses.field(default=None)
+    num_completion_tokens: int = dataclasses.field(default=None)
+    temperature: float = dataclasses.field(default=None)
+    retrieval_columns: list = dataclasses.field(default=None)
+    filter_columns: list = dataclasses.field(default=None)
+    include_general_knowledge: bool = dataclasses.field(default=None)
+    enable_web_search: bool = dataclasses.field(default=None)
+    behavior_instructions: str = dataclasses.field(default=None)
+    response_instructions: str = dataclasses.field(default=None)
+    enable_llm_rewrite: bool = dataclasses.field(default=None)
+    column_filtering_instructions: str = dataclasses.field(default=None)
+    keyword_requirement_instructions: str = dataclasses.field(default=None)
+    query_rewrite_instructions: str = dataclasses.field(default=None)
+    max_search_results: int = dataclasses.field(default=None)
+    data_feature_group_ids: List[str] = dataclasses.field(default=None)
+    data_prompt_context: str = dataclasses.field(default=None)
+    hide_generated_sql: bool = dataclasses.field(default=None)
+    disable_data_summarization: bool = dataclasses.field(default=None)
+    search_score_cutoff: float = dataclasses.field(default=None)
+    database_connector_id: str = dataclasses.field(default=None)
+    database_connector_tables: List[str] = dataclasses.field(default=None)
+    enable_code_execution: bool = dataclasses.field(default=None)
+    metadata_columns: list = dataclasses.field(default=None, metadata={'deprecated': True})
+    lookup_rewrite_instructions: str = dataclasses.field(default=None, metadata={'deprecated': True})
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.CHAT_LLM

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/batch_prediction.py

@@ -84,7 +84,7 @@ class BatchPrediction(AbstractApiClass):
             BatchPredictionArgs, globalPredictionArgs)
         self.batch_prediction_args = client._build_class(getattr(
             api_class, batchPredictionArgsType, BatchPredictionArgs) if batchPredictionArgsType else BatchPredictionArgs, batchPredictionArgs)
-        self.deprecated_keys = {'explanations', 'global_prediction_args'}
+        self.deprecated_keys = {'global_prediction_args', 'explanations'}
 
     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'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(self.summary_feature_group_table_name), f'csv_input_prefix': repr(

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/batch_prediction_version.py

@@ -100,7 +100,7 @@ class BatchPredictionVersion(AbstractApiClass):
             BatchPredictionArgs, globalPredictionArgs)
         self.batch_prediction_args = client._build_class(getattr(
             api_class, batchPredictionArgsType, BatchPredictionArgs) if batchPredictionArgsType else BatchPredictionArgs, batchPredictionArgs)
-        self.deprecated_keys = {'explanations', 'global_prediction_args'}
+        self.deprecated_keys = {'global_prediction_args', 'explanations'}
 
     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'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(self.monitor_warnings), f'csv_input_prefix': repr(

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/client.py

@@ -33,13 +33,13 @@ from .annotation_entry import AnnotationEntry
 from .annotations_status import AnnotationsStatus
 from .api_class import (
     AgentInterface, AlertActionConfig, AlertConditionConfig, ApiClass, ApiEnum,
-    BatchPredictionArgs, BlobInput, DatasetConfig,
+    ApplicationConnectorDatasetConfig, BatchPredictionArgs, BlobInput,
     DatasetDocumentProcessingConfig, DataType, DocumentProcessingConfig,
-    DocumentRetrieverConfig, EvalArtifactType, FeatureGroupExportConfig,
-    ForecastingMonitorConfig, IncrementalDatabaseConnectorConfig, LLMName,
-    MergeConfig, ParsingConfig, PredictionArguments, ProblemType,
-    ProjectFeatureGroupConfig, PythonFunctionType, SamplingConfig,
-    TrainingConfig, WorkflowGraph, get_clean_function_source_code
+    EvalArtifactType, FeatureGroupExportConfig, ForecastingMonitorConfig,
+    IncrementalDatabaseConnectorConfig, LLMName, MergeConfig, ParsingConfig,
+    PredictionArguments, ProblemType, ProjectFeatureGroupConfig,
+    PythonFunctionType, SamplingConfig, TrainingConfig, VectorStoreConfig,
+    WorkflowGraph, get_clean_function_source_code
 )
 from .api_class.abstract import get_clean_function_source_code, snake_case
 from .api_class.ai_agents import WorkflowGraph
@@ -74,7 +74,6 @@ from .deployment_conversation import DeploymentConversation
 from .deployment_conversation_export import DeploymentConversationExport
 from .document_data import DocumentData
 from .document_retriever import DocumentRetriever
-from .document_retriever_config import DocumentRetrieverConfig
 from .document_retriever_lookup_result import DocumentRetrieverLookupResult
 from .document_retriever_version import DocumentRetrieverVersion
 from .drift_distributions import DriftDistributions
@@ -168,6 +167,7 @@ from .upload_part import UploadPart
 from .use_case import UseCase
 from .use_case_requirements import UseCaseRequirements
 from .user import User
+from .web_search_response import WebSearchResponse
 from .webhook import Webhook
 
 
@@ -608,7 +608,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.5'
+    client_version = '1.3.2'
 
     def __init__(self, api_key: str = None, server: str = None, client_options: ClientOptions = None, skip_version_check: bool = False):
         self.api_key = api_key
@@ -5282,13 +5282,13 @@ Creates a new feature group defined as the union of other feature group versions
             Dataset: The created dataset."""
         return self._call_api('createDatasetFromDatabaseConnector', 'POST', query_params={}, body={'tableName': table_name, 'databaseConnectorId': database_connector_id, 'objectName': object_name, 'columns': columns, 'queryArguments': query_arguments, 'refreshSchedule': refresh_schedule, 'sqlQuery': sql_query, 'incremental': incremental, 'incrementalDatabaseConnectorConfig': incremental_database_connector_config}, parse_type=Dataset)
 
-    def create_dataset_from_application_connector(self, table_name: str, application_connector_id: str, dataset_config: Union[dict, DatasetConfig] = None, refresh_schedule: str = None) -> Dataset:
+    def create_dataset_from_application_connector(self, table_name: str, application_connector_id: str, dataset_config: Union[dict, ApplicationConnectorDatasetConfig] = None, refresh_schedule: str = None) -> Dataset:
         """Creates a dataset from an Application Connector.
 
         Args:
             table_name (str): Organization-unique table name.
             application_connector_id (str): Unique string identifier of the application connector to download data from.
-            dataset_config (DatasetConfig): Dataset config for the application connector.
+            dataset_config (ApplicationConnectorDatasetConfig): Dataset config for the application connector.
             refresh_schedule (str): Cron time string format that describes a schedule to retrieve the latest version of the imported dataset. The time is specified in UTC.
 
         Returns:
@@ -5309,12 +5309,12 @@ Creates a new feature group defined as the union of other feature group versions
             DatasetVersion: The new Dataset Version created."""
         return self._call_api('createDatasetVersionFromDatabaseConnector', 'POST', query_params={'datasetId': dataset_id}, body={'objectName': object_name, 'columns': columns, 'queryArguments': query_arguments, 'sqlQuery': sql_query}, parse_type=DatasetVersion)
 
-    def create_dataset_version_from_application_connector(self, dataset_id: str, dataset_config: Union[dict, DatasetConfig] = None) -> DatasetVersion:
+    def create_dataset_version_from_application_connector(self, dataset_id: str, dataset_config: Union[dict, ApplicationConnectorDatasetConfig] = None) -> DatasetVersion:
         """Creates a new version of the specified dataset.
 
         Args:
             dataset_id (str): The unique ID associated with the dataset.
-            dataset_config (DatasetConfig): Dataset config for the application connector. If any of the fields are not specified, the last values will be used.
+            dataset_config (ApplicationConnectorDatasetConfig): Dataset config for the application connector. If any of the fields are not specified, the last values will be used.
 
         Returns:
             DatasetVersion: The new Dataset Version created."""
@@ -7108,6 +7108,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('lookupMatches', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'data': data, 'filters': filters, 'num': num, 'resultColumns': result_columns, 'maxWords': max_words, 'numRetrievalMarginWords': num_retrieval_margin_words, 'maxWordsPerChunk': max_words_per_chunk, 'scoreMultiplierColumn': score_multiplier_column, 'minScore': min_score, 'requiredPhrases': required_phrases, 'filterClause': filter_clause, 'crowdingLimits': crowding_limits}, parse_type=DocumentRetrieverLookupResult, server_override=prediction_url)
 
+    def get_completion(self, deployment_token: str, deployment_id: str, prompt: str) -> Dict:
+        """Returns the finetuned LLM generated completion of the prompt.
+
+        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.
+            prompt (str): The prompt given to the finetuned LLM to generate the completion."""
+        prediction_url = self._get_prediction_endpoint(
+            deployment_id, deployment_token) if deployment_token else None
+        return self._call_api('getCompletion', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, body={'prompt': prompt}, server_override=prediction_url)
+
     def execute_agent_with_binary_data(self, deployment_token: str, deployment_id: str, arguments: list = None, keyword_arguments: dict = None, deployment_conversation_id: str = None, external_session_id: str = None, blobs: None = None) -> Dict:
         """Executes a deployed AI agent function with binary data as inputs.
 
@@ -8185,6 +8196,28 @@ Creates a new feature group defined as the union of other feature group versions
             ExtractedFields: The response from the document query."""
         return self._call_api('extractDataUsingLLM', 'POST', query_params={}, body={'fieldDescriptors': field_descriptors, 'documentId': document_id, 'documentText': document_text, 'llmName': llm_name}, parse_type=ExtractedFields)
 
+    def search_web_for_llm(self, queries: List, search_providers: List = None, max_results: int = 1, safe: bool = True, fetch_content: bool = False, max_page_tokens: int = 8192, convert_to_markdown: bool = True) -> WebSearchResponse:
+        """Access web search providers to fetch content related to the queries for use in large language model inputs.
+
+        This method can access multiple search providers and return information from them. If the provider supplies
+        URLs for the results then this method also supports fetching the contents of those URLs, optionally converting
+        them to markdown format, and returning them as part of the response. Set a token budget to limit the amount of
+        content returned in the response.
+
+
+        Args:
+            queries (List): List of queries to send to the search providers. At most 10 queries each less than 512 characters.
+            search_providers (List): Search providers to use for the search. If not provided a default provider is used. - BING - GOOGLE
+            max_results (int): Maximum number of results to fetch per provider. Must be in [1, 100]. Defaults to 1 (I'm feeling lucky).
+            safe (bool): Whether content safety is enabled for these search request. Defaults to True.
+            fetch_content (bool): If true fetches the content from the urls in the search results. Defailts to False.
+            max_page_tokens (int): Maximum number of tokens to accumulate if fetching search result contents.
+            convert_to_markdown (bool): Whether content should be converted to markdown. Defaults to True.
+
+        Returns:
+            WebSearchResponse: Results of running the search queries."""
+        return self._proxy_request('SearchWebForLlm', 'POST', query_params={}, body={'queries': queries, 'searchProviders': search_providers, 'maxResults': max_results, 'safe': safe, 'fetchContent': fetch_content, 'maxPageTokens': max_page_tokens, 'convertToMarkdown': convert_to_markdown}, parse_type=WebSearchResponse)
+
     def construct_agent_conversation_messages_for_llm(self, current_message: str = None, current_doc_ids: List = None, include_history: bool = True, include_document_contents: bool = True, deployment_conversation_id: str = None, external_session_id: str = None, max_document_words: int = None) -> ChatMessage:
         """Returns conversation history in a format for LLM calls.
 
@@ -8212,17 +8245,17 @@ Creates a new feature group defined as the union of other feature group versions
             LlmResponse: The response from the LLM App."""
         return self._call_api('getLLMAppResponse', 'POST', query_params={}, body={'llmAppName': llm_app_name, 'prompt': prompt}, parse_type=LlmResponse)
 
-    def create_document_retriever(self, project_id: str, name: str, feature_group_id: str, document_retriever_config: Union[dict, DocumentRetrieverConfig] = None) -> DocumentRetriever:
+    def create_document_retriever(self, project_id: str, name: str, feature_group_id: str, document_retriever_config: Union[dict, VectorStoreConfig] = None) -> DocumentRetriever:
         """Returns a document retriever that stores embeddings for document chunks in a feature group.
 
         Document columns in the feature group are broken into chunks. For cases with multiple document columns, chunks from all columns are combined together to form a single chunk.
 
 
         Args:
-            project_id (str): The ID of project that the vector store is created in.
-            name (str): The name of the vector store. Can be up to 120 characters long and can only contain alphanumeric characters and underscores.
-            feature_group_id (str): The ID of the feature group that the vector store is associated with.
-            document_retriever_config (DocumentRetrieverConfig): The configuration, including chunk_size and chunk_overlap_fraction, for document retrieval.
+            project_id (str): The ID of project that the Document Retriever is created in.
+            name (str): The name of the Document Retriever. Can be up to 120 characters long and can only contain alphanumeric characters and underscores.
+            feature_group_id (str): The ID of the feature group that the Document Retriever is associated with.
+            document_retriever_config (VectorStoreConfig): The configuration, including chunk_size and chunk_overlap_fraction, for document retrieval.
 
         Returns:
             DocumentRetriever: The newly created document retriever."""
@@ -8239,13 +8272,13 @@ Creates a new feature group defined as the union of other feature group versions
             DocumentRetriever: The updated document retriever."""
         return self._call_api('renameDocumentRetriever', 'POST', query_params={}, body={'documentRetrieverId': document_retriever_id, 'name': name}, parse_type=DocumentRetriever)
 
-    def create_document_retriever_version(self, document_retriever_id: str, feature_group_id: str = None, document_retriever_config: Union[dict, DocumentRetrieverConfig] = None) -> DocumentRetrieverVersion:
+    def create_document_retriever_version(self, document_retriever_id: str, feature_group_id: str = None, document_retriever_config: Union[dict, VectorStoreConfig] = None) -> DocumentRetrieverVersion:
         """Creates a document retriever version from the latest version of the feature group that the document retriever associated with.
 
         Args:
             document_retriever_id (str): The unique ID associated with the document retriever to create version with.
             feature_group_id (str): The ID of the feature group to update the document retriever with.
-            document_retriever_config (DocumentRetrieverConfig): The configuration, including chunk_size and chunk_overlap_fraction, for document retrieval.
+            document_retriever_config (VectorStoreConfig): The configuration, including chunk_size and chunk_overlap_fraction, for document retrieval.
 
         Returns:
             DocumentRetrieverVersion: The newly created document retriever version."""
@@ -8281,7 +8314,7 @@ 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, text: str = None) -> List[DocumentRetrieverLookupResult]:
+    def get_relevant_snippets(self, doc_ids: List = None, blobs: io.TextIOBase = None, query: str = None, document_retriever_config: Union[dict, VectorStoreConfig] = 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
@@ -8292,7 +8325,7 @@ Creates a new feature group defined as the union of other feature group versions
             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): 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.
+            document_retriever_config (VectorStoreConfig): 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.
             max_snippets_per_document (int): If provided, will limit the number of snippets retrieved from each document to the value specified.

abacusai-1.3.0/abacusai/agent_data_upload_result.py → abacusai-1.3.2/abacusai/data_upload_result.py

@@ -1,7 +1,7 @@
 from .return_class import AbstractApiClass
 
 
-class AgentDataUploadResult(AbstractApiClass):
+class DataUploadResult(AbstractApiClass):
     """
         Results of uploading data to agent.
 
@@ -17,7 +17,7 @@ class AgentDataUploadResult(AbstractApiClass):
 
     def __repr__(self):
         repr_dict = {f'doc_infos': repr(self.doc_infos)}
-        class_name = "AgentDataUploadResult"
+        class_name = "DataUploadResult"
         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])
         return f"{class_name}({repr_str})"

{abacusai-1.3.0 → abacusai-1.3.2}/abacusai/dataset.py

@@ -1,6 +1,9 @@
 from typing import Union
 
-from .api_class import DatasetConfig, DatasetDocumentProcessingConfig, DataType, DocumentProcessingConfig, ParsingConfig
+from .api_class import (
+    ApplicationConnectorDatasetConfig, DatasetDocumentProcessingConfig,
+    DataType, DocumentProcessingConfig, ParsingConfig
+)
 from .dataset_column import DatasetColumn
 from .dataset_version import DatasetVersion
 from .refresh_schedule import RefreshSchedule
@@ -118,12 +121,12 @@ class Dataset(AbstractApiClass):
         """
         return self.client.create_dataset_version_from_database_connector(self.dataset_id, object_name, columns, query_arguments, sql_query)
 
-    def create_version_from_application_connector(self, dataset_config: Union[dict, DatasetConfig] = None):
+    def create_version_from_application_connector(self, dataset_config: Union[dict, ApplicationConnectorDatasetConfig] = None):
         """
         Creates a new version of the specified dataset.
 
         Args:
-            dataset_config (DatasetConfig): Dataset config for the application connector. If any of the fields are not specified, the last values will be used.
+            dataset_config (ApplicationConnectorDatasetConfig): Dataset config for the application connector. If any of the fields are not specified, the last values will be used.
 
         Returns:
             DatasetVersion: The new Dataset Version created.