Update openai sdk

Author: collindutter

griptape/drivers/embedding/azure_openai_embedding_driver.py

@@ -3,6 +3,7 @@
 from attr import define, field, Factory
 from griptape.drivers import OpenAiEmbeddingDriver
 from griptape.tokenizers import OpenAiTokenizer
+import openai
 
 
 @define
@@ -28,6 +29,17 @@ class AzureOpenAiEmbeddingDriver(OpenAiEmbeddingDriver):
         ),
         kw_only=True,
     )
+    client: openai.AzureOpenAI = field(
+        init=False,
+        default=Factory(
+            lambda self: openai.AzureOpenAI(
+                api_key=self.api_key,
+                base_url=self.base_url,
+                organization=self.organization,
+            ),
+            takes_self=True,
+        ),
+    )
 
     def _params(self, chunk: list[int] | str) -> dict:
         return super()._params(chunk) | {"deployment_id": self.deployment_id}

griptape/drivers/embedding/openai_embedding_driver.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 import os
 from typing import Optional
-import openai
 from attr import define, field, Factory
 from griptape.drivers import BaseEmbeddingDriver
 from griptape.tokenizers import OpenAiTokenizer
+import openai
 
 
 @define
@@ -26,14 +26,21 @@ class OpenAiEmbeddingDriver(BaseEmbeddingDriver):
 
     model: str = field(default=DEFAULT_MODEL, kw_only=True)
     dimensions: int = field(default=DEFAULT_DIMENSIONS, kw_only=True)
-    api_type: str = field(default=openai.api_type, kw_only=True)
-    api_version: Optional[str] = field(default=openai.api_version, kw_only=True)
-    api_base: str = field(default=openai.api_base, kw_only=True)
+    base_url: str = field(default=None, kw_only=True)
     api_key: Optional[str] = field(
         default=Factory(lambda: os.environ.get("OPENAI_API_KEY")), kw_only=True
     )
-    organization: Optional[str] = field(
-        default=openai.organization, kw_only=True
+    organization: Optional[str] = field(default=None, kw_only=True)
+    client: openai.OpenAI = field(
+        init=False,
+        default=Factory(
+            lambda self: openai.OpenAI(
+                api_key=self.api_key,
+                base_url=self.base_url,
+                organization=self.organization,
+            ),
+            takes_self=True,
+        ),
     )
     tokenizer: OpenAiTokenizer = field(
         default=Factory(
@@ -42,29 +49,16 @@ class OpenAiEmbeddingDriver(BaseEmbeddingDriver):
         kw_only=True,
     )
 
-    def __attrs_post_init__(self) -> None:
-        openai.api_type = self.api_type
-        openai.api_version = self.api_version
-        openai.api_base = self.api_base
-        openai.api_key = self.api_key
-        openai.organization = self.organization
-
     def try_embed_chunk(self, chunk: str) -> list[float]:
         # Address a performance issue in older ada models
         # https://github.com/openai/openai-python/issues/418#issuecomment-1525939500
         if self.model.endswith("001"):
             chunk = chunk.replace("\n", " ")
-        return openai.Embedding.create(**self._params(chunk))["data"][0][
-            "embedding"
-        ]
+        return (
+            self.client.embeddings.create(**self._params(chunk))
+            .data[0]
+            .embedding
+        )
 
     def _params(self, chunk: str) -> dict:
-        return {
-            "input": chunk,
-            "model": self.model,
-            "api_key": self.api_key,
-            "organization": self.organization,
-            "api_version": self.api_version,
-            "api_base": self.api_base,
-            "api_type": self.api_type,
-        }
+        return {"input": chunk, "model": self.model}

griptape/drivers/prompt/azure_openai_chat_prompt_driver.py

@@ -1,31 +1,44 @@
+import os
 from attr import define, field, Factory
-from griptape.utils import PromptStack
+from typing import Optional
 from griptape.drivers import OpenAiChatPromptDriver
 from griptape.tokenizers import OpenAiTokenizer
+import openai
 
 
 @define
 class AzureOpenAiChatPromptDriver(OpenAiChatPromptDriver):
     """
     Attributes:
-        api_base: API URL.
+        api_version: API version.
+        azure_deployment: Azure deployment id.
+        azure_endpoint: Azure endpoint.
         deployment_id: Azure OpenAI deployment ID.
         model: OpenAI model name.
     """
 
-    api_base: str = field(kw_only=True)
-    model: str = field(kw_only=True)
-    deployment_id: str = field(kw_only=True)
-    api_type: str = field(default="azure", kw_only=True)
+    azure_deployment: str = field(kw_only=True)
+    azure_endpoint: str = field(kw_only=True)
     api_version: str = field(default="2023-05-15", kw_only=True)
+    api_key: Optional[str] = field(
+        default=Factory(lambda: os.environ.get("AZURE_OPENAI_API_KEY")),
+        kw_only=True,
+    )
+    client: openai.AzureOpenAI = field(
+        default=Factory(
+            lambda self: openai.AzureOpenAI(
+                organization=self.organization,
+                api_key=self.api_key,
+                api_version=self.api_version,
+                azure_endpoint=self.azure_endpoint,
+                azure_deployment=self.azure_deployment,
+            ),
+            takes_self=True,
+        )
+    )
     tokenizer: OpenAiTokenizer = field(
         default=Factory(
             lambda self: OpenAiTokenizer(model=self.model), takes_self=True
         ),
         kw_only=True,
     )
-
-    def _base_params(self, prompt_stack: PromptStack) -> dict:
-        return super()._base_params(prompt_stack) | {
-            "deployment_id": self.deployment_id
-        }

griptape/drivers/prompt/azure_openai_completion_prompt_driver.py

@@ -1,24 +1,28 @@
+import os
 from attr import define, field, Factory
-from griptape.utils import PromptStack
+from typing import Optional
 from griptape.drivers import OpenAiCompletionPromptDriver
-from griptape.tokenizers import OpenAiTokenizer
+import openai
 
 
 @define
 class AzureOpenAiCompletionPromptDriver(OpenAiCompletionPromptDriver):
-    api_base: str = field(kw_only=True)
-    model: str = field(kw_only=True)
-    deployment_id: str = field(kw_only=True)
-    api_type: str = field(default="azure", kw_only=True)
+    azure_deployment: str = field(kw_only=True)
+    azure_endpoint: str = field(kw_only=True)
     api_version: str = field(default="2023-05-15", kw_only=True)
-    tokenizer: OpenAiTokenizer = field(
-        default=Factory(
-            lambda self: OpenAiTokenizer(model=self.model), takes_self=True
-        ),
+    api_key: Optional[str] = field(
+        default=Factory(lambda: os.environ.get("AZURE_OPENAI_API_KEY")),
         kw_only=True,
     )
-
-    def _base_params(self, prompt_stack: PromptStack) -> dict:
-        return super()._base_params(prompt_stack) | {
-            "deployment_id": self.deployment_id
-        }
+    client: openai.AzureOpenAI = field(
+        default=Factory(
+            lambda self: openai.AzureOpenAI(
+                organization=self.organization,
+                api_key=self.api_key,
+                api_version=self.api_version,
+                azure_endpoint=self.azure_endpoint,
+                azure_deployment=self.azure_deployment,
+            ),
+            takes_self=True,
+        )
+    )

griptape/drivers/prompt/base_prompt_driver.py

@@ -32,7 +32,7 @@ class BasePromptDriver(ExponentialBackoffMixin, ABC):
     tokenizer: BaseTokenizer
     stream: bool = field(default=False, kw_only=True)
 
-    def max_output_tokens(self, text: str) -> int:
+    def max_output_tokens(self, text: str | list) -> int:
         tokens_left = self.tokenizer.count_tokens_left(text)
 
         if self.max_tokens:

griptape/drivers/prompt/openai_chat_prompt_driver.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import os
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Any, Literal
 import openai
 from attr import define, field, Factory
 from griptape.artifacts import TextArtifact
@@ -10,23 +10,22 @@
 from typing import Tuple, Type
 import dateparser
 from datetime import datetime, timedelta
-import requests
 
 
 @define
 class OpenAiChatPromptDriver(BasePromptDriver):
     """
     Attributes:
-        api_type: Can be changed to use OpenAI models on Azure.
-        api_version: API version.
-        api_base: API URL.
+        base_url: API URL.
         api_key: API key to pass directly; by default uses `OPENAI_API_KEY_PATH` environment variable.
         max_tokens: Optional maximum return tokens. If not specified, no value will be passed to the API. If set, the
             value will be bounded to the maximum possible as determined by the tokenizer.
         model: OpenAI model name. Uses `gpt-4` by default.
         organization: OpenAI organization.
         tokenizer: Custom `OpenAiTokenizer`.
         user: OpenAI user.
+        response_format: Optional response format. Currently only supports `json_object` which will enable OpenAi's JSON mode.
+        seed: Optional seed.
         _ratelimit_request_limit: The maximum number of requests allowed in the current rate limit window.
         _ratelimit_requests_remaining: The number of requests remaining in the current rate limit window.
         _ratelimit_requests_reset_at: The time at which the current rate limit window resets.
@@ -35,14 +34,23 @@ class OpenAiChatPromptDriver(BasePromptDriver):
         _ratelimit_tokens_reset_at: The time at which the current rate limit window resets.
     """
 
-    api_type: str = field(default=openai.api_type, kw_only=True)
-    api_version: Optional[str] = field(default=openai.api_version, kw_only=True)
-    api_base: str = field(default=openai.api_base, kw_only=True)
+    base_url: str = field(default=None, kw_only=True)
     api_key: Optional[str] = field(
         default=Factory(lambda: os.environ.get("OPENAI_API_KEY")), kw_only=True
     )
     organization: Optional[str] = field(
-        default=openai.organization, kw_only=True
+        default=os.environ.get("OPENAI_ORG_ID"), kw_only=True
+    )
+    seed: Optional[int] = field(default=None, kw_only=True)
+    client: openai.OpenAI = field(
+        default=Factory(
+            lambda self: openai.OpenAI(
+                api_key=self.api_key,
+                base_url=self.base_url,
+                organization=self.organization,
+            ),
+            takes_self=True,
+        )
     )
     model: str = field(kw_only=True)
     tokenizer: OpenAiTokenizer = field(
@@ -52,8 +60,11 @@ class OpenAiChatPromptDriver(BasePromptDriver):
         kw_only=True,
     )
     user: str = field(default="", kw_only=True)
+    response_format: Optional[Literal["json_object"]] = field(
+        default=None, kw_only=True
+    )
     ignored_exception_types: Tuple[Type[Exception], ...] = field(
-        default=Factory(lambda: openai.InvalidRequestError), kw_only=True
+        default=Factory(lambda: openai.BadRequestError), kw_only=True
     )
     _ratelimit_request_limit: Optional[int] = field(init=False, default=None)
     _ratelimit_requests_remaining: Optional[int] = field(
@@ -68,40 +79,36 @@ class OpenAiChatPromptDriver(BasePromptDriver):
         init=False, default=None
     )
 
-    def __attrs_post_init__(self) -> None:
-        # Define a hook to pull rate limit metadata from the OpenAI API response header.
-        openai.requestssession = requests.Session()
-        openai.requestssession.hooks = {
-            "response": self._extract_ratelimit_metadata
-        }
-
     def try_run(self, prompt_stack: PromptStack) -> TextArtifact:
-        result = openai.ChatCompletion.create(**self._base_params(prompt_stack))
+        result = self.client.chat.completions.with_raw_response.create(
+            **self._base_params(prompt_stack)
+        )
+
+        self._extract_ratelimit_metadata(result)
 
+        result = result.parse()
         if len(result.choices) == 1:
-            return TextArtifact(
-                value=result.choices[0]["message"]["content"].strip()
-            )
+            return TextArtifact(value=result.choices[0].message.content.strip())
         else:
             raise Exception(
                 "Completion with more than one choice is not supported yet."
             )
 
     def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]:
-        result = openai.ChatCompletion.create(
+        result = self.client.chat.completions.create(
             **self._base_params(prompt_stack), stream=True
         )
 
         for chunk in result:
             if len(chunk.choices) == 1:
-                delta = chunk.choices[0]["delta"]
+                delta = chunk.choices[0].delta
             else:
                 raise Exception(
                     "Completion with more than one choice is not supported yet."
                 )
 
-            if "content" in delta:
-                delta_content = delta["content"]
+            if delta.content is not None:
+                delta_content = delta.content
 
                 yield TextArtifact(value=delta_content)
 
@@ -112,33 +119,37 @@ def token_count(self, prompt_stack: PromptStack) -> int:
 
     def _prompt_stack_to_messages(
         self, prompt_stack: PromptStack
-    ) -> list[dict]:
+    ) -> list[dict[str, Any]]:
         return [
             {"role": self.__to_openai_role(i), "content": i.content}
             for i in prompt_stack.inputs
         ]
 
     def _base_params(self, prompt_stack: PromptStack) -> dict:
-        messages = self._prompt_stack_to_messages(prompt_stack)
-
         params = {
             "model": self.model,
             "temperature": self.temperature,
             "stop": self.tokenizer.stop_sequences,
             "user": self.user,
-            "api_key": self.api_key,
-            "organization": self.organization,
-            "api_version": self.api_version,
-            "api_base": self.api_base,
-            "api_type": self.api_type,
-            "messages": messages,
+            "seed": self.seed,
         }
 
+        if self.response_format == "json_object":
+            params["response_format"] = {"type": "json_object"}
+            # JSON mode still requires a system input instructing the LLM to output JSON.
+            prompt_stack.add_system_input(
+                "Provide your response as a valid JSON object."
+            )
+
+        messages = self._prompt_stack_to_messages(prompt_stack)
+
         # A max_tokens parameter is not required, but if it is specified by the caller, bound it to
         # the maximum value as determined by the tokenizer and pass it to the API.
         if self.max_tokens:
             params["max_tokens"] = self.max_output_tokens(messages)
 
+        params["messages"] = messages
+
         return params
 
     def __to_openai_role(self, prompt_input: PromptStack.Input) -> str:
@@ -149,7 +160,7 @@ def __to_openai_role(self, prompt_input: PromptStack.Input) -> str:
         else:
             return "user"
 
-    def _extract_ratelimit_metadata(self, response, *args, **kwargs):
+    def _extract_ratelimit_metadata(self, response):
         # The OpenAI SDK's requestssession variable is global, so this hook will fire for all API requests.
         # The following headers are not reliably returned in every API call, so we check for the presence of the
         # headers before reading and parsing their values to prevent other SDK users from encountering KeyErrors.