Merge branch 'mongodb-atlas-rm'

Author: jhyearsley

README.md

@@ -58,6 +58,7 @@ Or open our intro notebook in Google Colab: [<img align="center" src="https://co
 
 > _Note: If you're looking for Demonstrate-Search-Predict (DSP), which is the previous version of DSPy, you can find it on the [v1](https://github.com/stanfordnlp/dspy/tree/v1) branch of this repo._
 
+By default, DSPy depends on `openai==0.28`. However, if you install `openai>=1.0`, the library will use that just fine. Both are supported.
 
 For the optional Pinecone, Qdrant, [chromadb](https://github.com/chroma-core/chroma), or  [marqo](https://github.com/marqo-ai/marqo) retrieval integration(s), include the extra(s) below:
 

dsp/evaluation/utils.py

@@ -1,5 +1,3 @@
-from openai import InvalidRequestError
-from openai.error import APIError
 
 import dsp
 import tqdm

dsp/modules/gpt3.py

@@ -4,12 +4,23 @@
 
 import backoff
 import openai
-import openai.error
-from openai.openai_object import OpenAIObject
 
 from dsp.modules.cache_utils import CacheMemory, NotebookCacheMemory, cache_turn_on
 from dsp.modules.lm import LM
 
+try:
+    OPENAI_LEGACY = int(openai.version.__version__[0]) == 0
+except Exception:
+    OPENAI_LEGACY = True
+
+try:
+    from openai.openai_object import OpenAIObject
+    import openai.error
+    ERRORS = (openai.error.RateLimitError, openai.error.ServiceUnavailableError, openai.error.APIError)
+except Exception:
+    ERRORS = (openai.RateLimitError, openai.APIError)
+    OpenAIObject = dict
+
 
 def backoff_hdlr(details):
     """Handler from https://pypi.org/project/backoff/"""
@@ -36,13 +47,19 @@ def __init__(
         model: str = "gpt-3.5-turbo-instruct",
         api_key: Optional[str] = None,
         api_provider: Literal["openai", "azure"] = "openai",
+        api_base: Optional[str] = None,
         model_type: Literal["chat", "text"] = None,
         **kwargs,
     ):
         super().__init__(model)
         self.provider = "openai"
 
-        default_model_type = "chat" if ('gpt-3.5' in model or 'turbo' in model or 'gpt-4' in model) and ('instruct' not in model) else "text"
+        default_model_type = (
+            "chat"
+            if ("gpt-3.5" in model or "turbo" in model or "gpt-4" in model)
+            and ("instruct" not in model)
+            else "text"
+        )
         self.model_type = model_type if model_type else default_model_type
 
         if api_provider == "azure":
@@ -58,8 +75,8 @@ def __init__(
         if api_key:
             openai.api_key = api_key
 
-        if kwargs.get("api_base"):
-            openai.api_base = kwargs["api_base"]
+        if api_base:
+            openai.base_url = api_base
 
         self.kwargs = {
             "temperature": 0.0,
@@ -70,29 +87,27 @@ def __init__(
             "n": 1,
             **kwargs,
         }  # TODO: add kwargs above for </s>
-        
+
         if api_provider != "azure":
             self.kwargs["model"] = model
         self.history: list[dict[str, Any]] = []
 
-    def _openai_client():
+    def _openai_client(self):
         return openai
 
-    def basic_request(self, prompt: str, **kwargs) -> OpenAIObject:
+    def basic_request(self, prompt: str, **kwargs):
         raw_kwargs = kwargs
 
         kwargs = {**self.kwargs, **kwargs}
         if self.model_type == "chat":
             # caching mechanism requires hashable kwargs
             kwargs["messages"] = [{"role": "user", "content": prompt}]
-            kwargs = {
-                "stringify_request": json.dumps(kwargs)
-            }
-            response = cached_gpt3_turbo_request(**kwargs)
-            
+            kwargs = {"stringify_request": json.dumps(kwargs)}
+            response = chat_request(**kwargs)
+
         else:
             kwargs["prompt"] = prompt
-            response = cached_gpt3_request(**kwargs)
+            response = completions_request(**kwargs)
 
         history = {
             "prompt": prompt,
@@ -106,15 +121,15 @@ def basic_request(self, prompt: str, **kwargs) -> OpenAIObject:
 
     @backoff.on_exception(
         backoff.expo,
-        (openai.error.RateLimitError, openai.error.ServiceUnavailableError, openai.error.APIError),
+        ERRORS,
         max_time=1000,
         on_backoff=backoff_hdlr,
     )
-    def request(self, prompt: str, **kwargs) -> OpenAIObject:
+    def request(self, prompt: str, **kwargs):
         """Handles retreival of GPT-3 completions whilst handling rate limiting and caching."""
         if "model_type" in kwargs:
             del kwargs["model_type"]
-        
+
         return self.basic_request(prompt, **kwargs)
 
     def _get_choice_text(self, choice: dict[str, Any]) -> str:
@@ -150,6 +165,7 @@ def __call__(
         #         kwargs = {**kwargs, "logprobs": 5}
 
         response = self.request(prompt, **kwargs)
+
         choices = response["choices"]
 
         completed_choices = [c for c in choices if c["finish_reason"] != "length"]
@@ -158,7 +174,6 @@ def __call__(
             choices = completed_choices
 
         completions = [self._get_choice_text(c) for c in choices]
-
         if return_sorted and kwargs.get("n", 1) > 1:
             scored_completions = []
 
@@ -181,31 +196,57 @@ def __call__(
         return completions
 
 
+
 @CacheMemory.cache
 def cached_gpt3_request_v2(**kwargs):
     return openai.Completion.create(**kwargs)
 
-
 @functools.lru_cache(maxsize=None if cache_turn_on else 0)
 @NotebookCacheMemory.cache
 def cached_gpt3_request_v2_wrapped(**kwargs):
     return cached_gpt3_request_v2(**kwargs)
 
-
-cached_gpt3_request = cached_gpt3_request_v2_wrapped
-
-
 @CacheMemory.cache
 def _cached_gpt3_turbo_request_v2(**kwargs) -> OpenAIObject:
     if "stringify_request" in kwargs:
         kwargs = json.loads(kwargs["stringify_request"])
     return cast(OpenAIObject, openai.ChatCompletion.create(**kwargs))
 
-
 @functools.lru_cache(maxsize=None if cache_turn_on else 0)
 @NotebookCacheMemory.cache
 def _cached_gpt3_turbo_request_v2_wrapped(**kwargs) -> OpenAIObject:
     return _cached_gpt3_turbo_request_v2(**kwargs)
 
+@CacheMemory.cache
+def v1_cached_gpt3_request_v2(**kwargs):
+    return openai.completions.create(**kwargs)
+
+@functools.lru_cache(maxsize=None if cache_turn_on else 0)
+@NotebookCacheMemory.cache
+def v1_cached_gpt3_request_v2_wrapped(**kwargs):
+    return v1_cached_gpt3_request_v2(**kwargs)
+
+@CacheMemory.cache
+def v1_cached_gpt3_turbo_request_v2(**kwargs):
+    if "stringify_request" in kwargs:
+        kwargs = json.loads(kwargs["stringify_request"])
+    return openai.chat.completions.create(**kwargs)
+
+@functools.lru_cache(maxsize=None if cache_turn_on else 0)
+@NotebookCacheMemory.cache
+def v1_cached_gpt3_turbo_request_v2_wrapped(**kwargs):
+    return v1_cached_gpt3_turbo_request_v2(**kwargs)
+
+
+
+def chat_request(**kwargs):
+    if OPENAI_LEGACY:
+        return _cached_gpt3_turbo_request_v2_wrapped(**kwargs)
+
+    return v1_cached_gpt3_turbo_request_v2_wrapped(**kwargs).model_dump()
+
+def completions_request(**kwargs):
+    if OPENAI_LEGACY:
+        return cached_gpt3_request_v2_wrapped(**kwargs)
 
-cached_gpt3_turbo_request = _cached_gpt3_turbo_request_v2_wrapped
+    return v1_cached_gpt3_request_v2_wrapped(**kwargs).model_dump()

dsp/modules/sentence_vectorizer.py

@@ -109,6 +109,12 @@ def __call__(self, inp_examples: List["Example"]) -> np.ndarray:
         return embeddings
 
 
+try:
+    OPENAI_LEGACY = int(openai.version.__version__[0]) == 0
+except Exception:
+    OPENAI_LEGACY = True
+
+
 class OpenAIVectorizer(BaseSentenceVectorizer):
     '''
     This vectorizer uses OpenAI API to convert texts to embeddings. Changing `model` is not
@@ -124,6 +130,11 @@ def __init__(
         self.model = model
         self.embed_batch_size = embed_batch_size
 
+        if OPENAI_LEGACY:
+            self.Embedding = openai.Embedding
+        else:
+            self.Embedding = openai.embeddings
+
         if api_key:
             openai.api_key = api_key
 
@@ -138,7 +149,7 @@ def __call__(self, inp_examples: List["Example"]) -> np.ndarray:
             end_idx = (cur_batch_idx + 1) * self.embed_batch_size
             cur_batch = text_to_vectorize[start_idx: end_idx]
             # OpenAI API call:
-            response = openai.Embedding.create(
+            response = self.Embedding.create(
                 model=self.model,
                 input=cur_batch
             )
@@ -147,4 +158,4 @@ def __call__(self, inp_examples: List["Example"]) -> np.ndarray:
             embeddings_list.extend(cur_batch_embeddings)
 
         embeddings = np.array(embeddings_list, dtype=np.float32)
-        return embeddings
+        return embeddings

dspy/retrieve/chromadb_rm.py

@@ -11,6 +11,7 @@
 try:
     import chromadb
     from chromadb.config import Settings
+    from chromadb.utils import embedding_functions
 except ImportError:
     chromadb = None
 
@@ -70,17 +71,13 @@ def __init__(
 
         self._init_chromadb(collection_name, persist_directory)
 
-        # If not provided, defaults to env vars
-        if openai_api_key:
-            openai.api_key = openai_api_key
-        if openai_api_type:
-            openai.api_type = openai_api_type
-        if openai_api_base:
-            openai.api_base = openai_api_base
-        if openai_api_version:
-            openai.api_version = openai_api_version
-        if openai_api_provider:
-            self._openai_api_provider = openai_api_provider
+        self.openai_ef = embedding_functions.OpenAIEmbeddingFunction(
+            api_key=openai_api_key,
+            api_base=openai_api_base,
+            api_type=openai_api_type,
+            api_version=openai_api_version,
+            model_name=openai_embed_model,
+        )
 
         super().__init__(k=k)
 
@@ -111,7 +108,7 @@ def _init_chromadb(
 
     @backoff.on_exception(
         backoff.expo,
-        (openai.error.RateLimitError, openai.error.ServiceUnavailableError),
+        (openai.RateLimitError),
         max_time=15,
     )
     def _get_embeddings(self, queries: List[str]) -> List[List[float]]:
@@ -124,24 +121,10 @@ def _get_embeddings(self, queries: List[str]) -> List[List[float]]:
             List[List[float]]: List of embeddings corresponding to each query.
         """
 
-        if self._openai_api_provider == "azure":
-            model_args = {
-                "engine": self._openai_embed_model,
-                "deployment_id": self._openai_embed_model,
-                "api_version": openai.api_version,
-                "api_base": openai.api_base,
-            }
-            embedding = openai.Embedding.create(
-                input=queries,
-                model=self._openai_embed_model,
-                **model_args,
-                api_provider=self._openai_api_provider
-            )
-        else:
-            embedding = openai.Embedding.create(
-                input=queries, model=self._openai_embed_model
-            )
-        return [embedding["embedding"] for embedding in embedding["data"]]
+        embedding = self.openai_ef._client.create(
+            input=queries, model=self._openai_embed_model
+        )
+        return [embedding.embedding for embedding in embedding.data]
 
     def forward(
         self, query_or_queries: Union[str, List[str]], k: Optional[int] = None

dspy/retrieve/mongodb_atlas_rm.py

@@ -1,7 +1,13 @@
-from typing import List, Union, Any
+from typing import List, Optional, Union, Any
 import dspy
 import os
-import openai
+from openai import (
+    OpenAI,
+    APITimeoutError,
+    InternalServerError,
+    RateLimitError,
+    UnprocessableEntityError,
+)
 import backoff
 
 try:
@@ -39,24 +45,25 @@ def build_vector_search_pipeline(
 class Embedder:
     def __init__(self, provider: str, model: str):
         if provider == "openai":
-            openai.api_key = os.getenv("OPENAI_API_KEY")
-            if not openai.api_key:
+            api_key = os.getenv("OPENAI_API_KEY")
+            if not api_key:
                 raise ValueError("Environment variable OPENAI_API_KEY must be set")
-            self.client = openai
+            self.client = OpenAI()
             self.model = model
 
     @backoff.on_exception(
         backoff.expo,
         (
-            openai.error.RateLimitError,
-            openai.error.ServiceUnavailableError,
-            openai.error.APIError,
+            APITimeoutError,
+            InternalServerError,
+            RateLimitError,
+            UnprocessableEntityError,
         ),
         max_time=15,
     )
     def __call__(self, queries) -> Any:
-        embedding = self.client.Embedding.create(input=queries, model=self.model)
-        return [embedding["embedding"] for embedding in embedding["data"]]
+        embedding = self.client.embeddings.create(input=queries, model=self.model)
+        return [result.embedding for result in embedding.data]
 
 
 class MongoDBAtlasRM(dspy.Retrieve):
@@ -98,13 +105,8 @@ def __init__(
 
         self.embedder = Embedder(provider=embedding_provider, model=embedding_model)
 
-    def forward(self, query_or_queries: Union[str, List[str]]) -> dspy.Prediction:
-        queries = (
-            [query_or_queries]
-            if isinstance(query_or_queries, str)
-            else query_or_queries
-        )
-        query_vector = self.embedder(queries)
+    def forward(self, query_or_queries: str) -> dspy.Prediction:
+        query_vector = self.embedder([query_or_queries])
         pipeline = build_vector_search_pipeline(
             index_name=self.index_name,
             query_vector=query_vector[0],