Add OpenAI v1.x.x compatibility to embedding models and api base (#641)

* Fix the nil memory eviction when using the init_similar_cache method Signed-off-by: SimFG <[email protected]> * Update the version to 0.1.44 Signed-off-by: SimFG <[email protected]> * Add OpenAI v1.x.x compatibility to embedding models and api base - Support passing the client as the parameter - In OpenAI v1.x.x, `api_base` has been renamed into `base_url` openai/openai-python#1051 (comment) Signed-off-by: Hollow Man <[email protected]> --------- Signed-off-by: SimFG <[email protected]> Signed-off-by: Hollow Man <[email protected]> Co-authored-by: SimFG <[email protected]>
zilliztech · Sep 3, 2024 · 17646e1 · 17646e1
1 parent a507074
commit 17646e1
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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)