From f48e3ce0d5cbd7b5a8e146dc3983aeab020c5f70 Mon Sep 17 00:00:00 2001 From: Chong Luo Date: Wed, 25 Sep 2024 21:35:18 +0800 Subject: [PATCH] Add Voyage AI embedding API for Anthropic. --- .gitignore | 3 +- README.md | 1 + gptcache/embedding/__init__.py | 6 +- gptcache/embedding/voyageai.py | 88 ++++++++++++ gptcache/utils/__init__.py | 4 + tests/unit_tests/embedding/test_voyageai.py | 140 ++++++++++++++++++++ 6 files changed, 240 insertions(+), 2 deletions(-) create mode 100644 gptcache/embedding/voyageai.py create mode 100644 tests/unit_tests/embedding/test_voyageai.py diff --git a/.gitignore b/.gitignore index a972ff9a..38b4aace 100644 --- a/.gitignore +++ b/.gitignore @@ -138,4 +138,5 @@ dmypy.json **/example.db **/.chroma docs/references/* -!docs/references/index.rst \ No newline at end of file +!docs/references/index.rst +.vscode/ \ No newline at end of file diff --git a/README.md b/README.md index c5f6f955..e9988dba 100644 --- a/README.md +++ b/README.md @@ -331,6 +331,7 @@ This module is created to extract embeddings from requests for similarity search - [x] Support [fastText](https://fasttext.cc) embedding. - [x] Support [SentenceTransformers](https://www.sbert.net) embedding. - [x] Support [Timm](https://timm.fast.ai/) models for image embedding. + - [x] Support [VoyageAI](https://www.voyageai.com/) embedding API for Anthropic. - [ ] Support other embedding APIs. - **Cache Storage**: **Cache Storage** is where the response from LLMs, such as ChatGPT, is stored. Cached responses are retrieved to assist in evaluating similarity and are returned to the requester if there is a good semantic match. At present, GPTCache supports SQLite and offers a universally accessible interface for extension of this module. diff --git a/gptcache/embedding/__init__.py b/gptcache/embedding/__init__.py index e813021f..a78a0124 100644 --- a/gptcache/embedding/__init__.py +++ b/gptcache/embedding/__init__.py @@ -12,6 +12,7 @@ "Rwkv", "PaddleNLP", "UForm", + "VoyageAI", ] @@ -31,7 +32,7 @@ paddlenlp = LazyImport("paddlenlp", globals(), "gptcache.embedding.paddlenlp") uform = LazyImport("uform", globals(), "gptcache.embedding.uform") nomic = LazyImport("nomic", globals(), "gptcache.embedding.nomic") - +voyageai = LazyImport("voyageai", globals(), "gptcache.embedding.voyageai") def Nomic(model: str = "nomic-embed-text-v1.5", api_key: str = None, @@ -90,3 +91,6 @@ def PaddleNLP(model="ernie-3.0-medium-zh"): def UForm(model="unum-cloud/uform-vl-multilingual", embedding_type="text"): return uform.UForm(model, embedding_type) + +def VoyageAI(model: str="voyage-3", api_key: str=None, api_key_path:str=None, input_type:str=None, truncation:bool=True): + return voyageai.VoyageAI(model=model, api_key=api_key, api_key_path=api_key_path, input_type=input_type, truncation=truncation) diff --git a/gptcache/embedding/voyageai.py b/gptcache/embedding/voyageai.py new file mode 100644 index 00000000..4d88d65e --- /dev/null +++ b/gptcache/embedding/voyageai.py @@ -0,0 +1,88 @@ +import numpy as np + +from gptcache.utils import import_voyageai +from gptcache.embedding.base import BaseEmbedding + +import_voyageai() + +import voyageai + + +class VoyageAI(BaseEmbedding): + """Generate text embedding for given text using VoyageAI. + + :param model: The model name to use for generating embeddings. Defaults to 'voyage-3'. + :type model: str + :param api_key_path: The path to the VoyageAI API key file. + :type api_key_path: str + :param api_key: The VoyageAI API key. If it is None, the client will search for the API key in the following order: + 1. api_key_path, path to the file containing the key; + 2. environment variable VOYAGE_API_KEY_PATH, which can be set to the path to the file containing the key; + 3. environment variable VOYAGE_API_KEY. + This behavior is defined by the VoyageAI Python SDK. + :type api_key: str + :param input_type: The type of input data. Defaults to None. Default to None. Other options: query, document. + More details can be found in the https://docs.voyageai.com/docs/embeddings + :type input_type: str + :param truncation: Whether to truncate the input data. Defaults to True. + :type truncation: bool + + Example: + .. code-block:: python + + from gptcache.embedding import VoyageAI + + test_sentence = 'Hello, world.' + encoder = VoyageAI(model='voyage-3', api_key='your_voyageai_key') + embed = encoder.to_embeddings(test_sentence) + """ + + def __init__(self, model: str = "voyage-3", api_key_path: str = None, api_key: str = None, input_type: str = None, truncation: bool = True): + voyageai.api_key_path = api_key_path + voyageai.api_key = api_key + + self._vo = voyageai.Client() + self._model = model + self._input_type = input_type + self._truncation = truncation + + if self._model in self.dim_dict(): + self.__dimension = self.dim_dict()[model] + else: + self.__dimension = None + + def to_embeddings(self, data, **_): + """ + Generate embedding for the given text input. + + :param data: The input text. + :type data: str or list[str] + + :return: The text embedding in the shape of (dim,). + :rtype: numpy.ndarray + """ + if not isinstance(data, list): + data = [data] + result = self._vo.embed(texts=data, model=self._model, input_type=self._input_type, truncation=self._truncation) + embeddings = result.embeddings + return np.array(embeddings).astype("float32").squeeze(0) + + @property + def dimension(self): + """Embedding dimension. + + :return: embedding dimension + """ + if not self.__dimension: + foo_emb = self.to_embeddings("foo") + self.__dimension = len(foo_emb) + return self.__dimension + + @staticmethod + def dim_dict(): + return {"voyage-3": 1024, + "voyage-3-lite": 512, + "voyage-finance-2": 1024, + "voyage-multilingual-2": 1024, + "voyage-law-2": 1024, + "voyage-code-2": 1536} diff --git a/gptcache/utils/__init__.py b/gptcache/utils/__init__.py index 877e23f6..9f92185b 100644 --- a/gptcache/utils/__init__.py +++ b/gptcache/utils/__init__.py @@ -4,6 +4,7 @@ "import_sbert", "import_cohere", "import_nomic", + "import_voyageai", "import_fasttext", "import_huggingface", "import_uform", @@ -85,6 +86,9 @@ def import_cohere(): def import_nomic(): _check_library("nomic") +def import_voyageai(): + _check_library("voyageai") + def import_fasttext(): _check_library("fasttext", package="fasttext==0.9.2") diff --git a/tests/unit_tests/embedding/test_voyageai.py b/tests/unit_tests/embedding/test_voyageai.py new file mode 100644 index 00000000..a0e4da72 --- /dev/null +++ b/tests/unit_tests/embedding/test_voyageai.py @@ -0,0 +1,140 @@ +import os +import types +import pytest +import mock +from gptcache.utils import import_voyageai +from gptcache.embedding import VoyageAI + +import_voyageai() + + + +@mock.patch.dict(os.environ, {"VOYAGE_API_KEY": "API_KEY", "VOYAGE_API_KEY_PATH": "API_KEY_FILE_PATH_ENV"}) +@mock.patch("builtins.open", new_callable=mock.mock_open, read_data="API_KEY") +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_without_api_key(mock_created, mock_file): + dimension = 1024 + vo = VoyageAI() + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + + mock_file.assert_called_once_with("API_KEY_FILE_PATH_ENV", "rt") + mock_created.assert_called_once_with(texts=["foo"], model="voyage-3", input_type=None, truncation=True) + + +@mock.patch.dict(os.environ, {"VOYAGE_API_KEY": "API_KEY", "VOYAGE_API_KEY_PATH": "API_KEY_FILE_PATH_ENV"}) +@mock.patch("builtins.open", new_callable=mock.mock_open, read_data="API_KEY") +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_with_api_key_path(mock_create, mock_file): + dimension = 1024 + vo = VoyageAI(api_key_path="API_KEY_FILE_PATH") + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + + mock_file.assert_called_once_with("API_KEY_FILE_PATH", "rt") + mock_create.assert_called_once_with(texts=["foo"], model="voyage-3", input_type=None, truncation=True) + + +@mock.patch.dict(os.environ, {"VOYAGE_API_KEY": "API_KEY"}) +@mock.patch("builtins.open", new_callable=mock.mock_open, read_data="API_KEY") +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_with_api_key_in_envrion(mock_create, mock_file): + dimension = 1024 + vo = VoyageAI() + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_file.assert_not_called() + mock_create.assert_called_once_with(texts=["foo"], model="voyage-3", input_type=None, truncation=True) + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_with_api_key(mock_create): + dimension = 1024 + vo = VoyageAI(api_key="API_KEY") + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_create.assert_called_once_with(texts=["foo"], model="voyage-3", input_type=None, truncation=True) + + +@mock.patch.dict(os.environ) +@mock.patch("builtins.open", new_callable=mock.mock_open, read_data="API_KEY") +def test_voageai_without_api_key_or_api_key_file_path(mock_file): + with pytest.raises(Exception): + VoyageAI() + mock_file.assert_not_called() + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 512])) +def test_voageai_with_model_voyage_3_lite(mock_create): + dimension = 512 + model = "voyage-3-lite" + vo = VoyageAI(api_key="API_KEY", model=model) + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_create.assert_called_once_with(texts=["foo"], model=model, input_type=None, truncation=True) + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_with_model_voyage_finance_2(mock_create): + dimension = 1024 + model = "voyage-finance-2" + vo = VoyageAI(api_key="API_KEY", model=model) + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_create.assert_called_once_with(texts=["foo"], model=model, input_type=None, truncation=True) + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_with_model_voyage_multilingual_2(mock_create): + dimension = 1024 + model = "voyage-multilingual-2" + vo = VoyageAI(api_key="API_KEY", model=model) + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_create.assert_called_once_with(texts=["foo"], model=model, input_type=None, truncation=True) + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1024])) +def test_voageai_with_model_voyage_law_2(mock_create): + dimension = 1024 + model = "voyage-law-2" + vo = VoyageAI(api_key="API_KEY", model=model) + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_create.assert_called_once_with(texts=["foo"], model=model, input_type=None, truncation=True) + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1536])) +def test_voageai_with_model_voyage_code_2(mock_create): + dimension = 1536 + model = "voyage-code-2" + vo = VoyageAI(api_key="API_KEY", model=model) + + assert vo.dimension == dimension + assert len(vo.to_embeddings("foo")) == dimension + mock_create.assert_called_once_with(texts=["foo"], model=model, input_type=None, truncation=True) + + +@mock.patch("voyageai.Client.embed", return_value=types.SimpleNamespace(embeddings=[[0] * 1536])) +def test_voageai_with_general_parameters(mock_create): + dimension = 1536 + model = "voyage-code-2" + api_key = "API_KEY" + input_type = "query" + truncation = False + + mock_create.return_value = types.SimpleNamespace(embeddings=[[0] * dimension]) + + vo = VoyageAI(model=model, api_key=api_key, input_type=input_type, truncation=truncation) + assert vo.dimension == dimension + assert len(vo.to_embeddings(["foo"])) == dimension + + mock_create.assert_called_once_with(texts=["foo"], model=model, input_type=input_type, truncation=truncation)