Merge pull request #261 from DanielUH2019/backwards-compatibility-for-rm

openai 0.28 backwards compatibility for pinecone and chromadb

Author: okhat

dspy/retrieve/chromadb_rm.py

@@ -8,6 +8,12 @@
 import backoff
 from dsp.utils import dotdict
 
+try:
+    import openai.error
+    ERRORS = (openai.error.RateLimitError, openai.error.ServiceUnavailableError, openai.error.APIError)
+except Exception:
+    ERRORS = (openai.RateLimitError, openai.APIError)
+
 try:
     import chromadb
     from chromadb.config import Settings
@@ -108,7 +114,7 @@ def _init_chromadb(
 
     @backoff.on_exception(
         backoff.expo,
-        (openai.RateLimitError),
+        ERRORS,
         max_time=15,
     )
     def _get_embeddings(self, queries: List[str]) -> List[List[float]]:

dspy/retrieve/pinecone_rm.py

@@ -5,7 +5,6 @@
 
 from dsp.utils import dotdict
 from typing import Optional, List, Union
-import openai
 import dspy
 import backoff
 
@@ -19,6 +18,17 @@
         "The pinecone library is required to use PineconeRM. Install it with `pip install dspy-ai[pinecone]`"
     )
 
+import openai
+try:
+    OPENAI_LEGACY = int(openai.version.__version__[0]) == 0
+except Exception:
+    OPENAI_LEGACY = True
+
+try:
+    import openai.error
+    ERRORS = (openai.error.RateLimitError, openai.error.ServiceUnavailableError, openai.error.APIError)
+except Exception:
+    ERRORS = (openai.RateLimitError, openai.APIError)
 
 class PineconeRM(dspy.Retrieve):
     """
@@ -164,7 +174,7 @@ def _mean_pooling(
 
     @backoff.on_exception(
         backoff.expo,
-        (openai.RateLimitError),
+        ERRORS,
         max_time=15,
     )
     def _get_embeddings(
@@ -187,10 +197,15 @@ def _get_embeddings(
             ) from exc
 
         if not self.use_local_model:
-            embedding = openai.embeddings.create(
-                input=queries, model=self._openai_embed_model
-            )
-            return [embedding.embedding for embedding in embedding.data]
+            if OPENAI_LEGACY:
+                embedding = openai.Embedding.create(
+                    input=queries, model=self._openai_embed_model
+                )
+            else:
+                embedding = openai.embeddings.create(
+                    input=queries, model=self._openai_embed_model
+                ).model_dump()
+            return [embedding["embedding"] for embedding in embedding["data"]]
 
         # Use local model
         encoded_input = self._local_tokenizer(queries, padding=True, truncation=True, return_tensors="pt").to(self.device)