Merge branch 'docs-2' of https://github.com/PrefectHQ/marvin into docs-2

cookbook/azure/README.md

@@ -6,19 +6,23 @@ It is possible to use Azure OpenAI with _some_ of `marvin`'s functionality via:
 !!! Note
     Azure OpenAI often lags behind the latest version of OpenAI in terms of functionality, therefore some features may not work with Azure OpenAI. If you encounter problems, please check that the underlying functionality is supported by Azure OpenAI before reporting an issue.
 
-## Settings
-After setting up your Azure OpenAI account and deployment, it is recommended you save these settings in your `~/.marvin/.env` file.
+## Configuring with environment variables
+After setting up your Azure OpenAI account and deployment, set these environment variables in your environment, `~/.marvin/.env`, or `.env` file:
 
 ```bash
-» cat ~/.marvin/.env | rg AZURE
-MARVIN_USE_AZURE_OPENAI=true
+MARVIN_PROVIDER=azure_openai
 MARVIN_AZURE_OPENAI_API_KEY=<your-api-key>
-MARVIN_AZURE_OPENAI_ENDPOINT=https://<your-endpoint>.openai.azure.com/
-MARVIN_AZURE_OPENAI_API_VERSION=2023-12-01-preview # or whatever is the latest
-MARVIN_AZURE_OPENAI_DEPLOYMENT_NAME=gpt-35-turbo-0613 # or whatever you named your deployment
+MARVIN_AZURE_OPENAI_ENDPOINT="https://<your-endpoint>.openai.azure.com/"
+MARVIN_AZURE_OPENAI_API_VERSION=2023-12-01-preview # or latest
+
+MARVIN_CHAT_COMPLETION_MODEL=<your azure openai deployment name>
 ```
 
-## Passing a `MarvinClient` set up with `AzureOpenAI` manually
+Note that the chat completion model must be your Azure OpenAI deployment name.
+
+## Passing clients manually
+
+As an alternative to setting environment variables, you can pass the `AzureOpenAI` client to Marvin's components manually:
 
 ```python
 import marvin
@@ -27,14 +31,14 @@ from marvin.client import MarvinClient
 from openai import AzureOpenAI
 
 azure_openai_client = AzureOpenAI(
-    api_key="my-api-key",
-    azure_endpoint="https://my-endpoint.openai.azure.com/",
+    api_key="your-api-key",
+    azure_endpoint="https://your-endpoint.openai.azure.com/",
     api_version="2023-12-01-preview",
 )
 
 @marvin.fn(
     client=MarvinClient(client=azure_openai_client),
-    model_kwargs={"model": "gpt-35-turbo-0613"}
+    model_kwargs={"model": "your_deployment_name"}
 )
 def list_fruits(n: int) -> list[str]:
     """generate a list of fruits"""

cookbook/azure/usage.py

@@ -1,14 +1,19 @@
-"""Usage example of Marvin with Azure OpenAI
+"""
+Usage example of Marvin with Azure OpenAI
 
-If you'll be using Azure OpenAI exclusively, you can set the following environment variables in `~/.marvin/.env`:
+If you'll be using Azure OpenAI exclusively, you can set the following env vars in your environment, `~/.marvin/.env`, or `.env`:
 ```bash
-MARVIN_USE_AZURE_OPENAI=true
-MARVIN_AZURE_OPENAI_API_KEY=...
-MARVIN_AZURE_OPENAI_API_VERSION=...
-MARVIN_AZURE_OPENAI_ENDPOINT=...
-MARVIN_AZURE_OPENAI_DEPLOYMENT_NAME=...
+
+MARVIN_PROVIDER=azure_openai
+MARVIN_AZURE_OPENAI_API_KEY=<your-api-key>
+MARVIN_AZURE_OPENAI_ENDPOINT="https://<your-endpoint>.openai.azure.com/"
+MARVIN_AZURE_OPENAI_API_VERSION=2023-12-01-preview # or latest
+
+Note that you MUST set the LLM model name to be your Azure OpenAI deployment name, e.g.
+MARVIN_CHAT_COMPLETION_MODEL=<your Azure OpenAI deployment name>
 ```
 """
+
 from enum import Enum
 
 import marvin
@@ -34,8 +39,12 @@ def list_fruits(n: int = 3) -> list[str]:
 
 
 with temporary_settings(
-    use_azure_openai=True
-):  # or set MARVIN_USE_AZURE_OPENAI=true in `~/.marvin/.env`
+    provider="azure_openai",
+    azure_openai_api_key="...",
+    azure_openai_api_version="...",
+    azure_openai_endpoint="...",
+    chat_completion_model="<your Azure OpenAI deployment name>",
+):
     fruits = list_fruits()
     location = marvin.model(Location)("windy city")
     casted_location = marvin.cast("windy city", Location)

docs/docs/configuration/settings.md

@@ -30,13 +30,15 @@ A runtime settings object is accessible via `marvin.settings` and can be used to
 ## Settings for using Azure OpenAI models
 _Some_ of Marvin's functionality is supported by Azure OpenAI services.
 
-If you're exclusively using Marvin with Azure OpenAI services, you can set the following environment variables to avoid having to pass `AzureOpenAI` client instances to Marvin's components.
+After setting up your Azure OpenAI account and deployment, set these environment variables in your environment, `~/.marvin/.env`, or `.env` file:
+
 ```bash
-MARVIN_USE_AZURE_OPENAI=true
-MARVIN_AZURE_OPENAI_API_KEY=...
-MARVIN_AZURE_OPENAI_API_VERSION=...
-MARVIN_AZURE_OPENAI_ENDPOINT=...
-MARVIN_AZURE_OPENAI_DEPLOYMENT_NAME=...
+MARVIN_PROVIDER=azure_openai
+MARVIN_AZURE_OPENAI_API_KEY=<your-api-key>
+MARVIN_AZURE_OPENAI_ENDPOINT="https://<your-endpoint>.openai.azure.com/"
+MARVIN_AZURE_OPENAI_API_VERSION=2023-12-01-preview # or latest
+
+MARVIN_CHAT_COMPLETION_MODEL=<your azure openai deployment name>
 ```
 
-To selectively use Azure OpenAI services, you can pass an `AzureOpenAI` client to Marvin's components or use `temporary_settings` like this [example](https://github.com/PrefectHQ/marvin/blob/main/cookbook/azure/README.md).
+Note that the chat completion model must be your Azure OpenAI deployment name.

src/marvin/client/openai.py

@@ -1,4 +1,3 @@
-import inspect
 from functools import partial
 from pathlib import Path
 from typing import (
@@ -32,6 +31,7 @@
     VisionRequest,
 )
 from marvin.utilities.logging import get_logger
+from marvin.utilities.openai import get_openai_client
 
 if TYPE_CHECKING:
     from openai.types import ImagesResponse
@@ -143,65 +143,14 @@ async def should_fallback(e: NotFoundError, request: ChatRequest) -> bool:
         return False
 
 
-def _get_default_client(client_type: str) -> Union[Client, AsyncClient]:
-    if getattr(settings, "use_azure_openai", False):
-        from openai import AsyncAzureOpenAI, AzureOpenAI
-
-        client_class = AsyncAzureOpenAI if client_type == "async" else AzureOpenAI
-
-        try:
-            return client_class(
-                api_key=settings.azure_openai_api_key,
-                api_version=settings.azure_openai_api_version,
-                azure_endpoint=settings.azure_openai_endpoint,
-            )
-        except AttributeError:
-            raise ValueError(
-                inspect.cleandoc(
-                    """
-                To use Azure OpenAI, please set all of the following environment variables in `~/.marvin/.env`:
-
-                ```
-                MARVIN_USE_AZURE_OPENAI=true
-                MARVIN_AZURE_OPENAI_API_KEY=...
-                MARVIN_AZURE_OPENAI_API_VERSION=...
-                MARVIN_AZURE_OPENAI_ENDPOINT=...
-                MARVIN_AZURE_OPENAI_DEPLOYMENT_NAME=...
-                ```
-                """
-                )
-            )
-
-    api_key = (
-        settings.openai.api_key.get_secret_value() if settings.openai.api_key else None
-    )
-    if not api_key:
-        raise ValueError(
-            inspect.cleandoc(
-                """
-            OpenAI API key not found! Marvin will not work properly without it.
-            
-            You can either:
-                1. Set the `MARVIN_OPENAI_API_KEY` or `OPENAI_API_KEY` environment variables
-                2. Set `marvin.settings.openai.api_key` in your code (not recommended for production)
-                
-            If you do not have an OpenAI API key, you can create one at https://platform.openai.com/api-keys.
-            """
-            )
-        )
-    if client_type not in ["sync", "async"]:
-        raise ValueError(f"Invalid client type {client_type!r}")
-
-    client_class = Client if client_type == "sync" else AsyncClient
-    return client_class(api_key=api_key, organization=settings.openai.organization)
-
-
 class MarvinClient(pydantic.BaseModel):
     model_config = pydantic.ConfigDict(
         arbitrary_types_allowed=True, protected_namespaces=()
     )
 
-    client: Client = pydantic.Field(default_factory=lambda: _get_default_client("sync"))
+    client: Client = pydantic.Field(
+        default_factory=lambda: get_openai_client(is_async=False)
+    )
 
     @classmethod
     def wrap(cls, client: Client) -> "Client":
@@ -287,7 +236,7 @@ class AsyncMarvinClient(pydantic.BaseModel):
     )
 
     client: AsyncClient = pydantic.Field(
-        default_factory=lambda: _get_default_client("async")
+        default_factory=lambda: get_openai_client(is_async=True)
     )
 
     @classmethod

src/marvin/settings.py

@@ -181,7 +181,7 @@ class OpenAISettings(MarvinSettings):
     audio: AudioSettings = Field(default_factory=AudioSettings)
     assistants: AssistantSettings = Field(default_factory=AssistantSettings)
 
-    @field_validator("api_key")
+    @field_validator("api_key", mode="before")
     def discover_api_key(cls, v):
         if v is None:
             # check global OpenAI API key
@@ -224,9 +224,20 @@ class Settings(MarvinSettings):
         protected_namespaces=(),
     )
 
+    # providers
+    provider: Literal["openai", "azure_openai"] = Field(
+        default="openai",
+        description=(
+            'The LLM provider to use. Supports "openai" and "azure_openai" at this'
+            " time."
+        ),
+    )
     openai: OpenAISettings = Field(default_factory=OpenAISettings)
+
+    # ai settings
     ai: AISettings = Field(default_factory=AISettings)
 
+    # log settings
     log_level: str = Field(
         default="INFO",
         description="The log level to use.",

src/marvin/types.py

@@ -124,13 +124,7 @@ class ResponseModel(MarvinType):
 
 
 class ChatRequest(Prompt[T]):
-    model: str = Field(
-        default_factory=lambda: (
-            settings.openai.chat.completions.model
-            if not getattr(settings, "use_azure_openai", False)
-            else settings.azure_openai_deployment_name
-        )
-    )
+    model: str = Field(default_factory=lambda: settings.openai.chat.completions.model)
     frequency_penalty: Optional[
         Annotated[float, Field(strict=True, ge=-2.0, le=2.0)]
     ] = 0

src/marvin/utilities/openai.py

@@ -1,18 +1,23 @@
-"""Module for working with OpenAI."""
+"""Utilities for working with OpenAI."""
 
 import asyncio
+import inspect
 from functools import lru_cache
-from typing import Optional
+from typing import Any, Optional, Union
 
-from openai import AsyncClient
+from openai import AsyncAzureOpenAI, AsyncClient, AzureOpenAI, Client
 
+import marvin
 
-def get_openai_client() -> AsyncClient:
+
+def get_openai_client(
+    is_async: bool = True,
+) -> Union[AsyncClient, Client, AzureOpenAI, AsyncAzureOpenAI]:
     """
-    Retrieves an OpenAI client with the given api key and organization.
+    Retrieves an OpenAI client (sync or async) based on the current configuration.
 
     Returns:
-        The OpenAI client with the given api key and organization.
+        The OpenAI client
 
     Example:
         Retrieving an OpenAI client
@@ -22,33 +27,92 @@ def get_openai_client() -> AsyncClient:
         client = get_client()
         ```
     """
-    from marvin import settings
 
-    api_key: Optional[str] = (
-        settings.openai.api_key.get_secret_value() if settings.openai.api_key else None
-    )
-    organization: Optional[str] = settings.openai.organization
+    kwargs = {}
+
+    # --- Openai
+    if marvin.settings.provider == "openai":
+        client_class = AsyncClient if is_async else Client
+
+        api_key = (
+            marvin.settings.openai.api_key.get_secret_value()
+            if marvin.settings.openai.api_key
+            else None
+        )
+
+        if not api_key:
+            raise ValueError(
+                inspect.cleandoc(
+                    """
+                    OpenAI API key not found! Marvin will not work properly without it.
+                    
+                    You can either:
+                        1. Set the `MARVIN_OPENAI_API_KEY` or `OPENAI_API_KEY` environment variables
+                        2. Set `marvin.settings.openai.api_key` in your code (not recommended for production)
+                        
+                    If you do not have an OpenAI API key, you can create one at https://platform.openai.com/api-keys.
+                    """
+                )
+            )
+
+        kwargs.update(api_key=api_key, organization=marvin.settings.openai.organization)
+
+    # --- Azure OpenAI
+    elif marvin.settings.provider == "azure_openai":
+        api_key = getattr(marvin.settings, "azure_openai_api_key", None)
+        api_version = getattr(marvin.settings, "azure_openai_api_version", None)
+        azure_endpoint = getattr(marvin.settings, "azure_openai_endpoint", None)
+
+        if any(k is None for k in [api_key, api_version, azure_endpoint]):
+            raise ValueError(
+                inspect.cleandoc(
+                    """
+                Azure OpenAI configuration is missing. Marvin will not work properly without it.
+                
+                Please make sure to set the following environment variables:
+                    - MARVIN_AZURE_OPENAI_API_KEY
+                    - MARVIN_AZURE_OPENAI_API_VERSION
+                    - MARVIN_AZURE_OPENAI_ENDPOINT
+                    
+                In addition, you must set the LLM model name to your Azure OpenAI deployment name, e.g.
+                    - MARVIN_CHAT_COMPLETION_MODEL = <your Azure OpenAI deployment name>
+                """
+                )
+            )
+        client_class = AsyncAzureOpenAI if is_async else AzureOpenAI
+        kwargs.update(
+            api_key=api_key, api_version=api_version, azure_endpoint=azure_endpoint
+        )
+
+    # --- N/A
+    else:
+        raise ValueError(f"Unknown provider {marvin.settings.provider}")
+
     return _get_client_memoized(
-        api_key=api_key, organization=organization, loop=asyncio.get_event_loop()
+        cls=client_class,
+        loop=asyncio.get_event_loop(),
+        kwargs_items=tuple(kwargs.items()),
     )
 
 
 @lru_cache
 def _get_client_memoized(
-    api_key: Optional[str],
-    organization: Optional[str],
+    cls: type,
     loop: Optional[asyncio.AbstractEventLoop] = None,
-) -> AsyncClient:
+    kwargs_items: tuple[tuple[str, Any]] = None,
+) -> Union[Client, AsyncClient]:
     """
     This function is memoized to ensure that only one instance of the client is
-    created for a given api key / organization / loop tuple.
+    created for a given set of configuration parameters
+
+    It can return either a sync or an async client.
 
     The `loop` is an important key to ensure that the client is not re-used
     across multiple event loops (which can happen when using the `run_sync`
     function). Attempting to re-use the client across multiple event loops
     can result in a `RuntimeError: Event loop is closed` error or infinite hangs.
+
+    kwargs_items is a tuple of dict items to get around the fact that
+    memoization requires hashable arguments.
     """
-    return AsyncClient(
-        api_key=api_key,
-        organization=organization,
-    )
+    return cls(**dict(kwargs_items))