Version 1.0.5

abacusai · Jan 12, 2024 · b026817 · b026817
1 parent fa6b982
commit b026817
Show file tree

Hide file tree

Showing 248 changed files with 1,342 additions and 166 deletions.
diff --git a/abacusai/__init__.py b/abacusai/__init__.py
@@ -1,7 +1,7 @@
 from .api_class import *
-from .client import ApiClient, ApiException, ClientOptions, ReadOnlyClient, _request_context
+from .client import AgentResponse, ApiClient, ApiException, ClientOptions, ReadOnlyClient, _request_context
 from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.0.2"
+__version__ = "1.0.5"
diff --git a/abacusai/agent.py b/abacusai/agent.py
@@ -19,11 +19,12 @@ class Agent(AbstractApiClass):
             agentConfig (dict): The config options used to create this agent.
             memory (int): Memory in GB specified for the deployment resources for the agent.
             trainingRequired (bool): Whether training is required to deploy the latest agent code.
+            agentExecutionConfig (dict): The config for arguments used to execute the agent.
             latestAgentVersion (AgentVersion): The latest agent version.
             codeSource (CodeSource): If a python model, information on the source code
     """
 
-    def __init__(self, client, name=None, agentId=None, createdAt=None, projectId=None, notebookId=None, predictFunctionName=None, sourceCode=None, agentConfig=None, memory=None, trainingRequired=None, codeSource={}, latestAgentVersion={}):
+    def __init__(self, client, name=None, agentId=None, createdAt=None, projectId=None, notebookId=None, predictFunctionName=None, sourceCode=None, agentConfig=None, memory=None, trainingRequired=None, agentExecutionConfig=None, codeSource={}, latestAgentVersion={}):
         super().__init__(client, agentId)
         self.name = name
         self.agent_id = agentId
@@ -35,13 +36,14 @@ def __init__(self, client, name=None, agentId=None, createdAt=None, projectId=No
         self.agent_config = agentConfig
         self.memory = memory
         self.training_required = trainingRequired
+        self.agent_execution_config = agentExecutionConfig
         self.code_source = client._build_class(CodeSource, codeSource)
         self.latest_agent_version = client._build_class(
             AgentVersion, latestAgentVersion)
 
     def __repr__(self):
         repr_dict = {f'name': repr(self.name), f'agent_id': repr(self.agent_id), f'created_at': repr(self.created_at), f'project_id': repr(self.project_id), f'notebook_id': repr(self.notebook_id), f'predict_function_name': repr(self.predict_function_name), f'source_code': repr(
-            self.source_code), f'agent_config': repr(self.agent_config), f'memory': repr(self.memory), f'training_required': repr(self.training_required), f'code_source': repr(self.code_source), f'latest_agent_version': repr(self.latest_agent_version)}
+            self.source_code), f'agent_config': repr(self.agent_config), f'memory': repr(self.memory), f'training_required': repr(self.training_required), f'agent_execution_config': repr(self.agent_execution_config), f'code_source': repr(self.code_source), f'latest_agent_version': repr(self.latest_agent_version)}
         class_name = "Agent"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None])
@@ -54,8 +56,8 @@ def to_dict(self):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'name': self.name, 'agent_id': self.agent_id, 'created_at': self.created_at, 'project_id': self.project_id, 'notebook_id': self.notebook_id, 'predict_function_name': self.predict_function_name, 'source_code': self.source_code,
-                'agent_config': self.agent_config, 'memory': self.memory, 'training_required': self.training_required, 'code_source': self._get_attribute_as_dict(self.code_source), 'latest_agent_version': self._get_attribute_as_dict(self.latest_agent_version)}
+        resp = {'name': self.name, 'agent_id': self.agent_id, 'created_at': self.created_at, 'project_id': self.project_id, 'notebook_id': self.notebook_id, 'predict_function_name': self.predict_function_name, 'source_code': self.source_code, 'agent_config': self.agent_config,
+                'memory': self.memory, 'training_required': self.training_required, 'agent_execution_config': self.agent_execution_config, 'code_source': self._get_attribute_as_dict(self.code_source), 'latest_agent_version': self._get_attribute_as_dict(self.latest_agent_version)}
         return {key: value for key, value in resp.items() if value is not None}
 
     def refresh(self):

diff --git a/abacusai/agent_version.py b/abacusai/agent_version.py
@@ -17,10 +17,11 @@ class AgentVersion(AbstractApiClass):
             pendingDeploymentIds (list): List of deployment IDs where deployment is pending.
             failedDeploymentIds (list): List of failed deployment IDs.
             error (str): Relevant error if the status is FAILED.
+            agentExecutionConfig (dict): The config for arguments used to execute the agent.
             codeSource (CodeSource): If a python model, information on where the source code is located.
     """
 
-    def __init__(self, client, agentVersion=None, status=None, agentId=None, agentConfig=None, publishingStartedAt=None, publishingCompletedAt=None, pendingDeploymentIds=None, failedDeploymentIds=None, error=None, codeSource={}):
+    def __init__(self, client, agentVersion=None, status=None, agentId=None, agentConfig=None, publishingStartedAt=None, publishingCompletedAt=None, pendingDeploymentIds=None, failedDeploymentIds=None, error=None, agentExecutionConfig=None, codeSource={}):
         super().__init__(client, agentVersion)
         self.agent_version = agentVersion
         self.status = status
@@ -31,11 +32,12 @@ def __init__(self, client, agentVersion=None, status=None, agentId=None, agentCo
         self.pending_deployment_ids = pendingDeploymentIds
         self.failed_deployment_ids = failedDeploymentIds
         self.error = error
+        self.agent_execution_config = agentExecutionConfig
         self.code_source = client._build_class(CodeSource, codeSource)
 
     def __repr__(self):
         repr_dict = {f'agent_version': repr(self.agent_version), f'status': repr(self.status), f'agent_id': repr(self.agent_id), f'agent_config': repr(self.agent_config), f'publishing_started_at': repr(self.publishing_started_at), f'publishing_completed_at': repr(
-            self.publishing_completed_at), f'pending_deployment_ids': repr(self.pending_deployment_ids), f'failed_deployment_ids': repr(self.failed_deployment_ids), f'error': repr(self.error), f'code_source': repr(self.code_source)}
+            self.publishing_completed_at), f'pending_deployment_ids': repr(self.pending_deployment_ids), f'failed_deployment_ids': repr(self.failed_deployment_ids), f'error': repr(self.error), f'agent_execution_config': repr(self.agent_execution_config), f'code_source': repr(self.code_source)}
         class_name = "AgentVersion"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None])
@@ -49,7 +51,7 @@ def to_dict(self):
             dict: The dict value representation of the class parameters
         """
         resp = {'agent_version': self.agent_version, 'status': self.status, 'agent_id': self.agent_id, 'agent_config': self.agent_config, 'publishing_started_at': self.publishing_started_at, 'publishing_completed_at': self.publishing_completed_at,
-                'pending_deployment_ids': self.pending_deployment_ids, 'failed_deployment_ids': self.failed_deployment_ids, 'error': self.error, 'code_source': self._get_attribute_as_dict(self.code_source)}
+                'pending_deployment_ids': self.pending_deployment_ids, 'failed_deployment_ids': self.failed_deployment_ids, 'error': self.error, 'agent_execution_config': self.agent_execution_config, 'code_source': self._get_attribute_as_dict(self.code_source)}
         return {key: value for key, value in resp.items() if value is not None}
 
     def refresh(self):

diff --git a/abacusai/api_class/abstract.py b/abacusai/api_class/abstract.py
@@ -140,6 +140,10 @@ def from_dict(cls, input_dict: dict):
                 return builder.from_dict(input_dict)
             if not cls._upper_snake_case_keys:
                 input_dict = {snake_case(k): v for k, v in input_dict.items()}
+            if not cls._support_kwargs:
+                # only use keys that are valid fields in the ApiClass
+                field_names = set((field.name) for field in dataclasses.fields(cls))
+                input_dict = {k: v for k, v in input_dict.items() if k in field_names}
         return cls(**input_dict)
 
 

diff --git a/abacusai/api_class/enums.py b/abacusai/api_class/enums.py
@@ -391,9 +391,9 @@ class LLMName(ApiEnum):
     LLAMA2_CHAT = 'LLAMA2_CHAT'
     PALM = 'PALM'
     PALM_TEXT = 'PALM_TEXT'
+    GEMINI_PRO = 'GEMINI_PRO'
     MIXTRAL_CHAT = 'MIXTRAL_CHAT'
-    MISTRAL = 'MISTRAL'
-    GEMINI = 'GEMINI'
+    MISTRAL_MEDIUM = 'MISTRAL_MEDIUM'
 
 
 class MonitorAlertType(ApiEnum):
@@ -411,6 +411,9 @@ class FeatureDriftType(ApiEnum):
     KS = 'ks'
     WS = 'ws'
     JS = 'js'
+    PSI = 'psi'
+    CHI_SQUARE = 'chi_square'
+    CSI = 'csi'
 
 
 class DataIntegrityViolationType(ApiEnum):

diff --git a/abacusai/api_class/model.py b/abacusai/api_class/model.py
@@ -443,6 +443,8 @@ class ChatLLMTrainingConfig(TrainingConfig):
         behavior_instructions (str): Customize the overall role instructions for the LLM.
         response_instructions (str): Customize instructions for what the LLM responses should look like.
         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.
+        data_prompt_context (str): Prompt context for the data feature group ids.
     """
     document_retrievers: List[str] = None
     num_completion_tokens: int = None
@@ -452,6 +454,8 @@ class ChatLLMTrainingConfig(TrainingConfig):
     behavior_instructions: str = None
     response_instructions: str = None
     max_search_results: int = None
+    data_feature_group_ids: List[str] = None
+    data_prompt_context: str = None
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.CHAT_LLM