diff --git a/gptcache/__init__.py b/gptcache/__init__.py index 39df80f7..adbd3527 100644 --- a/gptcache/__init__.py +++ b/gptcache/__init__.py @@ -1,5 +1,5 @@ """gptcache version""" -__version__ = "0.1.43" +__version__ = "0.1.44" from gptcache.config import Config from gptcache.core import Cache diff --git a/gptcache/core.py b/gptcache/core.py index db096d7b..ed96fad3 100644 --- a/gptcache/core.py +++ b/gptcache/core.py @@ -126,7 +126,10 @@ def set_azure_openai_key(): openai.api_type = "azure" openai.api_key = os.getenv("OPENAI_API_KEY") - openai.api_base = os.getenv("OPENAI_API_BASE") + if hasattr(openai, "api_base"): + openai.api_base = os.getenv("OPENAI_API_BASE") + elif hasattr(openai, "base_url"): + openai.base_url = os.getenv("OPENAI_BASE_URL", os.getenv("OPENAI_API_BASE")) openai.api_version = os.getenv("OPENAI_API_VERSION") cache = Cache() diff --git a/gptcache/embedding/__init__.py b/gptcache/embedding/__init__.py index 65c2dde4..e813021f 100644 --- a/gptcache/embedding/__init__.py +++ b/gptcache/embedding/__init__.py @@ -44,8 +44,8 @@ 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(model="text-embedding-ada-002", api_key=None, api_base=None, client=None): + return openai.OpenAI(model, api_key, api_base, client) def Huggingface(model="distilbert-base-uncased"): diff --git a/gptcache/embedding/openai.py b/gptcache/embedding/openai.py index 6811144c..aa19e104 100644 --- a/gptcache/embedding/openai.py +++ b/gptcache/embedding/openai.py @@ -29,20 +29,23 @@ class OpenAI(BaseEmbedding): embed = encoder.to_embeddings(test_sentence) """ - def __init__(self, model: str = "text-embedding-ada-002", api_key: str = None, api_base: str = None): + def __init__(self, model: str = "text-embedding-ada-002", api_key: str = None, api_base: str = None, client = 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: + if hasattr(openai, "api_base") and openai.api_base: api_base = openai.api_base + elif hasattr(openai, "base_url") and openai.base_url: + api_base = openai.base_url else: - api_base = os.getenv("OPENAI_API_BASE") + api_base = os.getenv("OPENAI_API_BASE", os.getenv("OPENAI_BASE_URL")) 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 self.model = model + self.client = client if model in self.dim_dict(): self.__dimension = self.dim_dict()[model] else: @@ -56,8 +59,12 @@ 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") + if self.client: + sentence_embeddings = self.client.embeddings.create(model=self.model, input=data) + return np.array(sentence_embeddings.data[0].embedding).astype("float32") + else: + 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") @property def dimension(self): diff --git a/gptcache/manager/factory.py b/gptcache/manager/factory.py index 65d4bb2e..12a0468a 100644 --- a/gptcache/manager/factory.py +++ b/gptcache/manager/factory.py @@ -118,6 +118,12 @@ def manager_factory(manager="map", maxmemory_samples=eviction_params.get("maxmemory_samples", scalar_params.get("maxmemory_samples")), ) + if eviction_manager == "memory": + return get_data_manager(s, v, o, None, + eviction_params.get("max_size", 1000), + eviction_params.get("clean_size", None), + eviction_params.get("eviction", "LRU"),) + e = EvictionBase( name=eviction_manager, **eviction_params @@ -194,7 +200,7 @@ def get_data_manager( vector_base = VectorBase(name=vector_base) if isinstance(object_base, str): object_base = ObjectBase(name=object_base) - if isinstance(eviction_base, str): + if isinstance(eviction_base, str) and eviction_base != "memory": eviction_base = EvictionBase(name=eviction_base) assert cache_base and vector_base return SSDataManager(cache_base, vector_base, object_base, eviction_base, max_size, clean_size, eviction)