lancedb/python/python/lancedb/embeddings/bedrock.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors


import json
from functools import cached_property
from typing import List, Union

import numpy as np

from lancedb.pydantic import PYDANTIC_VERSION

from ..util import attempt_import_or_raise
from .base import TextEmbeddingFunction
from .registry import register
from .utils import TEXT


@register("bedrock-text")
class BedRockText(TextEmbeddingFunction):
    """
    Parameters
    ----------
    name : str, default "amazon.titan-embed-text-v1"
        The model ID of the bedrock model to use. Supported models for are:
        - amazon.titan-embed-text-v1
        - cohere.embed-english-v3
        - cohere.embed-multilingual-v3
    region : str, default "us-east-1"
        Optional name of the AWS Region in which the service should be called.
    profile_name : str, default None
        Optional name of the AWS profile to use for calling the Bedrock service.
        If not specified, the default profile will be used.
    assumed_role : str, default None
        Optional ARN of an AWS IAM role to assume for calling the Bedrock service.
        If not specified, the current active credentials will be used.
    role_session_name : str, default "lancedb-embeddings"
        Optional name of the AWS IAM role session to use for calling the Bedrock
        service. If not specified, "lancedb-embeddings" name will be used.

    Examples
    --------
    import lancedb
    import pandas as pd
    from lancedb.pydantic import LanceModel, Vector

    model = get_registry().get("bedrock-text").create()

    class TextModel(LanceModel):
        text: str = model.SourceField()
        vector: Vector(model.ndims()) = model.VectorField()

    df = pd.DataFrame({"text": ["hello world", "goodbye world"]})
    db = lancedb.connect("tmp_path")
    tbl = db.create_table("test", schema=TextModel, mode="overwrite")

    tbl.add(df)

    rs = tbl.search("hello").limit(1).to_pandas()
    """

    name: str = "amazon.titan-embed-text-v1"
    region: str = "us-east-1"
    assumed_role: Union[str, None] = None
    profile_name: Union[str, None] = None
    role_session_name: str = "lancedb-embeddings"
    source_input_type: str = "search_document"
    query_input_type: str = "search_query"

    if PYDANTIC_VERSION.major < 2:  # Pydantic 1.x compat

        class Config:
            keep_untouched = (cached_property,)
    else:
        model_config = dict()
        model_config["ignored_types"] = (cached_property,)

    def ndims(self):
        # return len(self._generate_embedding("test"))
        # TODO: fix hardcoding
        if self.name == "amazon.titan-embed-text-v1":
            return 1536
        elif self.name in [
            "amazon.titan-embed-text-v2:0",
            "cohere.embed-english-v3",
            "cohere.embed-multilingual-v3",
        ]:
            # TODO: "amazon.titan-embed-text-v2:0" model supports dynamic ndims
            return 1024
        else:
            raise ValueError(f"Model {self.name} not supported")

    def compute_query_embeddings(
        self, query: str, *args, **kwargs
    ) -> List[List[float]]:
        return self.compute_source_embeddings(query, input_type=self.query_input_type)

    def compute_source_embeddings(
        self, texts: TEXT, *args, **kwargs
    ) -> List[List[float]]:
        texts = self.sanitize_input(texts)
        # assume source input type if not passed by `compute_query_embeddings`
        kwargs["input_type"] = kwargs.get("input_type") or self.source_input_type

        return self.generate_embeddings(texts, **kwargs)

    def generate_embeddings(
        self, texts: Union[List[str], np.ndarray], *args, **kwargs
    ) -> List[List[float]]:
        """
        Get the embeddings for the given texts

        Parameters
        ----------
        texts: list[str] or np.ndarray (of str)
            The texts to embed

        Returns
        -------
        list[list[float]]
            The embeddings for the given texts
        """
        results = []
        for text in texts:
            response = self._generate_embedding(text, *args, **kwargs)
            results.append(response)
        return results

    def _generate_embedding(self, text: str, *args, **kwargs) -> List[float]:
        """
        Get the embeddings for the given texts

        Parameters
        ----------
        texts: str
            The texts to embed

        Returns
        -------
        list[float]
            The embeddings for the given texts
        """
        # format input body for provider
        provider = self.name.split(".")[0]
        input_body = {**kwargs}
        if provider == "cohere":
            input_body["texts"] = [text]
        else:
            # includes common provider == "amazon"
            input_body.pop("input_type", None)
            input_body["inputText"] = text
        body = json.dumps(input_body)

        try:
            # invoke bedrock API
            response = self.client.invoke_model(
                body=body,
                modelId=self.name,
                accept="application/json",
                contentType="application/json",
            )

            # format output based on provider
            response_body = json.loads(response.get("body").read())
            if provider == "cohere":
                return response_body.get("embeddings")[0]
            else:
                # includes common provider == "amazon"
                return response_body.get("embedding")
        except Exception as e:
            help_txt = """
                boto3 client failed to invoke the bedrock API. In case of
                AWS credentials error:
                    - Please check your AWS credentials and ensure that you have access.
                    You can set up aws credentials using `aws configure` command and
                    verify by running `aws sts get-caller-identity` in your terminal.
                """
            raise ValueError(f"Error raised by boto3 client: {e}. \n {help_txt}")

    @cached_property
    def client(self):
        """Create a boto3 client for Amazon Bedrock service

        Returns
        -------
        boto3.client
            The boto3 client for Amazon Bedrock service
        """
        botocore = attempt_import_or_raise("botocore")
        boto3 = attempt_import_or_raise("boto3")

        session_kwargs = {"region_name": self.region}
        client_kwargs = {**session_kwargs}

        if self.profile_name:
            session_kwargs["profile_name"] = self.profile_name

        retry_config = botocore.config.Config(
            region_name=self.region,
            retries={
                "max_attempts": 0,  # disable this as retries retries are handled
                "mode": "standard",
            },
        )
        session = (
            boto3.Session(**session_kwargs) if self.profile_name else boto3.Session()
        )
        if self.assumed_role:  # if not using default credentials
            sts = session.client("sts")
            response = sts.assume_role(
                RoleArn=str(self.assumed_role),
                RoleSessionName=self.role_session_name,
            )
            client_kwargs["aws_access_key_id"] = response["Credentials"]["AccessKeyId"]
            client_kwargs["aws_secret_access_key"] = response["Credentials"][
                "SecretAccessKey"
            ]
            client_kwargs["aws_session_token"] = response["Credentials"]["SessionToken"]

        service_name = "bedrock-runtime"

        bedrock_client = session.client(
            service_name=service_name, config=retry_config, **client_kwargs
        )

        return bedrock_client