Remove init=False, improve comments

collindutter · collindutter · commit 95213e8215ed · 2023-11-09T12:21:05.000-08:00
diff --git a/griptape/drivers/embedding/azure_openai_embedding_driver.py b/griptape/drivers/embedding/azure_openai_embedding_driver.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import Optional
 from attr import define, field, Factory
 from griptape.drivers import OpenAiEmbeddingDriver
 from griptape.tokenizers import OpenAiTokenizer
@@ -10,29 +11,36 @@
 class AzureOpenAiEmbeddingDriver(OpenAiEmbeddingDriver):
     """
     Attributes:
-        model: OpenAI embedding model name.
-        deployment_id: Azure OpenAI deployment ID.
-        api_base: API URL.
-        api_type: OpenAI API type. Defaults to 'azure'.
-        api_version: API version. Defaults to '2023-05-15'.
-        tokenizer: Optionally provide custom `OpenAiTokenizer`.
+        azure_deployment: An Azure OpenAi deployment id.
+        azure_endpoint: An Azure OpenAi endpoint.
+        azure_ad_token: An optional Azure Active Directory token.
+        azure_ad_token_provider: An optional Azure Active Directory token provider.
+        api_version: An Azure OpenAi API version.
+        tokenizer: An `OpenAiTokenizer`.
+        client: An `openai.AzureOpenAI` client.
     """
 
-    deployment_id: str = field(kw_only=True)
-    api_base: 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)
+    azure_ad_token: Optional[str] = field(kw_only=True, default=None)
+    azure_ad_token_provider: Optional[str] = field(kw_only=True, default=None)
     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), 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
+                organization=self.organization,
+                api_key=self.api_key,
+                api_version=self.api_version,
+                azure_endpoint=self.azure_endpoint,
+                azure_deployment=self.azure_deployment,
+                azure_ad_token=self.azure_ad_token,
+                azure_ad_token_provider=self.azure_ad_token_provider,
             ),
             takes_self=True,
-        ),
+        )
     )
 
     def _params(self, chunk: str) -> dict:
diff --git a/griptape/drivers/embedding/openai_embedding_driver.py b/griptape/drivers/embedding/openai_embedding_driver.py
@@ -16,6 +16,12 @@ class OpenAiEmbeddingDriver(BaseEmbeddingDriver):
         api_key: API key to pass directly. Defaults to `OPENAI_API_KEY` environment variable.
         organization: OpenAI organization. Defaults to 'OPENAI_ORGANIZATION' environment variable.
         tokenizer: Optionally provide custom `OpenAiTokenizer`.
+        client: Optionally provide custom `openai.OpenAI` client.
+        azure_deployment: An Azure OpenAi deployment id.
+        azure_endpoint: An Azure OpenAi endpoint.
+        azure_ad_token: An optional Azure Active Directory token.
+        azure_ad_token_provider: An optional Azure Active Directory token provider.
+        api_version: An Azure OpenAi API version.
     """
 
     DEFAULT_MODEL = "text-embedding-ada-002"
@@ -27,11 +33,10 @@ class OpenAiEmbeddingDriver(BaseEmbeddingDriver):
     api_key: Optional[str] = field(default=None, 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(lambda self: OpenAiTokenizer(model=self.model), takes_self=True), kw_only=True
diff --git a/griptape/drivers/prompt/azure_openai_chat_prompt_driver.py b/griptape/drivers/prompt/azure_openai_chat_prompt_driver.py
@@ -2,19 +2,19 @@
 from typing import Optional
 from griptape.utils import PromptStack
 from griptape.drivers import OpenAiChatPromptDriver
-from griptape.tokenizers import OpenAiTokenizer
 import openai
 
 
 @define
 class AzureOpenAiChatPromptDriver(OpenAiChatPromptDriver):
     """
     Attributes:
-        azure_deployment: Azure deployment id.
-        azure_endpoint: Azure endpoint.
-        azure_ad_token: Azure Active Directory token.
-        azure_ad_token_provider: Azure Active Directory token provider.
-        api_version: API version.
+        azure_deployment: An Azure OpenAi deployment id.
+        azure_endpoint: An Azure OpenAi endpoint.
+        azure_ad_token: An optional Azure Active Directory token.
+        azure_ad_token_provider: An optional Azure Active Directory token provider.
+        api_version: An Azure OpenAi API version.
+        client: An `openai.AzureOpenAI` client.
     """
 
     azure_deployment: str = field(kw_only=True)
@@ -36,9 +36,6 @@ class AzureOpenAiChatPromptDriver(OpenAiChatPromptDriver):
             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:
         params = super()._base_params(prompt_stack)
diff --git a/griptape/drivers/prompt/azure_openai_completion_prompt_driver.py b/griptape/drivers/prompt/azure_openai_completion_prompt_driver.py
@@ -6,6 +6,16 @@
 
 @define
 class AzureOpenAiCompletionPromptDriver(OpenAiCompletionPromptDriver):
+    """
+    Attributes:
+        azure_deployment: An Azure OpenAi deployment id.
+        azure_endpoint: An Azure OpenAi endpoint.
+        azure_ad_token: An optional Azure Active Directory token.
+        azure_ad_token_provider: An optional Azure Active Directory token provider.
+        api_version: An Azure OpenAi API version.
+        client: An `openai.AzureOpenAI` client.
+    """
+
     azure_deployment: str = field(kw_only=True)
     azure_endpoint: str = field(kw_only=True)
     azure_ad_token: Optional[str] = field(kw_only=True, default=None)
diff --git a/griptape/drivers/prompt/base_prompt_driver.py b/griptape/drivers/prompt/base_prompt_driver.py
@@ -14,6 +14,19 @@
 
 @define
 class BasePromptDriver(ExponentialBackoffMixin, ABC):
+    """Base class for Prompt Drivers.
+
+    Attributes:
+        temperature: The temperature to use for the completion.
+        max_tokens: The maximum number of tokens to generate. If not specified, the value will be automatically generated based by the tokenizer.
+        structure: An optional `Structure` to publish events to.
+        prompt_stack_to_string: A function that converts a `PromptStack` to a string.
+        ignored_exception_types: A tuple of exception types to ignore.
+        model: The model name.
+        tokenizer: An instance of `BaseTokenizer` to when calculating tokens.
+        stream: Whether to stream the completion or not. `CompletionChunkEvent`s will be published to the `Structure` if one is provided.
+    """
+
     temperature: float = field(default=0.1, kw_only=True)
     max_tokens: Optional[int] = field(default=None, kw_only=True)
     structure: Optional[Structure] = field(default=None, kw_only=True)
diff --git a/griptape/drivers/prompt/openai_chat_prompt_driver.py b/griptape/drivers/prompt/openai_chat_prompt_driver.py
@@ -15,17 +15,16 @@
 class OpenAiChatPromptDriver(BasePromptDriver):
     """
     Attributes:
-        base_url: API URL.
-        api_key: API key to pass directly. Defaults to `OPENAI_API_KEY` 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.
-        organization: OpenAI organization. Defaults to `OPENAI_ORG_ID` environment variable.
-        client: OpenAI client. Defaults to `openai.OpenAI`.
-        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.
+        base_url: An optional OpenAi API URL.
+        api_key: An optional OpenAi API key. If not provided, the `OPENAI_API_KEY` environment variable will be used.
+        organization: An optional OpenAI organization. If not provided, the `OPENAI_ORG_ID` environment variable will be used.
+        client: An `openai.OpenAI` client.
+        model: An OpenAI model name.
+        tokenizer: An `OpenAiTokenizer`.
+        user: An optional user id. Can be used to track requests by user.
+        response_format: An optional OpenAi Chat Completion response format. Currently only supports `json_object` which will enable OpenAi's JSON mode.
+        seed: An optional OpenAi Chat Completion seed.
+        ignored_exception_types: An optional tuple of exception types to ignore. Defaults to OpenAI's known exception types.
         _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.
@@ -34,10 +33,9 @@ class OpenAiChatPromptDriver(BasePromptDriver):
         _ratelimit_tokens_reset_at: The time at which the current rate limit window resets.
     """
 
-    base_url: str = field(default=None, kw_only=True)
+    base_url: Optional[str] = field(default=None, kw_only=True)
     api_key: Optional[str] = field(default=None, kw_only=True)
     organization: Optional[str] = field(default=None, 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),
@@ -48,8 +46,9 @@ class OpenAiChatPromptDriver(BasePromptDriver):
     tokenizer: OpenAiTokenizer = field(
         default=Factory(lambda self: OpenAiTokenizer(model=self.model), takes_self=True), kw_only=True
     )
-    user: str = field(default="", kw_only=True)
+    user: Optional[str] = field(default=None, kw_only=True)
     response_format: Optional[Literal["json_object"]] = field(default=None, kw_only=True)
+    seed: Optional[int] = field(default=None, kw_only=True)
     ignored_exception_types: Tuple[Type[Exception], ...] = field(
         default=Factory(
             lambda: (
diff --git a/griptape/drivers/prompt/openai_completion_prompt_driver.py b/griptape/drivers/prompt/openai_completion_prompt_driver.py
@@ -12,17 +12,17 @@
 class OpenAiCompletionPromptDriver(BasePromptDriver):
     """
     Attributes:
-        base_url: API URL.
-        api_key: API key to pass directly. Defaults to `OPENAI_API_KEY` environment variable.
-        max_tokens: Optional maximum return tokens. If not specified, the value will be automatically generated based by the tokenizer.
-        model: OpenAI model name. Uses `gpt-4` by default.
-        organization: OpenAI organization. Defaults to `OPENAI_ORG_ID` environment variable.
-        client: OpenAI client. Defaults to `openai.OpenAI`.
-        tokenizer: Custom `OpenAiTokenizer`.
-        user: OpenAI user.
+        base_url: An optional OpenAi API URL.
+        api_key: An optional OpenAi API key. If not provided, the `OPENAI_API_KEY` environment variable will be used.
+        organization: An optional OpenAI organization. If not provided, the `OPENAI_ORG_ID` environment variable will be used.
+        client: An `openai.OpenAI` client.
+        model: An OpenAI model name.
+        tokenizer: An `OpenAiTokenizer`.
+        user: An optional user id. Can be used to track requests by user.
+        ignored_exception_types: An optional tuple of exception types to ignore. Defaults to OpenAI's known exception types.
     """
 
-    base_url: str = field(default=None, kw_only=True)
+    base_url: Optional[str] = field(default=None, kw_only=True)
     api_key: Optional[str] = field(default=None, kw_only=True)
     organization: Optional[str] = field(default=None, kw_only=True)
     client: openai.OpenAI = field(
@@ -35,7 +35,7 @@ class OpenAiCompletionPromptDriver(BasePromptDriver):
     tokenizer: OpenAiTokenizer = field(
         default=Factory(lambda self: OpenAiTokenizer(model=self.model), takes_self=True), kw_only=True
     )
-    user: str = field(default="", kw_only=True)
+    user: Optional[str] = field(default=None, kw_only=True)
     ignored_exception_types: Tuple[Type[Exception], ...] = field(
         default=Factory(
             lambda: (
diff --git a/griptape/tokenizers/openai_tokenizer.py b/griptape/tokenizers/openai_tokenizer.py
@@ -15,6 +15,7 @@ class OpenAiTokenizer(BaseTokenizer):
     DEFAULT_MAX_TOKENS = 2049
     TOKEN_OFFSET = 8
 
+    # https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo
     MODEL_PREFIXES_TO_MAX_TOKENS = {
         "gpt-4-1106": 128000,
         "gpt-4-32k": 32768,
