ENG-1393: Index improvements (#360)

* Index improvements

* Remove add/update methods

Author: thiago-aixplain

aixplain/modules/model/index_model.py

@@ -3,7 +3,7 @@
 from aixplain.utils import config
 from aixplain.modules.model.response import ModelResponse
 from typing import Text, Optional, Union, Dict
-from aixplain.modules.model.document_index import DocumentIndex
+from aixplain.modules.model.record import Record
 from typing import List
 
 
@@ -55,26 +55,14 @@ def search(self, query: str, top_k: int = 10) -> ModelResponse:
         data = {"action": "search", "data": query, "payload": {"filters": {}, "top_k": top_k}}
         return self.run(data=data)
 
-    def add(self, documents: List[DocumentIndex]) -> ModelResponse:
+    def upsert(self, documents: List[Record]) -> ModelResponse:
         payloads = [doc.to_dict() for doc in documents]
         data = {"action": "ingest", "data": "", "payload": {"payloads": payloads}}
         response = self.run(data=data)
         if response.status == ResponseStatus.SUCCESS:
             response.data = payloads
             return response
-        raise Exception(f"Failed to add documents: {response.error_message}")
-
-    def update(self, documents: List[DocumentIndex]) -> ModelResponse:
-        payloads = [
-            {"value": doc.value, "value_type": doc.value_type, "id": str(doc.id), "uri": doc.uri, "attributes": doc.attributes}
-            for doc in documents
-        ]
-        data = {"action": "update", "data": "", "payload": {"payloads": payloads}}
-        response = self.run(data=data)
-        if response.status == ResponseStatus.SUCCESS:
-            response.data = payloads
-            return response
-        raise Exception(f"Failed to update documents: {response.error_message}")
+        raise Exception(f"Failed to upsert documents: {response.error_message}")
 
     def count(self) -> float:
         data = {"action": "count", "data": ""}

aixplain/modules/model/document_index.py aixplain/modules/model/record.py

@@ -2,7 +2,7 @@
 from uuid import uuid4
 
 
-class DocumentIndex:
+class Record:
     def __init__(self, value: str, value_type: str = "text", id: Optional[str] = None, uri: str = "", attributes: dict = {}):
         self.value = value
         self.value_type = value_type

tests/functional/model/run_model_test.py

@@ -55,11 +55,11 @@ def test_run_async():
 
 
 def test_index_model():
-    from aixplain.modules.model.document_index import DocumentIndex
+    from aixplain.modules.model.record import Record
     from aixplain.factories import IndexFactory
 
     index_model = IndexFactory.create("test", "test")
-    index_model.add([DocumentIndex(value="Hello, world!", value_type="text", uri="", attributes={})])
+    index_model.upsert([Record(value="Hello, world!", value_type="text", uri="", attributes={})])
     response = index_model.search("Hello")
     assert str(response.status) == "SUCCESS"
     assert index_model.count() == 1

tests/unit/index_model_test.py

@@ -1,6 +1,6 @@
 import requests_mock
 from aixplain.enums import Function, ResponseStatus
-from aixplain.modules.model.document_index import DocumentIndex
+from aixplain.modules.model.record import Record
 from aixplain.modules.model.response import ModelResponse
 from aixplain.modules.model.index_model import IndexModel
 from aixplain.utils import config
@@ -28,16 +28,16 @@ def test_add_success():
     mock_response = {"status": "SUCCESS"}
 
     mock_documents = [
-        DocumentIndex(value="Sample document content 1", value_type="text", id=0, uri="", attributes={}),
-        DocumentIndex(value="Sample document content 2", value_type="text", id=1, uri="", attributes={}),
+        Record(value="Sample document content 1", value_type="text", id=0, uri="", attributes={}),
+        Record(value="Sample document content 2", value_type="text", id=1, uri="", attributes={}),
     ]
 
     with requests_mock.Mocker() as mock:
         mock.post(execute_url, json=mock_response, status_code=200)
 
         index_model = IndexModel(id=index_id, data=data, name="name", function=Function.SEARCH)
 
-        response = index_model.add(mock_documents)
+        response = index_model.upsert(mock_documents)
 
     assert isinstance(response, ModelResponse)
     assert response.status == ResponseStatus.SUCCESS
@@ -47,8 +47,8 @@ def test_update_success():
     mock_response = {"status": "SUCCESS"}
 
     mock_documents = [
-        DocumentIndex(value="Updated document content 1", value_type="text", id=0, uri="", attributes={}),
-        DocumentIndex(value="Updated document content 2", value_type="text", id=1, uri="", attributes={}),
+        Record(value="Updated document content 1", value_type="text", id=0, uri="", attributes={}),
+        Record(value="Updated document content 2", value_type="text", id=1, uri="", attributes={}),
     ]
 
     with requests_mock.Mocker() as mock:
@@ -57,7 +57,7 @@ def test_update_success():
 
         index_model = IndexModel(id=index_id, data=data, name="name", function=Function.SEARCH)
 
-        response = index_model.update(mock_documents)
+        response = index_model.upsert(mock_documents)
 
     assert isinstance(response, ModelResponse)
     assert response.status == ResponseStatus.SUCCESS