From 63bcfb872b42ee2612c812bf2ed448a1a00cf3e1 Mon Sep 17 00:00:00 2001 From: Skylar Payne Date: Wed, 9 Oct 2024 15:31:37 -0700 Subject: [PATCH] Add a util for loading Mistral API key --- .../learn/calls/basic_usage/mistral/official_sdk_call.py | 4 ++-- .../calls/custom_client/mistral/base_message_param.py | 6 ++---- examples/learn/calls/custom_client/mistral/messages.py | 6 ++---- examples/learn/calls/custom_client/mistral/shorthand.py | 6 ++---- .../learn/calls/custom_client/mistral/string_template.py | 6 ++---- .../response_models/basic_usage/mistral/official_sdk.py | 7 +++---- mirascope/core/mistral/__init__.py | 2 ++ mirascope/core/mistral/_utils/_load_api_key.py | 6 ++++++ mirascope/core/mistral/_utils/_setup_call.py | 4 ++-- 9 files changed, 23 insertions(+), 24 deletions(-) create mode 100644 mirascope/core/mistral/_utils/_load_api_key.py diff --git a/examples/learn/calls/basic_usage/mistral/official_sdk_call.py b/examples/learn/calls/basic_usage/mistral/official_sdk_call.py index b14fb8227..bf6495a17 100644 --- a/examples/learn/calls/basic_usage/mistral/official_sdk_call.py +++ b/examples/learn/calls/basic_usage/mistral/official_sdk_call.py @@ -1,7 +1,7 @@ +from mirascope.core import mistral from mistralai import Mistral -import os -client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")) +client = Mistral(api_key=mistral.load_api_key()) def recommend_book(genre: str) -> str: diff --git a/examples/learn/calls/custom_client/mistral/base_message_param.py b/examples/learn/calls/custom_client/mistral/base_message_param.py index 01fa5181c..8d9c5b5d0 100644 --- a/examples/learn/calls/custom_client/mistral/base_message_param.py +++ b/examples/learn/calls/custom_client/mistral/base_message_param.py @@ -1,12 +1,10 @@ -import os - from mirascope.core import BaseMessageParam, mistral from mistralai import Mistral @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) def recommend_book(genre: str) -> list[BaseMessageParam]: return [BaseMessageParam(role="user", content=f"Recommend a {genre} book")] @@ -14,7 +12,7 @@ def recommend_book(genre: str) -> list[BaseMessageParam]: @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) async def recommend_book_async(genre: str) -> list[BaseMessageParam]: return [BaseMessageParam(role="user", content=f"Recommend a {genre} book")] diff --git a/examples/learn/calls/custom_client/mistral/messages.py b/examples/learn/calls/custom_client/mistral/messages.py index 569ca229a..818ed277d 100644 --- a/examples/learn/calls/custom_client/mistral/messages.py +++ b/examples/learn/calls/custom_client/mistral/messages.py @@ -1,12 +1,10 @@ -import os - from mirascope.core import Messages, mistral from mistralai import Mistral @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) def recommend_book(genre: str) -> Messages.Type: return Messages.User(f"Recommend a {genre} book") @@ -14,7 +12,7 @@ def recommend_book(genre: str) -> Messages.Type: @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) async def recommend_book_async(genre: str) -> Messages.Type: return Messages.User(f"Recommend a {genre} book") diff --git a/examples/learn/calls/custom_client/mistral/shorthand.py b/examples/learn/calls/custom_client/mistral/shorthand.py index 38571b32b..f61346ab2 100644 --- a/examples/learn/calls/custom_client/mistral/shorthand.py +++ b/examples/learn/calls/custom_client/mistral/shorthand.py @@ -1,12 +1,10 @@ -import os - from mirascope.core import mistral from mistralai import Mistral @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) def recommend_book(genre: str) -> str: return f"Recommend a {genre} book" @@ -14,7 +12,7 @@ def recommend_book(genre: str) -> str: @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) async def recommend_book_async(genre: str) -> str: return f"Recommend a {genre} book" diff --git a/examples/learn/calls/custom_client/mistral/string_template.py b/examples/learn/calls/custom_client/mistral/string_template.py index 8743a1a60..0fcb2904b 100644 --- a/examples/learn/calls/custom_client/mistral/string_template.py +++ b/examples/learn/calls/custom_client/mistral/string_template.py @@ -1,12 +1,10 @@ -import os - from mirascope.core import mistral, prompt_template from mistralai import Mistral @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) @prompt_template("Recommend a {genre} book") def recommend_book(genre: str): ... @@ -14,7 +12,7 @@ def recommend_book(genre: str): ... @mistral.call( "mistral-large-latest", - client=Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")), + client=Mistral(api_key=mistral.load_api_key()), ) @prompt_template("Recommend a {genre} book") async def recommend_book_async(genre: str): ... diff --git a/examples/learn/response_models/basic_usage/mistral/official_sdk.py b/examples/learn/response_models/basic_usage/mistral/official_sdk.py index 7131d46dd..102c95139 100644 --- a/examples/learn/response_models/basic_usage/mistral/official_sdk.py +++ b/examples/learn/response_models/basic_usage/mistral/official_sdk.py @@ -1,10 +1,9 @@ -import os - -from mistralai.client import Mistral +from mirascope.core import mistral +from mistralai import Mistral from mistralai.models import ToolChoice from pydantic import BaseModel -client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")) +client = Mistral(api_key=mistral.load_api_key()) class Book(BaseModel): diff --git a/mirascope/core/mistral/__init__.py b/mirascope/core/mistral/__init__.py index 39ce1b4ed..b08cd7293 100644 --- a/mirascope/core/mistral/__init__.py +++ b/mirascope/core/mistral/__init__.py @@ -12,6 +12,7 @@ from ..base import BaseMessageParam from ._call import mistral_call from ._call import mistral_call as call +from ._utils._load_api_key import load_api_key from .call_params import MistralCallParams from .call_response import MistralCallResponse from .call_response_chunk import MistralCallResponseChunk @@ -25,6 +26,7 @@ __all__ = [ "call", + "load_api_key", "MistralDynamicConfig", "MistralCallParams", "MistralCallResponse", diff --git a/mirascope/core/mistral/_utils/_load_api_key.py b/mirascope/core/mistral/_utils/_load_api_key.py new file mode 100644 index 000000000..f06f3bcc4 --- /dev/null +++ b/mirascope/core/mistral/_utils/_load_api_key.py @@ -0,0 +1,6 @@ +import os + + +def load_api_key() -> str: + """Load the API key from the standard environment variable.""" + return os.environ.get("MISTRAL_API_KEY", "") diff --git a/mirascope/core/mistral/_utils/_setup_call.py b/mirascope/core/mistral/_utils/_setup_call.py index c8bb09689..b8d96cf84 100644 --- a/mirascope/core/mistral/_utils/_setup_call.py +++ b/mirascope/core/mistral/_utils/_setup_call.py @@ -1,6 +1,5 @@ """This module contains the setup_call function for Mistral tools.""" -import os from collections.abc import ( Awaitable, Callable, @@ -19,6 +18,7 @@ UserMessage, ) +from ... import mistral from ...base import BaseTool, _utils from ...base._utils import AsyncCreateFn, CreateFn, get_async_create_fn, get_create_fn from ...base._utils._protocols import fn_is_async @@ -114,7 +114,7 @@ def setup_call( call_kwargs |= {"model": model, "messages": messages} if client is None: - client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY", "")) + client = Mistral(api_key=mistral.load_api_key()) if fn_is_async(fn): create_or_stream = get_async_create_fn( client.chat.complete_async, client.chat.stream_async