diff --git a/aixplain/modules/model/index_model.py b/aixplain/modules/model/index_model.py index 67b3f8f7..0dcc7b4c 100644 --- a/aixplain/modules/model/index_model.py +++ b/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": ""} diff --git a/aixplain/modules/model/document_index.py b/aixplain/modules/model/record.py similarity index 96% rename from aixplain/modules/model/document_index.py rename to aixplain/modules/model/record.py index 12562931..a3c57173 100644 --- a/aixplain/modules/model/document_index.py +++ b/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 diff --git a/tests/functional/model/run_model_test.py b/tests/functional/model/run_model_test.py index 04e5da0d..16239726 100644 --- a/tests/functional/model/run_model_test.py +++ b/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 diff --git a/tests/unit/index_model_test.py b/tests/unit/index_model_test.py index be9acc6f..dbf698cc 100644 --- a/tests/unit/index_model_test.py +++ b/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,8 +28,8 @@ 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: @@ -37,7 +37,7 @@ def test_add_success(): 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