diff --git a/gptcache/embedding/__init__.py b/gptcache/embedding/__init__.py index 08b255c7..3ebd8677 100644 --- a/gptcache/embedding/__init__.py +++ b/gptcache/embedding/__init__.py @@ -18,6 +18,7 @@ from gptcache.utils.lazy_import import LazyImport openai = LazyImport("openai", globals(), "gptcache.embedding.openai") +azureopenai = LazyImport("azureopenai", globals(), "gptcache.embedding.azureopenai") huggingface = LazyImport("huggingface", globals(), "gptcache.embedding.huggingface") sbert = LazyImport("sbert", globals(), "gptcache.embedding.sbert") onnx = LazyImport("onnx", globals(), "gptcache.embedding.onnx") @@ -36,8 +37,12 @@ def Cohere(model="large", api_key=None): return cohere.Cohere(model, api_key) -def OpenAI(model="text-embedding-ada-002", api_key=None): - return openai.OpenAI(model, api_key) +def OpenAI(client, model="text-embedding-ada-002"): + return openai.OpenAIEmbedding(client, model) + + +def AzureOpenAI(client, azure_deployment, model="text-embedding-ada-002"): + return azureopenai.AzureOpenAIEmbedding(client, azure_deployment, model) def Huggingface(model="distilbert-base-uncased"): diff --git a/gptcache/embedding/azureopenai.py b/gptcache/embedding/azureopenai.py new file mode 100644 index 00000000..73dfaed9 --- /dev/null +++ b/gptcache/embedding/azureopenai.py @@ -0,0 +1,82 @@ +from typing import Union + +import numpy as np + +from gptcache.embedding.base import BaseEmbedding +from gptcache.utils import import_openai + +import_openai() + +from openai.lib.azure import AzureOpenAI, AsyncAzureOpenAI + + +class AzureOpenAIEmbedding(BaseEmbedding): + """Generate text embedding for given text using Azure's OpenAI service. + + :param model: OpenAI Client with any modifications you intend to use. + :type model: str + + :param model: model id from the API, defaults to 'text-embedding-ada-002'. + :type model: str + + Example: + .. code-block:: python + + from gptcache.embedding import AzureOpenAIEmbedding + from openai import AzureOpenAI + + test_sentence = 'Hello, world.' + client = AzureOpenAI() + + # You can create different deployments for different embedding models on Azure. + encoder = OpenAIEmbedding(client, azure_deployment='my_embedding_azure_deployment') + + embed = encoder.to_embeddings(test_sentence) + """ + + def __init__(self, + client: Union[AzureOpenAI, AsyncAzureOpenAI], + azure_deployment: str, + model: str = 'text-embedding-ada-002', + ): + """ + + :param client: Azure OpenAI Client class + :type client: Union[AzureOpenAI, AsyncAzureOpenAI] + :param azure_deployment: The deployment name for the embedding; used to generate the endpoint url. + :type azure_deployment: str + :param model: The name of the embedding model; only used for determining the dimensions. Defaults to "text-embedding-ada-002" + :type model: str + """ + self.model = model + self.__azure_embedding_model_deployment = azure_deployment + self.__dimension = self.dim_dict().get(self.model, None) + self.client = client + + def to_embeddings(self, data, **_): + """ + + :param data: String that you wish to convert to an embedding. + :param _: + :return: Array of Float32 numbers that represent the string + """ + sentence_embeddings = self.client.embeddings.create( + input=data, + model=self.__azure_embedding_model_deployment, + ) + return np.array(sentence_embeddings.data[0].embedding).astype("float32") + + @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 {"text-embedding-ada-002": 1536} diff --git a/gptcache/embedding/openai.py b/gptcache/embedding/openai.py index 96767f48..7e8a4270 100644 --- a/gptcache/embedding/openai.py +++ b/gptcache/embedding/openai.py @@ -1,4 +1,4 @@ -import os +from typing import Union import numpy as np @@ -7,39 +7,35 @@ import_openai() -import openai # pylint: disable=C0413 +from openai import OpenAI, AsyncOpenAI -class OpenAI(BaseEmbedding): + +class OpenAIEmbedding(BaseEmbedding): """Generate text embedding for given text using OpenAI. - :param model: model name, defaults to 'text-embedding-ada-002'. + :param model: OpenAI Client with any modifications you intend to use. + :type model: str + + :param model: model id from the API, defaults to 'text-embedding-ada-002'. :type model: str - :param api_key: OpenAI API Key. When the parameter is not specified, it will load the key by default if it is available. - :type api_key: str Example: .. code-block:: python - from gptcache.embedding import OpenAI + from gptcache.embedding import OpenAIEmbedding + from openai import OpenAI test_sentence = 'Hello, world.' - encoder = OpenAI(api_key='your_openai_key') + client = OpenAI(api_key='your_openai_key') + encoder = OpenAIEmbedding(client, model="MyEmbeddingModelId") embed = encoder.to_embeddings(test_sentence) """ - def __init__(self, model: str = "text-embedding-ada-002", api_key: str = None, api_base: str = None): - if not api_key: - if openai.api_key: - api_key = openai.api_key - else: - api_key = os.getenv("OPENAI_API_KEY") - if not api_base: - if openai.api_base: - api_base = openai.api_base - else: - api_base = os.getenv("OPENAI_API_BASE") - openai.api_key = api_key - self.api_base = api_base # don't override all of openai as we may just want to override for say embeddings + def __init__(self, + client: Union[OpenAI, AsyncOpenAI], + model: str = "text-embedding-ada-002", + ): + self.client = client self.model = model if model in self.dim_dict(): self.__dimension = self.dim_dict()[model] @@ -54,8 +50,8 @@ def to_embeddings(self, data, **_): :return: a text embedding in shape of (dim,). """ - sentence_embeddings = openai.Embedding.create(model=self.model, input=data, api_base=self.api_base) - return np.array(sentence_embeddings["data"][0]["embedding"]).astype("float32") + sentence_embeddings = self.client.embeddings.create(model=self.model, input=data) + return np.array(sentence_embeddings.data[0].embedding).astype("float32") @property def dimension(self): diff --git a/requirements.txt b/requirements.txt index f16e4145..4d48d626 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,6 @@ -numpy -cachetools -requests \ No newline at end of file +numpy~=1.26.2 +cachetools~=5.3.2 +requests~=2.31.0 +pydantic~=2.5.1 +httpx~=0.25.1 +setuptools~=60.2.0 \ No newline at end of file