[Feat] Integrate Bedrock & Vertex AI Embedding Adapters (#156)

* Initial commit

* Added bedrock embedding support

* Fix: Inherit EmbeddingAdapter in Bedrock

* added method to instantiate Bedrock class

* Added vertexAI embedding

* Added schema for vertex ai embeddings

* updated json schema for bedrock embeddings

* updated sdk's version

* Updated SDK's version

* Update src/unstract/sdk/adapters/exceptions.py

Co-authored-by: Chandrasekharan M <117059509+chandrasekharan-zipstack@users.noreply.github.com>
Signed-off-by: Praveen Kumar <praveen@zipstack.com>

* Update __init__.py

Signed-off-by: Praveen Kumar <praveen@zipstack.com>

* Imported EmbeddinError from parent exception class

* moved test_embedding function to parent class

* Moved get_jason_schema to base Adapter class

* added pdm.lock

---------

Signed-off-by: Praveen Kumar <praveen@zipstack.com>
Co-authored-by: Chandrasekharan M <117059509+chandrasekharan-zipstack@users.noreply.github.com>

Author: pk-zipstack

pdm.lock

@@ -24,6 +24,8 @@ dependencies = [
     #        "llama-index-embeddings-fastembed==0.1.4",
     "llama-index-embeddings-openai==0.3.1",
     "llama-index-embeddings-ollama==0.5.0",
+    "llama-index-embeddings-bedrock==0.5.0",
+    "llama-index-embeddings-vertex==0.3.1",
     "llama-index-vector-stores-postgres==0.4.1",
     # Including Supabase conflicts with postgres on pg-vector.
     # Hence, commenting it out at the moment

pyproject.toml

@@ -1,4 +1,4 @@
-__version__ = "0.56.0rc4"
+__version__ = "0.57.0rc1"
 
 
 def get_sdk_version():

src/unstract/sdk/__init__.py

@@ -30,10 +30,13 @@ def get_description() -> str:
     def get_icon() -> str:
         return ""
 
-    @staticmethod
-    @abstractmethod
-    def get_json_schema() -> str:
-        return ""
+    @classmethod
+    def get_json_schema(cls) -> str:
+        schema_path = getattr(cls, 'SCHEMA_PATH', None)
+        if schema_path is None:
+            raise ValueError(f"SCHEMA_PATH not defined for {cls.__name__}")
+        with open(schema_path) as f:
+            return f.read()
 
     @staticmethod
     @abstractmethod

src/unstract/sdk/adapters/base.py

@@ -27,6 +27,8 @@ def __init__(self, settings: dict[str, Any]):
         super().__init__("AzureOpenAIEmbedding")
         self.config = settings
 
+    SCHEMA_PATH = f"{os.path.dirname(__file__)}/static/json_schema.json"
+
     @staticmethod
     def get_id() -> str:
         return "azureopenai|9770f3f6-f8ba-4fa0-bb3a-bef48a00e66f"
@@ -47,13 +49,6 @@ def get_provider() -> str:
     def get_icon() -> str:
         return "/icons/adapter-icons/AzureopenAI.png"
 
-    @staticmethod
-    def get_json_schema() -> str:
-        f = open(f"{os.path.dirname(__file__)}/static/json_schema.json")
-        schema = f.read()
-        f.close()
-        return schema
-
     def get_embedding_instance(self) -> BaseEmbedding:
         try:
             embedding_batch_size = EmbeddingHelper.get_embedding_batch_size(
@@ -77,7 +72,4 @@ def get_embedding_instance(self) -> BaseEmbedding:
         except Exception as e:
             raise AdapterError(str(e))
 
-    def test_connection(self) -> bool:
-        embedding = self.get_embedding_instance()
-        test_result: bool = EmbeddingHelper.test_embedding_instance(embedding)
-        return test_result
+     

src/unstract/sdk/adapters/embedding/azure_open_ai/src/azure_open_ai.py

@@ -0,0 +1 @@
+# Unstract Bedrock Embeddings

src/unstract/sdk/adapters/embedding/bedrock/README.md

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["pdm-backend"]
+build-backend = "pdm.backend"
+
+
+[project]
+name = "unstract-bedrock-embedding"
+version = "0.0.1"
+description = "Bedrock Embedding"
+authors = [
+    {name = "Zipstack Inc.", email = "devsupport@zipstack.com"},
+]
+dependencies = [
+]
+requires-python = ">=3.9"
+readme = "README.md"
+classifiers = [
+  "Programming Language :: Python"
+]
+license = {text = "MIT"}
+
+[tool.pdm.build]
+includes = ["src"]
+package-dir = "src"
+# source-includes = ["tests"]

src/unstract/sdk/adapters/embedding/bedrock/pyproject.toml

@@ -0,0 +1,9 @@
+from .bedrock import Bedrock
+
+metadata = {
+    "name": Bedrock.__name__,
+    "version": "1.0.0",
+    "adapter": Bedrock,
+    "description": "Bedrock embedding adapter",
+    "is_active": True,
+}

src/unstract/sdk/adapters/embedding/bedrock/src/__init__.py

@@ -0,0 +1,70 @@
+import os
+from typing import Any
+
+from llama_index.core.embeddings import BaseEmbedding
+from llama_index.embeddings.bedrock import BedrockEmbedding
+
+from unstract.sdk.adapters.embedding.embedding_adapter import EmbeddingAdapter
+from unstract.sdk.adapters.embedding.helper import EmbeddingHelper
+from unstract.sdk.adapters.exceptions import AdapterError
+
+class Constants:
+    MODEL = "model"
+    TIMEOUT = "timeout"
+    MAX_RETRIES = "max_retries"
+    SECRET_ACCESS_KEY = "aws_secret_access_key"
+    ACCESS_KEY_ID = "aws_access_key_id"
+    REGION_NAME = "region_name"
+    DEFAULT_TIMEOUT = 240
+    DEFAULT_MAX_RETRIES = 3
+
+class Bedrock(EmbeddingAdapter):
+    def __init__(self, settings: dict[str, Any]):
+        super().__init__("Bedrock")
+        self.config = settings
+
+    SCHEMA_PATH = f"{os.path.dirname(__file__)}/static/json_schema.json"
+
+    @staticmethod
+    def get_id() -> str:
+        return "bedrock|88199741-8d7e-4e8c-9d92-d76b0dc20c91"
+
+    @staticmethod
+    def get_name() -> str:
+        return "Bedrock"
+
+    @staticmethod
+    def get_description() -> str:
+        return "Bedrock Embedding"
+
+    @staticmethod
+    def get_provider() -> str:
+        return "bedrock"
+
+    @staticmethod
+    def get_icon() -> str:
+        return "/icons/adapter-icons/Bedrock.png"
+
+    def get_embedding_instance(self) -> BaseEmbedding:
+        try:
+            embedding_batch_size = EmbeddingHelper.get_embedding_batch_size(
+                config=self.config
+            )
+            embedding: BaseEmbedding = BedrockEmbedding(
+                model_name=self.config.get(Constants.MODEL),
+                aws_access_key_id=self.config.get(Constants.ACCESS_KEY_ID),
+                aws_secret_access_key=self.config.get(Constants.SECRET_ACCESS_KEY),
+                region_name=self.config.get(Constants.REGION_NAME),
+                timeout=float(
+                    self.config.get(Constants.TIMEOUT, Constants.DEFAULT_TIMEOUT)
+                ),
+                max_retries=int(
+                    self.config.get(Constants.MAX_RETRIES, Constants.DEFAULT_MAX_RETRIES)
+                ),
+                embed_batch_size=embedding_batch_size,
+            )
+            return embedding
+        except Exception as e:
+            raise AdapterError(str(e))
+        
+     

src/unstract/sdk/adapters/embedding/bedrock/src/bedrock.py

@@ -0,0 +1,66 @@
+{
+    "title": "Bedrock Embeddings",
+    "type": "object",
+    "required": [
+      "aws_secret_access_key",
+      "region_name",
+      "aws_access_key_id",
+      "model",
+      "adapter_name"
+    ],
+    "properties": {
+      "adapter_name": {
+        "type": "string",
+        "title": "Name",
+        "default": "",
+        "description": "Provide a unique name for this adapter instance. Example: Bedrock-Embedding-1"
+      },
+      "model": {
+        "type": "string",
+        "title": "Model",
+        "default": "amazon.titan-embed-text-v2:0",
+        "description": "Model name. Refer to [Bedrock's documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) for the list of available models."
+      },
+      "aws_access_key_id": {
+        "type": "string",
+        "title": "AWS Access Key ID",
+        "description": "Provide your AWS Access Key ID",
+        "format": "password"
+      },
+      "aws_secret_access_key": {
+        "type": "string",
+        "title": "AWS Secret Access Key",
+        "description": "Provide your AWS Secret Access Key",
+        "format": "password"
+      },
+      "region_name": {
+        "type": "string",
+        "title": "AWS Region name",
+        "description": "Provide the AWS Region name where the service is running. Eg. us-east-1"
+      },
+      "embed_batch_size": {
+      "type": "number",
+      "minimum": 0,
+      "multipleOf": 1,
+      "title": "Embedding Batch Size",
+      "default": 10
+      },
+      "max_retries": {
+        "type": "number",
+        "minimum": 0,
+        "multipleOf": 1,
+        "title": "Max Retries",
+        "default": 5,
+        "description": "Maximum number of retries to attempt when a request fails."
+      },
+      "timeout": {
+        "type": "number",
+        "minimum": 0,
+        "multipleOf": 1,
+        "title": "Timeout",
+        "default": 900,
+        "description": "Timeout in seconds"
+      }
+    }
+  }
+  

src/unstract/sdk/adapters/embedding/bedrock/src/static/json_schema.json

@@ -7,6 +7,7 @@
 from unstract.sdk.adapters.base import Adapter
 from unstract.sdk.adapters.enums import AdapterTypes
 
+from unstract.sdk.adapters.embedding.helper import EmbeddingHelper
 
 class EmbeddingAdapter(Adapter, ABC):
     def __init__(self, name: str):
@@ -34,10 +35,6 @@ def get_provider() -> str:
     def get_icon() -> str:
         return ""
 
-    @staticmethod
-    def get_json_schema() -> str:
-        return ""
-
     @staticmethod
     def get_adapter_type() -> AdapterTypes:
         return AdapterTypes.EMBEDDING
@@ -50,3 +47,8 @@ def get_embedding_instance(self, embed_config: dict[str, Any]) -> BaseEmbedding:
             Raises exceptions for any error
         """
         return MockEmbedding(embed_dim=1)
+    
+    def test_connection(self) -> bool:
+        embedding = self.get_embedding_instance()
+        test_result: bool = EmbeddingHelper.test_embedding_instance(embedding)
+        return test_result

src/unstract/sdk/adapters/embedding/embedding_adapter.py

@@ -22,6 +22,8 @@ def __init__(self, settings: dict[str, Any]):
         super().__init__("HuggingFace")
         self.config = settings
 
+    SCHEMA_PATH = f"{os.path.dirname(__file__)}/static/json_schema.json"
+
     @staticmethod
     def get_id() -> str:
         return "huggingface|90ec9ec2-1768-4d69-8fb1-c88b95de5e5a"
@@ -38,13 +40,6 @@ def get_description() -> str:
     def get_icon() -> str:
         return "/icons/adapter-icons/huggingface.png"
 
-    @staticmethod
-    def get_json_schema() -> str:
-        f = open(f"{os.path.dirname(__file__)}/static/json_schema.json")
-        schema = f.read()
-        f.close()
-        return schema
-
     def get_embedding_instance(self) -> BaseEmbedding:
         try:
             embedding_batch_size = EmbeddingHelper.get_embedding_batch_size(
@@ -67,7 +62,4 @@ def get_embedding_instance(self) -> BaseEmbedding:
         except Exception as e:
             raise AdapterError(str(e))
 
-    def test_connection(self) -> bool:
-        embedding = self.get_embedding_instance()
-        test_result: bool = EmbeddingHelper.test_embedding_instance(embedding)
-        return test_result
+     

src/unstract/sdk/adapters/embedding/hugging_face/src/hugging_face.py

@@ -15,6 +15,8 @@ def __init__(self, settings: dict[str, Any]):
         super().__init__("NoOpCustomEmbedding")
         self.config = settings
 
+    SCHEMA_PATH = f"{os.path.dirname(__file__)}/static/json_schema.json"
+
     @staticmethod
     def get_id() -> str:
         return "noOpEmbedding|ff223003-fee8-4079-b288-e86215e6b39a"
@@ -35,13 +37,6 @@ def get_icon() -> str:
     def get_provider() -> str:
         return "NoOp"
 
-    @staticmethod
-    def get_json_schema() -> str:
-        f = open(f"{os.path.dirname(__file__)}/static/json_schema.json")
-        schema = f.read()
-        f.close()
-        return schema
-
     def get_embedding_instance(self) -> BaseEmbedding:
         embedding: BaseEmbedding = NoOpCustomEmbedding(
             embed_dim=1, wait_time=self.config.get("wait_time")

src/unstract/sdk/adapters/embedding/no_op/src/no_op_embedding.py

@@ -20,6 +20,8 @@ def __init__(self, settings: dict[str, Any]):
         super().__init__("Ollama")
         self.config = settings
 
+    SCHEMA_PATH = f"{os.path.dirname(__file__)}/static/json_schema.json"
+
     @staticmethod
     def get_id() -> str:
         return "ollama|d58d7080-55a9-4542-becd-8433528e127b"
@@ -40,13 +42,6 @@ def get_provider() -> str:
     def get_icon() -> str:
         return "/icons/adapter-icons/ollama.png"
 
-    @staticmethod
-    def get_json_schema() -> str:
-        f = open(f"{os.path.dirname(__file__)}/static/json_schema.json")
-        schema = f.read()
-        f.close()
-        return schema
-
     def get_embedding_instance(self) -> BaseEmbedding:
         try:
             embedding_batch_size = EmbeddingHelper.get_embedding_batch_size(
@@ -61,7 +56,4 @@ def get_embedding_instance(self) -> BaseEmbedding:
         except Exception as e:
             raise AdapterError(str(e))
 
-    def test_connection(self) -> bool:
-        embedding = self.get_embedding_instance()
-        test_result: bool = EmbeddingHelper.test_embedding_instance(embedding)
-        return test_result
+     

src/unstract/sdk/adapters/embedding/ollama/src/ollama.py

@@ -27,6 +27,8 @@ def __init__(self, settings: dict[str, Any]):
         super().__init__("OpenAI")
         self.config = settings
 
+    SCHEMA_PATH = f"{os.path.dirname(__file__)}/static/json_schema.json"
+
     @staticmethod
     def get_id() -> str:
         return "openai|717a0b0e-3bbc-41dc-9f0c-5689437a1151"
@@ -47,13 +49,6 @@ def get_provider() -> str:
     def get_icon() -> str:
         return "/icons/adapter-icons/OpenAI.png"
 
-    @staticmethod
-    def get_json_schema() -> str:
-        f = open(f"{os.path.dirname(__file__)}/static/json_schema.json")
-        schema = f.read()
-        f.close()
-        return schema
-
     def get_embedding_instance(self) -> BaseEmbedding:
         try:
             timeout = int(self.config.get(Constants.TIMEOUT, Constants.DEFAULT_TIMEOUT))
@@ -73,7 +68,4 @@ def get_embedding_instance(self) -> BaseEmbedding:
         except Exception as e:
             raise AdapterError(str(e))
 
-    def test_connection(self) -> bool:
-        embedding = self.get_embedding_instance()
-        test_result: bool = EmbeddingHelper.test_embedding_instance(embedding)
-        return test_result
+