Add Amazon Bedrock Embedding function #1361

chezou · 2023-11-09T00:05:42Z

https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html

Description of changes

New functionality
- Support Amazon Bedrock embedding function

Test plan

Tests pass locally with pytest for python, yarn test for js

Tested locally by given profile_name with appropreate ~/.aws/config

>>> import boto3
>>> from chromadb.utils.embedding_functions import AmazonBedrockEmbeddingFunction
>>> session = boto3.Session(profile_name="myprofile", region_name="us-east-1")
>>> ef = AmazonBedrockEmbeddingFunction(session=session)
>>> ef(["Hello Bedrock"])
[[-0.73046875, 0.390625, 0.24511719, 0.111816406, 0.83203125, 0.79296875,...,]]

Documentation Changes

Written docstrings as much as possible.

https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html

github-actions · 2023-11-09T00:05:57Z

beggers

This is sick, thanks for expanding Chroma's functionality. I have a few notes but once we get them sorted I will be very happy to merge this in.

beggers · 2023-12-18T23:17:51Z

chromadb/utils/embedding_functions.py

+        Returns:
+            Embeddings: The embeddings for the texts.
+
+        Example:


Could we move this example into the docstring for the class instead of its __call__ method?

Fixed 0b4bfd0

beggers · 2023-12-18T23:19:00Z

chromadb/utils/embedding_functions.py

+            embeddings.append(self._invoke(text))
+        return embeddings
+
+    def _invoke(


I don't think we need __call__() to be a thin wrapper around _invoke(). Would you mind moving this logic into __call__ directly?

Fixed 0b4bfd0

beggers · 2023-12-18T23:20:29Z

chromadb/utils/embedding_functions.py

+                contentType=content_type,
+            )
+            embedding = json.load(response.get("body")).get("embedding")
+        except Exception as e:


Could we remove this and the re-throw? I want clients to be able to handle exceptions differently depending on the type of exception, e.g. trying again in the case of a timeout.

Fixed 0b4bfd0

beggers · 2023-12-18T23:30:49Z

chromadb/utils/embedding_functions.py

+        self,
+        profile_name: Optional[str] = None,
+        region: Optional[str] = None,
+        model_name: str = "amazon.titan-embed-text-v1",


I think there's a cleaner API boundary here and I'd like to hear your thoughts.

This constructor could accept a boto3.Session and optional kwargs to pass into the client() call. That way we give users full access to the configuration options for both their Session and their client, and we have less code to maintain (we don't need to read from os.environ for example). I may be missing something here -- is there a reason we need to accept these and construct the Session ourselves? If not could we accept a Session and config kwargs?

I'm okay with just passing a boto3.Session and necessary kwargs for the client() call instead of creating a session within the class.

I don't have a specific reason to create a boto3.Session inside the class. I tried to align with the existing codes. Looking at the other classes creates a client internally, and HuggingFaceEmbeddingFunction creates requests.Session within the class.

btw, do you think we should support all the possible config kwargs as much as possible?

Looking at the boto3 document, I think most of the kwargs are covered by the session, so we may need to pass api_version, use_ssl, verify, and endpoint_url, at least. Or, we may postpone introducing those arguments at this moment.

Tentatively, I removed most of the redundant arguments without adding client-specific arguments. Let me know your thoughts on it.
eedf49d

This is already looking way better -- thank you!

One more small ask and then we can check it in: could we have the constructor accept **kwargs and pass them directly to the client() we create? That would give us three wins:
1 - We don't have to manage the retry_config which I imagine some folks will want control over.
2 - Users can specify whatever they want to override in their Session config without us needing to make any code changes.
3 - We don't need to import boto3 in this codebase anymore.

Fair point. Fixed by using **kwargs a53e415

- Remove unnecessary function - Remove re-throw an exception - Move comment to the right under the class

beggers

🚀 🚀 🚀

tazarov · 2023-12-21T11:19:10Z

@chezou, it would be really nice if you can raise a PR against the docs repo - https://github.com/chroma-core/docs/pulls to update the embedding functions and how to use them - https://docs.trychroma.com/embeddings

tazarov · 2023-12-21T11:21:55Z

chromadb/utils/embedding_functions.py

+            >>> texts = ["Hello, world!", "How are you?"]
+            >>> embeddings = bedrock(texts)
+        """
+


@chezou, it would be nice to check if boto3 package is installed and if not guide to user with appropriate error message. Check other EFs as ref.

I originally did checking it, but I followed @beggers suggestion to remove the dependency of boto3 from chroma itself.
See also #1361 (comment)

tazarov · 2023-12-21T11:22:31Z

chromadb/utils/embedding_functions.py

+        )
+
+    def __call__(self, input: Documents) -> Embeddings:
+        import json


Do we need to import json here instead of at top-level?

Good call, #1562

tazarov · 2023-12-21T11:22:54Z

@chezou, on second thought it'd be absolutely awesome if you can also add the JS client equivalent of the EF.

chezou · 2024-01-20T02:01:41Z

@tazarov I made a PR for the docs repo. Please have a look when your team has a chance chroma-core/docs#204

For JS client, I'm not familiar with the JS AWS client, so it may take a while.

[edit] I made a PR for JS client #1659

Add Amazon Bedrock Embedding function

49a33b3

https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html

beggers reviewed Dec 18, 2023

View reviewed changes

beggers assigned chezou Dec 18, 2023

chezou added 2 commits December 18, 2023 16:04

Refactor based on review comments

0b4bfd0

- Remove unnecessary function - Remove re-throw an exception - Move comment to the right under the class

Pass boto3 session instead of building it in the function

eedf49d

chezou removed their assignment Dec 19, 2023

chezou and others added 2 commits December 19, 2023 11:54

Delegate boto3 client configuration to the user by accepting kwargs

a53e415

Merge branch 'main' into bedrock

cc77bb7

beggers approved these changes Dec 20, 2023

View reviewed changes

beggers merged commit 2202df8 into chroma-core:main Dec 20, 2023

chezou deleted the bedrock branch December 21, 2023 01:36

tazarov reviewed Dec 21, 2023

View reviewed changes

chezou mentioned this pull request Jan 20, 2024

feat: Amazon Bedrock Embedding doc chroma-core/docs#204

Open

chezou mentioned this pull request Jan 20, 2024

[ENH] AmazonBedrockEmbeddingFunction for JS #1659

Closed

1 task

This was referenced Jun 18, 2024

Update dependency chromadb to v0.6.3 - autoclosed behroozazarkhalili/Chatbot-DL#5

Closed

Update dependency chromadb to v0.6.3 sky-dust-intelligence/langchain-search#60

Open

Add Amazon Bedrock Embedding function #1361

Add Amazon Bedrock Embedding function #1361

Uh oh!

Conversation

chezou commented Nov 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

Test plan

Documentation Changes

Uh oh!

github-actions bot commented Nov 9, 2023

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

beggers left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chezou Dec 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chezou Dec 19, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

beggers left a comment

Choose a reason for hiding this comment

Uh oh!

tazarov commented Dec 21, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chezou Dec 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tazarov commented Dec 21, 2023

Uh oh!

chezou commented Jan 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

chezou commented Nov 9, 2023 •

edited

Loading

chezou Dec 18, 2023 •

edited

Loading

chezou Dec 19, 2023 •

edited

Loading

chezou Dec 21, 2023 •

edited

Loading

chezou commented Jan 20, 2024 •

edited

Loading