Merge pull request #62 from AitoDotAI/feature/estimate-and-modify-endpoints

Add estimate, aggregate, modify, and batch API endpoints

Author: arauhala

.circleci/config.yml

@@ -167,15 +167,15 @@ commonYAMLStructure:
           (E'Johnson\'s Automotive', 2344, NULL);"
   - &run_sdk_and_cli_tests
     run:
-      name: run all tests execluding sql functions tests
+      name: run all tests excluding sql functions tests
       command: |
         [ "$CIRCLE_BRANCH" != master ] && git diff --quiet HEAD origin/master -- . ":!docs" && echo "NO CHANGE - SKIP" && exit 0
         . venv/bin/activate
         python -m tests -v suite sdk
         python -m tests -v suite cli
   - &run_sdk_and_cli_tests_windows
     run:
-      name: run all tests execluding sql functions tests
+      name: run all tests excluding sql functions tests
       command: |
         if ($Env:CIRCLE_BRANCH -ne "master") {
           (git diff --quiet HEAD origin/master -- . ":!docs") -and (echo "NO CHANGE - SKIP") -and (exit 0)

.gitignore

@@ -117,6 +117,9 @@ aito-python-tools.iml
 # vscode
 .vscode/
 
+# AI assistant files
+.ai/
+
 # test ouput file
 *_out.*
 

aito/__init__.py

@@ -1 +1 @@
-__version__ = "0.5.1"
+__version__ = "0.5.2"

aito/api.py

@@ -1167,3 +1167,153 @@ def generic_query(
     if use_job:
         return job_request(client=client, request_obj=req)
     return client.request(request_obj=req, raise_for_status=raise_for_status)
+
+
+def estimate(
+        client: AitoClient, query: Dict, raise_for_status: Optional[bool] = None, use_job: bool = False
+) -> Union[aito_responses.EstimateResponse, RequestError]:
+    """send a query to the `Estimate API <https://aito.ai/docs/api/#post-api-v1-estimate>`__
+
+    The Estimate operation predicts numeric values using K-NN with transformation pipelines.
+
+    :param client: the AitoClient instance
+    :type client: AitoClient
+    :param query: the estimate query
+    :type query: Dict
+    :param raise_for_status: raise :class:`.RequestError` if the request fails instead of returning the error
+        If set to None, value from Client will be used. Defaults to None
+    :type raise_for_status: Optional[bool]
+    :param use_job: use job fo request that takes longer than 30 seconds, defaults to False
+    :type use_job: bool
+    :return: :class:`.EstimateResponse` or :class:`.RequestError` if an error occurred and not raise_for_status
+    :rtype: Union[EstimateResponse, RequestError]
+    """
+    req = aito_requests.EstimateRequest(query)
+    if use_job:
+        return job_request(client=client, request_obj=req)
+    return client.request(request_obj=req, raise_for_status=raise_for_status)
+
+
+def aggregate(
+        client: AitoClient, query: Dict, raise_for_status: Optional[bool] = None, use_job: bool = False
+) -> Union[aito_responses.AggregateResponse, RequestError]:
+    """send a query to the `Aggregate API <https://aito.ai/docs/api/#post-api-v1-aggregate>`__
+
+    Returns aggregated values like averages, sums, counts, etc.
+
+    :param client: the AitoClient instance
+    :type client: AitoClient
+    :param query: the aggregate query
+    :type query: Dict
+    :param raise_for_status: raise :class:`.RequestError` if the request fails instead of returning the error
+        If set to None, value from Client will be used. Defaults to None
+    :type raise_for_status: Optional[bool]
+    :param use_job: use job fo request that takes longer than 30 seconds, defaults to False
+    :type use_job: bool
+    :return: :class:`.AggregateResponse` or :class:`.RequestError` if an error occurred and not raise_for_status
+    :rtype: Union[AggregateResponse, RequestError]
+    """
+    req = aito_requests.AggregateRequest(query)
+    if use_job:
+        return job_request(client=client, request_obj=req)
+    return client.request(request_obj=req, raise_for_status=raise_for_status)
+
+
+def modify(
+        client: AitoClient,
+        query: Union[Dict, 'aito_requests.ModifyOperation', List['aito_requests.ModifyOperation']],
+        raise_for_status: Optional[bool] = None,
+        use_job: bool = False
+) -> Union[aito_responses.ModifyResponse, RequestError]:
+    """perform atomic modifications on table data using the
+    `Modify API <https://aito.ai/docs/api/#post-api-v1-data-modify>`__
+
+    The modify endpoint allows individual modifications or a sequence
+    of modifications in one atomic operation.
+
+    Can accept a raw query dict, a single ModifyOperation, or a list of ModifyOperations::
+
+        from aito.client.requests import Insert, Update, Delete
+
+        # Insert a single entry
+        api.modify(client, Insert("products", {"id": "1", "name": "Apple"}))
+
+        # Update entries matching a condition
+        api.modify(client, Update("products").where({"id": "1"}).set({"name": "New Name"}))
+
+        # Delete entries matching a condition
+        api.modify(client, Delete("products", {"id": "1"}))
+
+        # Multiple operations atomically
+        api.modify(client, [
+            Insert("products", {"id": "1", "name": "Apple"}),
+            Update("products").where({"id": "2"}).set({"name": "Banana"}),
+            Delete("products", {"id": "3"})
+        ])
+
+    :param client: the AitoClient instance
+    :type client: AitoClient
+    :param query: the modify query - can be a Dict, a ModifyOperation (Insert/Update/Delete),
+        or a list of ModifyOperations
+    :type query: Union[Dict, ModifyOperation, List[ModifyOperation]]
+    :param raise_for_status: raise :class:`.RequestError` if the request fails instead of returning the error
+        If set to None, value from Client will be used. Defaults to None
+    :type raise_for_status: Optional[bool]
+    :param use_job: use job fo request that takes longer than 30 seconds, defaults to False
+    :type use_job: bool
+    :return: :class:`.ModifyResponse` or :class:`.RequestError` if an error occurred and not raise_for_status
+    :rtype: Union[ModifyResponse, RequestError]
+    """
+    # Convert ModifyOperation(s) to query dict
+    if isinstance(query, aito_requests.ModifyOperation):
+        query = query.to_query()
+    elif isinstance(query, list) and len(query) > 0 and isinstance(query[0], aito_requests.ModifyOperation):
+        # Multiple operations must be wrapped in {"operations": [...]}
+        query = {"operations": [op.to_query() for op in query]}
+    elif isinstance(query, list):
+        # Raw list of dicts also needs wrapping
+        query = {"operations": query}
+
+    req = aito_requests.ModifyRequest(query)
+    if use_job:
+        return job_request(client=client, request_obj=req)
+    return client.request(request_obj=req, raise_for_status=raise_for_status)
+
+
+def batch(
+        client: AitoClient,
+        queries: List[Dict],
+        raise_for_status: Optional[bool] = None
+) -> Union[aito_responses.BatchResponse, RequestError]:
+    """execute multiple queries in a single request using the
+    `Batch API <https://aito.ai/docs/api/#post-api-v1-batch>`__
+
+    The batch endpoint allows executing multiple queries (predict, similarity,
+    search, etc.) in a single HTTP request for improved performance.
+
+    Example::
+
+        import aito.api as api
+
+        # Execute multiple queries in one request
+        results = api.batch(client, [
+            {"from": "products", "where": {"name": "rye bread"}, "predict": "category"},
+            {"from": "products", "similarity": {"name": "rye bread"}, "limit": 5}
+        ])
+
+        # Access individual results
+        predict_result = results[0]  # First query result
+        similarity_result = results[1]  # Second query result
+
+    :param client: the AitoClient instance
+    :type client: AitoClient
+    :param queries: a list of query dictionaries to execute in batch
+    :type queries: List[Dict]
+    :param raise_for_status: raise :class:`.RequestError` if the request fails instead of returning the error
+        If set to None, value from Client will be used. Defaults to None
+    :type raise_for_status: Optional[bool]
+    :return: :class:`.BatchResponse` containing results for each query, or :class:`.RequestError` if an error occurred
+    :rtype: Union[BatchResponse, RequestError]
+    """
+    req = aito_requests.BatchRequest(queries)
+    return client.request(request_obj=req, raise_for_status=raise_for_status)

aito/cli/main_parser.py

@@ -17,7 +17,7 @@
     UploadEntriesSubCommand, UploadBatchSubCommand, UploadFileSubCommand, UploadDataFromSQLSubCommand, OptimizeTableSubCommand, \
     QuickAddTableFromSQLSubCommand, GetDatabaseSubCommand, GetTableSubCommand, QuickPredictSubCommand, SearchSubCommand, \
     PredictSubCommand, RecommendSubCommand, EvaluateSubCommand, SimilaritySubCommand, MatchSubCommand, RelateSubCommand, \
-    GenericQuerySubCommand, CreateDatabaseSubCommand
+    GenericQuerySubCommand, CreateDatabaseSubCommand, EstimateSubCommand, AggregateSubCommand
 from .sub_commands.infer_table_schema_sub_command import InferTableSchemaSubCommand
 from .sub_commands.sub_command import SubCommand
 
@@ -51,7 +51,9 @@ class MainParser(ArgParser):
         SimilaritySubCommand(),
         MatchSubCommand(),
         RelateSubCommand(),
-        GenericQuerySubCommand()
+        GenericQuerySubCommand(),
+        EstimateSubCommand(),
+        AggregateSubCommand()
     ]
 
     def __init__(self, commands: List[SubCommand] = None):

aito/cli/sub_commands/database_sub_command.py

@@ -445,3 +445,5 @@ def parse_and_execute(self, parsed_args: Dict):
 GenericQuerySubCommand = type(
     'GenericQuerySubCommand', (QueryToEndpointSubCommand,), {'api_method_name': 'generic_query'}
 )
+EstimateSubCommand = type('EstimateSubCommand', (QueryToEndpointSubCommand,), {'api_method_name': 'estimate'})
+AggregateSubCommand = type('AggregateSubCommand', (QueryToEndpointSubCommand,), {'api_method_name': 'aggregate'})

aito/client/requests/__init__.py

@@ -3,10 +3,12 @@
 """
 
 from .data_api_request import DataAPIRequest, UploadEntriesRequest, DeleteEntriesRequest, InitiateFileUploadRequest, \
-    TriggerFileProcessingRequest, GetFileProcessingRequest
+    TriggerFileProcessingRequest, GetFileProcessingRequest, ModifyRequest
+from .modify_operations import ModifyOperation, Insert, Update, Delete
 from .job_api_request import CreateJobRequest, GetJobResultRequest, GetJobStatusRequest
 from .query_api_request import QueryAPIRequest, SearchRequest, PredictRequest, RecommendRequest, EvaluateRequest, \
-    SimilarityRequest, MatchRequest, RelateRequest, GenericQueryRequest
+    SimilarityRequest, MatchRequest, RelateRequest, GenericQueryRequest, EstimateRequest, AggregateRequest, \
+    BatchRequest
 from .schema_api_request import GetDatabaseSchemaRequest, GetTableSchemaRequest, GetColumnSchemaRequest, \
     CreateDatabaseSchemaRequest, CreateTableSchemaRequest, CreateColumnSchemaRequest, DeleteDatabaseSchemaRequest, \
     DeleteTableSchemaRequest, DeleteColumnSchemaRequest

aito/client/requests/data_api_request.py

@@ -7,6 +7,7 @@
 from aito.client import responses as aito_resp
 from .aito_request import AitoRequest, _FinalRequest, _PatternEndpoint, _GetRequest, _PostRequest
 from aito.schema import AitoSchema
+from ..responses.query_api_response import ModifyResponse
 
 
 class DataAPIRequest(AitoRequest, ABC):
@@ -174,3 +175,28 @@ def path_matches(cls, path_to_match):
 
 class GetFileProcessingRequest(_GetRequest, TriggerFileProcessingRequest, DataAPIRequest):
     """Request to `Initiate File Upload <https://aito.ai/docs/api/#post-api-v1-data-table-file>`__"""
+
+
+class ModifyRequest(_PostRequest, _FinalRequest, DataAPIRequest):
+    """Request to `Modify data atomically <https://aito.ai/docs/api/#post-api-v1-data-modify>`__
+
+    Allows individual modifications or a sequence of modifications in one atomic operation.
+    """
+    _path_suffix = '_modify'
+    endpoint = f'{DataAPIRequest.endpoint_prefix}/{_path_suffix}'
+    response_cls = ModifyResponse
+
+    def __init__(self, query: Dict):
+        """
+        :param query: The modify query with operations
+        :type query: Dict
+        """
+        super().__init__(query=query)
+
+    @classmethod
+    def _endpoint_pattern(cls):
+        return re.compile(f'^{cls.endpoint}$')
+
+    @classmethod
+    def path_matches(cls, path_to_match):
+        return path_to_match.endswith(f'/{cls._path_suffix}')