Use SecretStr type for api key (#5939)

To prevent accidental export of API keys
This commit is contained in:
Eric Zhu 2025-03-13 21:29:19 -07:00 committed by GitHub
parent 84c622a4cc
commit a4b6372813
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 37 additions and 4 deletions

View File

@ -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

View File

@ -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

View File

@ -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."""

View File

@ -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"):