Merge pull request #9 from XpressAI/fahreza/vecto-data

Data Access API Update

Author: MFA-X-AI

tests/test_management.py

@@ -22,7 +22,7 @@
 
 from vecto.schema import (VectoIngestData, VectoEmbeddingData, VectoAttribute, VectoAnalogyStartEnd,
                     IngestResponse, LookupResult, VectoModel, VectoVectorSpace, VectoUser,
-                    VectoToken, VectoNewTokenResponse)
+                    VectoToken, VectoNewTokenResponse, DataPage, DataEntry)
 
 '''
 Please ensure that you have token, vecto_base_url and vector_space_id
@@ -51,6 +51,7 @@
 test_vector_space = ""
 test_vs_name = "management_sdk_test"
 test_vs_token = ""
+test_vs_data = ""
 
 # Currently test disabled as a bug prevents certain VS to be deleted.
 # @pytest.mark.management
@@ -152,6 +153,23 @@ def test_delete_token():
     token_list = user_vecto.list_tokens()
     assert not any(token.id == test_vs_token.id for token in token_list)
 
+@pytest.mark.management
+def test_listing():
+    global test_vs_data
+    logger.info("Check if vector space data can be listed and the return types are correct.")
+    test_vs_data = user_vecto.list_vector_space_data(vector_space_id, 10, 0)
+    assert isinstance(test_vs_data, DataPage)
+    assert isinstance(test_vs_data.elements[0], DataEntry)
+
+@pytest.mark.management
+def test_delete_data():
+    global test_vs_data
+    logger.info("Check if vector space data can be deleted.")
+    dataEntry = test_vs_data.elements[0]
+    user_vecto.delete_vector_space_entry(vector_space_id, dataEntry.id)
+    updated_vs_data = user_vecto.list_vector_space_data(vector_space_id, 10, 0)
+    updatedDataEntry = updated_vs_data.elements[0]
+    assert dataEntry != updatedDataEntry
 
 @pytest.mark.management
 def test_delete_vector_space():

vecto/schema.py

@@ -119,4 +119,14 @@ class MonthlyUsageResponse(NamedTuple):
     '''A named tuple that contains the usage metrics for a specified vector space and month.'''
     year: int
     month: int
-    usage: VectoUsageMetrics
+    usage: VectoUsageMetrics
+
+class DataEntry(NamedTuple):
+    '''A named tuple that represents an individual data entry with an ID and attributes.'''
+    id: int
+    attributes: dict
+
+class DataPage(NamedTuple):
+    '''A named tuple of data entries, contains the total count and a list of 'DataEntry' instances.'''
+    count: int
+    elements: List[DataEntry]

vecto/vecto_requests.py

@@ -28,7 +28,8 @@
 from .schema import (VectoIngestData, VectoEmbeddingData, VectoAttribute, VectoAnalogyStartEnd,
                     IngestResponse, LookupResult, VectoModel, VectoVectorSpace, VectoUser,
                     VectoToken, VectoNewTokenResponse, MODEL_MAP, VectoAnalogy, 
-                    DailyUsageMetric, UsageMetric, VectoUsageMetrics, MonthlyUsageResponse)
+                    DailyUsageMetric, UsageMetric, VectoUsageMetrics, MonthlyUsageResponse, 
+                    DataEntry, DataPage)
 
 from .client import Client
 import vecto
@@ -919,6 +920,44 @@ def delete_analogy(self, vector_space_id:int, **kwargs):
         url = f"/api/v0/account/space/{vector_space_id}/analogy"
         self._client.delete(url, **kwargs)
 
+    def list_vector_space_data(self, vector_space_id: int, limit: int = None, offset: int = None, **kwargs):
+        '''
+        List the attributes of all entries in the given vector space.
+
+        Args:
+            vector_space_id (int): The ID of the vector space.
+            limit (int, optional): The maximum number of entries to return. If not specified, it will be set to 0.
+            offset (int, optional): The offset from the start of the list to begin returning entries.
+            **kwargs: Other keyword arguments for clients other than `requests`
+
+        Returns:
+            DataPage: A DataPage instance containing the list of entries and their attributes.
+        '''
+
+        url = f"/api/v0/space/{vector_space_id}/data"
+        params = {'limit': limit, 'offset': offset}
+        response = self._client.get(url, params=params, **kwargs)
+        response_json = response.json()
+
+        # Create DataEntry instances for each element in the response
+        data_entries = [DataEntry(**entry) for entry in response_json["elements"]]
+
+        # Create and return the DataPage instance
+        return DataPage(count=response_json["count"], elements=data_entries)
+
+    def delete_vector_space_entry(self, vector_space_id: int, entry_id: int, **kwargs):
+        '''
+        Delete an entry in a vector space.
+
+        Args:
+            vector_space_id (int): The ID of the vector space.
+            entry_id (int): The ID of the entry to be deleted.
+            **kwargs: Other keyword arguments for clients other than `requests`
+        '''
+
+        url = f"/api/v0/space/{vector_space_id}/data/{entry_id}"
+        self._client.delete(url, **kwargs)
+
     ###############
     # Metrics API #
     ###############