From a4b63728131ca38dcee95eeac51fedbb2ae0e001 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Thu, 13 Mar 2025 21:29:19 -0700 Subject: [PATCH] Use SecretStr type for api key (#5939) To prevent accidental export of API keys --- .../models/anthropic/config/__init__.py | 4 ++-- .../models/openai/config/__init__.py | 4 ++-- .../models/test_anthropic_model_client.py | 19 +++++++++++++++++++ .../tests/models/test_openai_model_client.py | 14 ++++++++++++++ 4 files changed, 37 insertions(+), 4 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py index 58ff24848..0324f66b1 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py @@ -1,7 +1,7 @@ from typing import Any, Dict, List, Literal, Optional, Union from autogen_core.models import ModelCapabilities, ModelInfo # type: ignore -from pydantic import BaseModel +from pydantic import BaseModel, SecretStr from typing_extensions import TypedDict @@ -49,7 +49,7 @@ class CreateArgumentsConfigModel(BaseModel): class BaseAnthropicClientConfigurationConfigModel(CreateArgumentsConfigModel): - api_key: str | None = None + api_key: SecretStr | None = None base_url: str | None = None model_capabilities: ModelCapabilities | None = None # type: ignore model_info: ModelInfo | None = None diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py index 62c25c800..22dff2450 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py @@ -2,7 +2,7 @@ from typing import Awaitable, Callable, Dict, List, Literal, Optional, Union from autogen_core import ComponentModel from autogen_core.models import ModelCapabilities, ModelInfo # type: ignore -from pydantic import BaseModel +from pydantic import BaseModel, SecretStr from typing_extensions import Required, TypedDict @@ -77,7 +77,7 @@ class CreateArgumentsConfigModel(BaseModel): class BaseOpenAIClientConfigurationConfigModel(CreateArgumentsConfigModel): model: str - api_key: str | None = None + api_key: SecretStr | None = None timeout: float | None = None max_retries: int | None = None model_capabilities: ModelCapabilities | None = None # type: ignore diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 2ab8b9f40..4f4135b4b 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -28,6 +28,25 @@ def _add_numbers(a: int, b: int) -> int: return a + b +@pytest.mark.asyncio +async def test_anthropic_serialization_api_key() -> None: + client = AnthropicChatCompletionClient( + model="claude-3-haiku-20240307", # Use haiku for faster/cheaper testing + api_key="sk-password", + temperature=0.0, # Added temperature param to test + stop_sequences=["STOP"], # Added stop sequence + ) + assert client + config = client.dump_component() + assert config + assert "sk-password" not in str(config) + serialized_config = config.model_dump_json() + assert serialized_config + assert "sk-password" not in serialized_config + client2 = AnthropicChatCompletionClient.load_component(config) + assert client2 + + @pytest.mark.asyncio async def test_anthropic_basic_completion(caplog: pytest.LogCaptureFixture) -> None: """Test basic message completion with Claude.""" diff --git a/python/packages/autogen-ext/tests/models/test_openai_model_client.py b/python/packages/autogen-ext/tests/models/test_openai_model_client.py index 5cef46edb..e63bca30e 100644 --- a/python/packages/autogen-ext/tests/models/test_openai_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_openai_model_client.py @@ -189,6 +189,20 @@ async def test_openai_chat_completion_client_with_gemini_model() -> None: assert client +@pytest.mark.asyncio +async def test_openai_chat_completion_client_serialization() -> None: + client = OpenAIChatCompletionClient(model="gpt-4o", api_key="sk-password") + assert client + config = client.dump_component() + assert config + assert "sk-password" not in str(config) + serialized_config = config.model_dump_json() + assert serialized_config + assert "sk-password" not in serialized_config + client2 = OpenAIChatCompletionClient.load_component(config) + assert client2 + + @pytest.mark.asyncio async def test_openai_chat_completion_client_raise_on_unknown_model() -> None: with pytest.raises(ValueError, match="model_info is required"):