Merge to test

aixplain/modules/model/index_model.py

@@ -151,7 +151,7 @@ def search(self, query: str, top_k: int = 10, filters: List[IndexFilter] = []) -
             "data": query or uri,
             "dataType": value_type,
             "filters": [filter.to_dict() for filter in filters],
-            "payload": {"uri": uri, "value_type": value_type, "top_k": top_k}
+            "payload": {"uri": uri, "value_type": value_type, "top_k": top_k},
         }
         return self.run(data=data)
 
@@ -246,3 +246,49 @@ def delete_record(self, record_id: Text) -> ModelResponse:
         if response.status == "SUCCESS":
             return response
         raise Exception(f"Failed to delete record: {response.error_message}")
+
+    def retrieve_records_with_filter(self, filter: IndexFilter) -> ModelResponse:
+        """
+        Retrieve records from the index that match the given filter.
+
+        Args:
+            filter (IndexFilter): The filter criteria to apply when retrieving records.
+
+        Returns:
+            ModelResponse: Response containing the retrieved records.
+
+        Raises:
+            Exception: If retrieval fails.
+
+        Example:
+            >>> from aixplain.modules.model.index_model import IndexFilter, IndexFilterOperator
+            >>> my_filter = IndexFilter(field="category", value="world", operator=IndexFilterOperator.EQUALS)
+            >>> index_model.retrieve_records_with_filter(my_filter)
+        """
+        data = {"action": "retrieve_by_filter", "data": filter.to_dict()}
+        response = self.run(data=data)
+        if response.status == "SUCCESS":
+            return response
+        raise Exception(f"Failed to retrieve records with filter: {response.error_message}")
+
+    def delete_records_by_date(self, date: float) -> ModelResponse:
+        """
+        Delete records from the index that match the given date.
+
+        Args:
+            date (float): The date (as a timestamp) to match records for deletion.
+
+        Returns:
+            ModelResponse: Response containing the result of the deletion operation.
+
+        Raises:
+            Exception: If deletion fails.
+
+        Example:
+            >>> index_model.delete_records_by_date(1717708800)
+        """
+        data = {"action": "delete_by_date", "data": date}
+        response = self.run(data=data)
+        if response.status == "SUCCESS":
+            return response
+        raise Exception(f"Failed to delete records by date: {response.error_message}")

aixplain/modules/model/integration.py

@@ -5,48 +5,24 @@
 from typing import Text, Optional, Union, Dict
 from enum import Enum
 from pydantic import BaseModel
-
+import json
 
 class AuthenticationSchema(Enum):
-    BEARER = "BEARER_TOKEN"
-    OAUTH = "OAUTH"
+    BEARER_TOKEN = "BEARER_TOKEN"
+    OAUTH1 = "OAUTH1"
     OAUTH2 = "OAUTH2"
-
-
+    API_KEY = "API_KEY"
+    BASIC = "BASIC"
+    NO_AUTH = "NO_AUTH"
+
 class BaseAuthenticationParams(BaseModel):
     name: Optional[Text] = None
-    authentication_schema: AuthenticationSchema = AuthenticationSchema.OAUTH2
     connector_id: Optional[Text] = None
 
-
-class BearerAuthenticationParams(BaseAuthenticationParams):
-    token: Text
-    authentication_schema: AuthenticationSchema = AuthenticationSchema.BEARER
-
-
-class OAuthAuthenticationParams(BaseAuthenticationParams):
-    client_id: Text
-    client_secret: Text
-    authentication_schema: AuthenticationSchema = AuthenticationSchema.OAUTH
-
-
-class OAuth2AuthenticationParams(BaseAuthenticationParams):
-    authentication_schema: AuthenticationSchema = AuthenticationSchema.OAUTH2
-
-
 def build_connector_params(**kwargs) -> BaseAuthenticationParams:
     name = kwargs.get("name")
-    token = kwargs.get("token")
-    client_id = kwargs.get("client_id")
-    client_secret = kwargs.get("client_secret")
     connector_id = kwargs.get("connector_id")
-    if token:
-        args = BearerAuthenticationParams(name=name, token=token, connector_id=connector_id)
-    elif client_id and client_secret:
-        args = OAuthAuthenticationParams(name=name, client_id=client_id, client_secret=client_secret, connector_id=connector_id)
-    else:
-        args = OAuth2AuthenticationParams(name=name, connector_id=connector_id)
-    return args
+    return BaseAuthenticationParams(name=name, connector_id=connector_id)
 
 
 class Integration(Model):
@@ -94,8 +70,10 @@ def __init__(
         )
         self.url = config.MODELS_RUN_URL
         self.backend_url = config.BACKEND_URL
+        self.authentication_methods = json.loads([item for item in additional_info['attributes'] if item['name'] == 'auth_schemes'][0]['code'])
+
 
-    def connect(self, args: Optional[BaseAuthenticationParams] = None, **kwargs) -> ModelResponse:
+    def connect(self, authentication_schema: AuthenticationSchema, args: Optional[BaseAuthenticationParams] = None, data: Optional[Dict] = None, **kwargs) -> ModelResponse:
         """Connect to the integration
 
         Examples:
@@ -119,29 +97,29 @@ def connect(self, args: Optional[BaseAuthenticationParams] = None, **kwargs) ->
         if args is None:
             args = build_connector_params(**kwargs)
 
-        authentication_schema = args.authentication_schema
-        if authentication_schema == AuthenticationSchema.BEARER:
+        if authentication_schema.value not in self.authentication_methods:
+            raise ValueError(f"Authentication schema {authentication_schema.value} is not supported for this integration. Supported authentication methods: {self.authentication_methods}")
+
+        if data is None:
+            data = {}
+
+        if authentication_schema not in [AuthenticationSchema.OAUTH2, AuthenticationSchema.OAUTH1, AuthenticationSchema.NO_AUTH]:
+            required_params = json.loads([item for item in self.additional_info['attributes'] if item['name'] == authentication_schema.value + "-inputs"][0]['code'])
+            required_params_names = [param['name'] for param in required_params]
+            for param in required_params_names:
+                if param not in data:
+                    if len(required_params_names) == 1:
+                        raise ValueError(f"Parameter '{param}' is required for {self.name} {authentication_schema.value} authentication. Please provide the parameter in the data dictionary.")
+                    else:
+                        raise ValueError(f"Parameters {required_params_names} are required for {self.name} {authentication_schema.value} authentication. Please provide the parameters in the data dictionary.")
             return self.run(
                 {
                     "name": args.name,
                     "authScheme": authentication_schema.value,
-                    "data": {
-                        "token": args.token,
-                    },
+                    "data": data,
                 }
             )
-        elif authentication_schema == AuthenticationSchema.OAUTH:
-            return self.run(
-                {
-                    "name": args.name,
-                    "authScheme": authentication_schema.value,
-                    "data": {
-                        "client_id": args.client_id,
-                        "client_secret": args.client_secret,
-                    },
-                }
-            )
-        elif authentication_schema == AuthenticationSchema.OAUTH2:
+        else:
             response = self.run(
                 {
                     "name": args.name,
@@ -154,8 +132,9 @@ def connect(self, args: Optional[BaseAuthenticationParams] = None, **kwargs) ->
                 )
             return response
 
+
     def __repr__(self):
         try:
             return f"Integration: {self.name} by {self.supplier['name']} (id={self.id})"
         except Exception:
-            return f"Integration: {self.name} by {self.supplier} (id={self.id})"
+            return f"Integration: {self.name} by {self.supplier} (id={self.id})"

pyproject.toml

@@ -67,7 +67,7 @@ aixplain = "aixplain.cli_groups:run_cli"
 
 [project.optional-dependencies]
 model-builder = [
-    "model-interfaces~=0.0.1"
+    "model-interfaces~=0.0.2"
 ]
 test = [
     "pytest>=6.1.0",

tests/functional/model/run_model_test.py

@@ -322,3 +322,97 @@ def test_index_model_air_with_splitter(embedding_model, supplier_params):
     assert str(response.status) == "SUCCESS"
     assert "berlin" in response.data.lower()
     index_model.delete()
+
+
+def _test_records():
+    from aixplain.modules.model.record import Record
+    from aixplain.enums import DataType
+
+    return [
+        Record(
+            value="Artificial intelligence is transforming industries worldwide, from healthcare to finance.",
+            value_type=DataType.TEXT,
+            id="doc1",
+            uri="",
+            attributes={"category": "technology", "date": 1751464788},
+        ),
+        Record(
+            value="The Mona Lisa, painted by Leonardo da Vinci, is one of the most famous artworks in history.",
+            value_type=DataType.TEXT,
+            id="doc2",
+            uri="",
+            attributes={"category": "art", "date": 1751464790},
+        ),
+        Record(
+            value="Machine learning algorithms are being used to predict patient outcomes in hospitals.",
+            value_type=DataType.TEXT,
+            id="doc3",
+            uri="",
+            attributes={"category": "technology", "date": 1751464795},
+        ),
+        Record(
+            value="The Earth orbits the Sun once every 365.25 days, creating the calendar year.",
+            value_type=DataType.TEXT,
+            id="doc4",
+            uri="",
+            attributes={"category": "science", "date": 1751464798},
+        ),
+        Record(
+            value="Quantum computing promises to solve complex problems that are currently intractable for classical computers.",
+            value_type=DataType.TEXT,
+            id="doc5",
+            uri="",
+            attributes={"category": "technology", "date": 1751464801},
+        ),
+    ]
+
+
+@pytest.fixture(scope="function")
+def setup_index_with_test_records():
+    from aixplain.factories import IndexFactory
+    from aixplain.enums import EmbeddingModel
+    from aixplain.factories.index_factory.utils import AirParams
+    from uuid import uuid4
+    import time
+
+    # Clean up all existing indexes
+    for index in IndexFactory.list()["results"]:
+        index.delete()
+
+    params = AirParams(
+        name=f"Test Index {uuid4()}",
+        description="Test index for filter/date tests",
+        embedding_model=EmbeddingModel.OPENAI_ADA002,
+    )
+    index_model = IndexFactory.create(params=params)
+    records = _test_records()
+
+    index_model.upsert(records)
+
+    yield index_model
+    index_model.delete()
+
+
+def test_retrieve_records_with_filter(setup_index_with_test_records):
+    from aixplain.modules.model.index_model import IndexFilter, IndexFilterOperator
+
+    index_model = setup_index_with_test_records
+    filter_ = IndexFilter(field="category", value="technology", operator=IndexFilterOperator.EQUALS)
+    response = index_model.retrieve_records_with_filter(filter_)
+    assert response.status == "SUCCESS"
+    assert len(response.details) == 3
+    for item in response.details:
+        assert item["metadata"]["category"] == "technology"
+
+
+def test_delete_records_by_date(setup_index_with_test_records):
+    from aixplain.modules.model.index_model import IndexFilter, IndexFilterOperator
+
+    index_model = setup_index_with_test_records
+    response = index_model.delete_records_by_date(1751464796)
+    assert response.status == "SUCCESS"
+    assert response.data == "2"  # 2 records should remain
+    filter_all = IndexFilter(field="date", value=0, operator=IndexFilterOperator.GREATER_THAN)
+    response = index_model.retrieve_records_with_filter(filter_all)
+    assert response.status == "SUCCESS"
+    assert len(response.details) == 2