[ENH] Add Baseten integration

[philipkiely-baseten]

Description of changes

This PR adds a Baseten integration for using Chroma with models deployed on Baseten

Test plan

Tested locally with deployed model, output is identical to OpenAI integration.

Documentation Changes

Integration is documented in appropriate file.

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[jairad26]

thank you for the PR @philipkiely-baseten. though i don't believe this is complete. It has to implement call to allow users to invoke it. And if you want to allow js client users to use this embedding function, there also needs to be an implementation in the js client, as well as schema, and a couple other functions. Happy to iterate on this, but until this is complete i would recommend moving it to draft. This is not ready to be merged yet.

[philipkiely-baseten]

@jairad26 this works as implemented as it inherits the OpenAI integration's __call__ and other functions. You can technically achieve the same thing just using the OpenAI client with a Baseten api_base, this is a simple wrapper to make explicit that you're using Baseten functionality.

Is a JS client implementation required for all integrations?

[jairad26]

ah missed that it implements OpenAI embedding function instead of the EmbeddingFunction, thanks for pointing that out. The JS client is not required. Will run the workflows

[jairad26]

Could you update the PR title, and the default env var to match other embedding functions? The format should be CHROMA_BASETEN_ENV_VAR

[philipkiely-baseten]

@jairad26 thank you! We're hoping to go live tomorrow with a joint GTM announcement of this integration, if there are any issues with the workflows I will update right away

[philipkiely-baseten]

Could you update the PR title, and the default env var to match other embedding functions? The format should be CHROMA_BASETEN_ENV_VAR

On it!

[philipkiely-baseten]

@jairad26 done on both counts!

[philipkiely-baseten]

@jairad26 do we need to add a CHROMA_BASETEN_API_KEY somewhere for these tests to pass?

[jairad26]

yea just looked at the tests, youll need to add BasetenEmbeddingFunction to the list of expected_builtins in chromadb/test/ef/test_ef.py to pass that test

[philipkiely-baseten]

yea just looked at the tests, youll need to add BasetenEmbeddingFunction to the list of expected_builtins in chromadb/test/ef/test_ef.py to pass that test

Thank you, added

[philipkiely-baseten]

@jairad26 hoping to go live in a couple hours -- everything is approved on the content side w/ jeff just need to get this merged to unblock

[jairad26]

hey @philipkiely-baseten sorry was not aware of the launch, running the workflows. will prioritize this

[philipkiely-baseten]

@jairad26 looking at the failing tests its something with a zip file -- is this related to my change?

[jairad26]

not related to you this was fixed yday. could you rebase off of main

[philipkiely-baseten]

@jairad26 done -- used the sync fork feature on GitHub, lmk if it looks correct

[jairad26]

the bench flakes we can rerun. also to double check, you took a pass and tested the embedding function correct? ill also take a pass once these go through to double check

[philipkiely-baseten]

@jairad26 yes, I tested the workflow that is documented in baseten.md and it works!

[philipkiely-baseten]

@jairad26 here it is in my notebook using a model deployed in my baseten acct

[jairad26]

lovely, will also set up and test this as well just for sanity check. Excited for your launch!

[philipkiely-baseten]

LMK if you have any issues getting a model deployed on Baseten or getting the model ID

[jairad26]

looks like theres some errors on a couple tests.. looking

[jairad26]

theres a couple of changes necessary. i misspoke earlier when i said the other functions are not required. heres the full file, can you replace your baseten_embedding_function with this

import os
from chromadb.utils.embedding_functions.openai_embedding_function import OpenAIEmbeddingFunction
from chromadb.api.types import Documents, EmbeddingFunction
from typing import Dict, Any, Optional

class BasetenEmbeddingFunction(OpenAIEmbeddingFunction):

    def __init__(
            self,
            api_key: Optional[str],
            api_base: str,
            api_key_env_var: str = "CHROMA_BASETEN_API_KEY",
            ):
        """
        Initialize the BasetenEmbeddingFunction.
        Args:
            api_key (str, optional): The API key for your Baseten account
            api_base (str, required): The Baseten URL of the deployment
            api_key_env_var (str, optional): The environment variable to use for the API key. Defaults to "CHROMA_BASETEN_API_KEY".
        """
        try:
            import openai
        except ImportError:
            raise ValueError(
                "The openai python package is not installed. Please install it with `pip install openai`"
            )

        self.api_key_env_var = api_key_env_var
        # Prioritize api_key argument, then environment variable
        resolved_api_key = api_key or os.getenv(api_key_env_var)
        if not resolved_api_key:
            raise ValueError(f"API key not provided and {api_key_env_var} environment variable is not set.")
        self.api_key = resolved_api_key
        if not api_base:
            raise ValueError("The api_base argument must be provided.")
        self.api_base = api_base
        self.model_name = "baseten-embedding-model"
        self.dimensions = None

        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url=self.api_base
        )
        
    @staticmethod
    def name() -> str:
        return "baseten"
    
    def get_config(self) -> Dict[str, Any]:
        return {
            "api_base": self.api_base,
            "api_key_env_var": self.api_key_env_var
        }
        

    @staticmethod
    def build_from_config(config: Dict[str, Any]) -> "BasetenEmbeddingFunction":
        """
        Build the BasetenEmbeddingFunction from a configuration dictionary.

        Args:
            config (Dict[str, Any]): A dictionary containing the configuration parameters.
                                     Expected keys: 'api_key', 'api_base', 'api_key_env_var'.

        Returns:
            BasetenEmbeddingFunction: An instance of BasetenEmbeddingFunction.
        """
        api_key_env_var = config.get("api_key_env_var")
        api_base = config.get("api_base")
        if api_key_env_var is None or api_base is None:
            raise ValueError("Missing 'api_key_env_var' or 'api_base' in configuration for BasetenEmbeddingFunction.")

        # Note: We rely on the __init__ method to handle potential missing api_key
        # by checking the environment variable if the config value is None.
        # However, api_base must be present either in config or have a default.
        if api_base is None:
             raise ValueError("Missing 'api_base' in configuration for BasetenEmbeddingFunction.")

        return BasetenEmbeddingFunction(
            api_key=None, # Pass None if not in config, __init__ will check env var
            api_base=api_base,
            api_key_env_var=api_key_env_var,
        )

[jairad26]

and could you rerun the embedding function verify it works, i can get to it after lunch

[philipkiely-baseten]

@jairad26 sorry, just messed up the history. I'm sitting on a flight about to take off. Will ping when fixed.

[philipkiely-baseten]

@jairad26 Sorry for the weirdness on Git -- never using GitHub's "sync fork" feature again.

Anyway, the PR should be in a good state now, using your version of the file with the extra functions! Thank you for that!

[jairad26]

you also need to add it to known_embedding functions

[philipkiely-baseten]

you also need to add it to known_embedding functions

In what file? Its in init

[philipkiely-baseten]

Also, rebuilt the package locally and verified in notebook, worked as expected

[jairad26]

sorry missed that file, yep looks good! thanks for all the iteration

[philipkiely-baseten]

@jairad26 thank you for your help! Fingers crossed there are no issues on the checks