ENG-1789: Add Vectara Support

aixplain/enums/__init__.py

@@ -16,3 +16,4 @@
 from .sort_order import SortOrder
 from .response_status import ResponseStatus
 from .database_source import DatabaseSourceType
+from .index_stores import IndexStores

aixplain/enums/index_stores.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class IndexStores(Enum):
+    AIR = {"name": "air", "id": "66eae6656eb56311f2595011"}
+    VECTARA = {"name": "vectara", "id": "655e20f46eb563062a1aa301"}
+    # GRAPHRAG = {"name": "graphrag", "id": ""}
+    # ZERO_ENTROPY = {"name": "zero_entropy", "id": ""}
+
+    def __str__(self):
+        return self.value["name"]
+
+    def get_model_id(self):
+        return self.value["id"]

aixplain/factories/index_factory.py

@@ -1,14 +1,17 @@
 from aixplain.modules.model.index_model import IndexModel
 from aixplain.factories import ModelFactory
-from aixplain.enums import Function, ResponseStatus, SortBy, SortOrder, OwnershipType, Supplier
+from aixplain.enums import Function, ResponseStatus, SortBy, SortOrder, OwnershipType, Supplier, IndexStores
 from typing import Optional, Text, Union, List, Tuple
 
 
 class IndexFactory(ModelFactory):
     @classmethod
-    def create(cls, name: Text, description: Text) -> IndexModel:
+    def create(cls, name: Text, description: Text, host: IndexStores = IndexStores.AIR) -> IndexModel:
         """Create a new index collection"""
-        model = cls.get("66eae6656eb56311f2595011")
+        if isinstance(host, str):
+            host = IndexStores[host.upper()]
+        model_id = host.get_model_id()
+        model = cls.get(model_id)
 
         data = {"data": name, "description": description}
         response = model.run(data=data)