From c4cb6aa8d99ad94afd2f18c0ba69a40cb09c0080 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Fri, 8 Aug 2025 15:54:45 +0200 Subject: [PATCH] fix: prevent telemetry from leaking sensitive info MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prevent sensitive information from being logged in telemetry output by assigning SecretStr type to sensitive fields. API keys, password from KV store are now covered. All providers have been converted. Signed-off-by: Sébastien Han --- docs/docs/providers/files/remote_s3.mdx | 2 +- .../providers/inference/remote_anthropic.mdx | 2 +- .../providers/inference/remote_bedrock.mdx | 4 ++-- .../providers/inference/remote_gemini.mdx | 2 +- docs/docs/providers/inference/remote_groq.mdx | 2 +- .../inference/remote_llama-openai-compat.mdx | 2 +- .../providers/inference/remote_openai.mdx | 2 +- .../providers/inference/remote_runpod.mdx | 2 +- docs/docs/providers/inference/remote_vllm.mdx | 2 +- docs/docs/providers/safety/remote_bedrock.mdx | 4 ++-- .../providers/scoring/inline_braintrust.mdx | 2 +- .../tool_runtime/remote_bing-search.mdx | 2 +- .../tool_runtime/remote_brave-search.mdx | 2 +- .../tool_runtime/remote_tavily-search.mdx | 2 +- .../providers/vector_io/remote_pgvector.mdx | 2 +- .../inline/scoring/braintrust/__init__.py | 4 ++-- .../inline/scoring/braintrust/braintrust.py | 6 ++--- .../inline/scoring/braintrust/config.py | 4 ++-- .../meta_reference/console_span_processor.py | 4 +++- .../providers/remote/files/s3/config.py | 4 ++-- .../providers/remote/files/s3/files.py | 2 +- .../remote/inference/anthropic/anthropic.py | 3 ++- .../remote/inference/anthropic/config.py | 6 ++--- .../remote/inference/gemini/config.py | 6 ++--- .../remote/inference/gemini/gemini.py | 3 ++- .../providers/remote/inference/groq/config.py | 6 ++--- .../inference/llama_openai_compat/config.py | 6 ++--- .../remote/inference/openai/config.py | 6 ++--- .../remote/inference/runpod/config.py | 4 ++-- .../remote/inference/runpod/runpod.py | 5 ++++- .../remote/inference/vertexai/vertexai.py | 3 ++- .../providers/remote/inference/vllm/config.py | 6 ++--- .../providers/remote/inference/vllm/vllm.py | 2 +- .../tool_runtime/bing_search/bing_search.py | 2 +- .../remote/tool_runtime/bing_search/config.py | 4 ++-- .../tool_runtime/brave_search/brave_search.py | 2 +- .../tool_runtime/brave_search/config.py | 4 ++-- .../tool_runtime/tavily_search/config.py | 4 ++-- .../tavily_search/tavily_search.py | 2 +- .../wolfram_alpha/wolfram_alpha.py | 2 +- .../remote/vector_io/pgvector/config.py | 4 ++-- .../remote/vector_io/pgvector/pgvector.py | 2 +- llama_stack/providers/utils/bedrock/client.py | 4 ++-- llama_stack/providers/utils/bedrock/config.py | 10 ++++----- .../utils/inference/litellm_openai_mixin.py | 13 ++++++----- .../providers/utils/inference/openai_mixin.py | 5 +++-- llama_stack/providers/utils/kvstore/config.py | 6 ++--- .../utils/kvstore/mongodb/mongodb.py | 2 +- .../utils/kvstore/postgres/postgres.py | 2 +- .../providers/utils/sqlstore/sqlstore.py | 6 ++--- .../test_inference_client_caching.py | 6 ++--- .../inference/test_litellm_openai_mixin.py | 22 +++++++++---------- .../inference/test_openai_base_url_config.py | 14 +++++++----- 53 files changed, 121 insertions(+), 109 deletions(-) diff --git a/docs/docs/providers/files/remote_s3.mdx b/docs/docs/providers/files/remote_s3.mdx index 353cedbfb..a065da5f8 100644 --- a/docs/docs/providers/files/remote_s3.mdx +++ b/docs/docs/providers/files/remote_s3.mdx @@ -17,7 +17,7 @@ AWS S3-based file storage provider for scalable cloud file management with metad | `bucket_name` | `` | No | | S3 bucket name to store files | | `region` | `` | No | us-east-1 | AWS region where the bucket is located | | `aws_access_key_id` | `str \| None` | No | | AWS access key ID (optional if using IAM roles) | -| `aws_secret_access_key` | `str \| None` | No | | AWS secret access key (optional if using IAM roles) | +| `aws_secret_access_key` | `pydantic.types.SecretStr \| None` | No | | AWS secret access key (optional if using IAM roles) | | `endpoint_url` | `str \| None` | No | | Custom S3 endpoint URL (for MinIO, LocalStack, etc.) | | `auto_create_bucket` | `` | No | False | Automatically create the S3 bucket if it doesn't exist | | `metadata_store` | `utils.sqlstore.sqlstore.SqliteSqlStoreConfig \| utils.sqlstore.sqlstore.PostgresSqlStoreConfig` | No | sqlite | SQL store configuration for file metadata | diff --git a/docs/docs/providers/inference/remote_anthropic.mdx b/docs/docs/providers/inference/remote_anthropic.mdx index 6bd636c92..c8115c7a8 100644 --- a/docs/docs/providers/inference/remote_anthropic.mdx +++ b/docs/docs/providers/inference/remote_anthropic.mdx @@ -14,7 +14,7 @@ Anthropic inference provider for accessing Claude models and Anthropic's AI serv | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | API key for Anthropic models | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | API key for Anthropic models | ## Sample Configuration diff --git a/docs/docs/providers/inference/remote_bedrock.mdx b/docs/docs/providers/inference/remote_bedrock.mdx index 04c2154a9..3a4c296cf 100644 --- a/docs/docs/providers/inference/remote_bedrock.mdx +++ b/docs/docs/providers/inference/remote_bedrock.mdx @@ -15,8 +15,8 @@ AWS Bedrock inference provider for accessing various AI models through AWS's man | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `aws_access_key_id` | `str \| None` | No | | The AWS access key to use. Default use environment variable: AWS_ACCESS_KEY_ID | -| `aws_secret_access_key` | `str \| None` | No | | The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY | -| `aws_session_token` | `str \| None` | No | | The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN | +| `aws_secret_access_key` | `pydantic.types.SecretStr \| None` | No | | The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY | +| `aws_session_token` | `pydantic.types.SecretStr \| None` | No | | The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN | | `region_name` | `str \| None` | No | | The default AWS Region to use, for example, us-west-1 or us-west-2.Default use environment variable: AWS_DEFAULT_REGION | | `profile_name` | `str \| None` | No | | The profile name that contains credentials to use.Default use environment variable: AWS_PROFILE | | `total_max_attempts` | `int \| None` | No | | An integer representing the maximum number of attempts that will be made for a single request, including the initial attempt. Default use environment variable: AWS_MAX_ATTEMPTS | diff --git a/docs/docs/providers/inference/remote_gemini.mdx b/docs/docs/providers/inference/remote_gemini.mdx index 0505c69da..ce7c797c4 100644 --- a/docs/docs/providers/inference/remote_gemini.mdx +++ b/docs/docs/providers/inference/remote_gemini.mdx @@ -14,7 +14,7 @@ Google Gemini inference provider for accessing Gemini models and Google's AI ser | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | API key for Gemini models | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | API key for Gemini models | ## Sample Configuration diff --git a/docs/docs/providers/inference/remote_groq.mdx b/docs/docs/providers/inference/remote_groq.mdx index 1797035c1..c2745ed98 100644 --- a/docs/docs/providers/inference/remote_groq.mdx +++ b/docs/docs/providers/inference/remote_groq.mdx @@ -14,7 +14,7 @@ Groq inference provider for ultra-fast inference using Groq's LPU technology. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | The Groq API key | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | The Groq API key | | `url` | `` | No | https://api.groq.com | The URL for the Groq AI server | ## Sample Configuration diff --git a/docs/docs/providers/inference/remote_llama-openai-compat.mdx b/docs/docs/providers/inference/remote_llama-openai-compat.mdx index cb624ad87..b7b0204e6 100644 --- a/docs/docs/providers/inference/remote_llama-openai-compat.mdx +++ b/docs/docs/providers/inference/remote_llama-openai-compat.mdx @@ -14,7 +14,7 @@ Llama OpenAI-compatible provider for using Llama models with OpenAI API format. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | The Llama API key | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | The Llama API key | | `openai_compat_api_base` | `` | No | https://api.llama.com/compat/v1/ | The URL for the Llama API server | ## Sample Configuration diff --git a/docs/docs/providers/inference/remote_openai.mdx b/docs/docs/providers/inference/remote_openai.mdx index 56ca94233..2c7650bb5 100644 --- a/docs/docs/providers/inference/remote_openai.mdx +++ b/docs/docs/providers/inference/remote_openai.mdx @@ -14,7 +14,7 @@ OpenAI inference provider for accessing GPT models and other OpenAI services. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | API key for OpenAI models | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | API key for OpenAI models | | `base_url` | `` | No | https://api.openai.com/v1 | Base URL for OpenAI API | ## Sample Configuration diff --git a/docs/docs/providers/inference/remote_runpod.mdx b/docs/docs/providers/inference/remote_runpod.mdx index 2e8847dc5..5d3655474 100644 --- a/docs/docs/providers/inference/remote_runpod.mdx +++ b/docs/docs/providers/inference/remote_runpod.mdx @@ -15,7 +15,7 @@ RunPod inference provider for running models on RunPod's cloud GPU platform. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `url` | `str \| None` | No | | The URL for the Runpod model serving endpoint | -| `api_token` | `str \| None` | No | | The API token | +| `api_token` | `pydantic.types.SecretStr \| None` | No | | The API token | ## Sample Configuration diff --git a/docs/docs/providers/inference/remote_vllm.mdx b/docs/docs/providers/inference/remote_vllm.mdx index 77b8e1355..1faf131b9 100644 --- a/docs/docs/providers/inference/remote_vllm.mdx +++ b/docs/docs/providers/inference/remote_vllm.mdx @@ -16,7 +16,7 @@ Remote vLLM inference provider for connecting to vLLM servers. |-------|------|----------|---------|-------------| | `url` | `str \| None` | No | | The URL for the vLLM model serving endpoint | | `max_tokens` | `` | No | 4096 | Maximum number of tokens to generate. | -| `api_token` | `str \| None` | No | fake | The API token | +| `api_token` | `pydantic.types.SecretStr \| None` | No | ********** | The API token | | `tls_verify` | `bool \| str` | No | True | Whether to verify TLS certificates. Can be a boolean or a path to a CA certificate file. | | `refresh_models` | `` | No | False | Whether to refresh models periodically | diff --git a/docs/docs/providers/safety/remote_bedrock.mdx b/docs/docs/providers/safety/remote_bedrock.mdx index 5461d7cdc..85bc662e8 100644 --- a/docs/docs/providers/safety/remote_bedrock.mdx +++ b/docs/docs/providers/safety/remote_bedrock.mdx @@ -15,8 +15,8 @@ AWS Bedrock safety provider for content moderation using AWS's safety services. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| | `aws_access_key_id` | `str \| None` | No | | The AWS access key to use. Default use environment variable: AWS_ACCESS_KEY_ID | -| `aws_secret_access_key` | `str \| None` | No | | The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY | -| `aws_session_token` | `str \| None` | No | | The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN | +| `aws_secret_access_key` | `pydantic.types.SecretStr \| None` | No | | The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY | +| `aws_session_token` | `pydantic.types.SecretStr \| None` | No | | The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN | | `region_name` | `str \| None` | No | | The default AWS Region to use, for example, us-west-1 or us-west-2.Default use environment variable: AWS_DEFAULT_REGION | | `profile_name` | `str \| None` | No | | The profile name that contains credentials to use.Default use environment variable: AWS_PROFILE | | `total_max_attempts` | `int \| None` | No | | An integer representing the maximum number of attempts that will be made for a single request, including the initial attempt. Default use environment variable: AWS_MAX_ATTEMPTS | diff --git a/docs/docs/providers/scoring/inline_braintrust.mdx b/docs/docs/providers/scoring/inline_braintrust.mdx index d12f9de25..67de0667a 100644 --- a/docs/docs/providers/scoring/inline_braintrust.mdx +++ b/docs/docs/providers/scoring/inline_braintrust.mdx @@ -14,7 +14,7 @@ Braintrust scoring provider for evaluation and scoring using the Braintrust plat | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `openai_api_key` | `str \| None` | No | | The OpenAI API Key | +| `openai_api_key` | `pydantic.types.SecretStr \| None` | No | | The OpenAI API Key | ## Sample Configuration diff --git a/docs/docs/providers/tool_runtime/remote_bing-search.mdx b/docs/docs/providers/tool_runtime/remote_bing-search.mdx index ec06bc20f..3085ab9ec 100644 --- a/docs/docs/providers/tool_runtime/remote_bing-search.mdx +++ b/docs/docs/providers/tool_runtime/remote_bing-search.mdx @@ -14,7 +14,7 @@ Bing Search tool for web search capabilities using Microsoft's search engine. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | | | `top_k` | `` | No | 3 | | ## Sample Configuration diff --git a/docs/docs/providers/tool_runtime/remote_brave-search.mdx b/docs/docs/providers/tool_runtime/remote_brave-search.mdx index 3aeed67d5..c71ac3d7a 100644 --- a/docs/docs/providers/tool_runtime/remote_brave-search.mdx +++ b/docs/docs/providers/tool_runtime/remote_brave-search.mdx @@ -14,7 +14,7 @@ Brave Search tool for web search capabilities with privacy-focused results. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | The Brave Search API Key | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | The Brave Search API Key | | `max_results` | `` | No | 3 | The maximum number of results to return | ## Sample Configuration diff --git a/docs/docs/providers/tool_runtime/remote_tavily-search.mdx b/docs/docs/providers/tool_runtime/remote_tavily-search.mdx index fdca31bbe..0fe263a5b 100644 --- a/docs/docs/providers/tool_runtime/remote_tavily-search.mdx +++ b/docs/docs/providers/tool_runtime/remote_tavily-search.mdx @@ -14,7 +14,7 @@ Tavily Search tool for AI-optimized web search with structured results. | Field | Type | Required | Default | Description | |-------|------|----------|---------|-------------| -| `api_key` | `str \| None` | No | | The Tavily Search API Key | +| `api_key` | `pydantic.types.SecretStr \| None` | No | | The Tavily Search API Key | | `max_results` | `` | No | 3 | The maximum number of results to return | ## Sample Configuration diff --git a/docs/docs/providers/vector_io/remote_pgvector.mdx b/docs/docs/providers/vector_io/remote_pgvector.mdx index d21810c68..cf0c28f12 100644 --- a/docs/docs/providers/vector_io/remote_pgvector.mdx +++ b/docs/docs/providers/vector_io/remote_pgvector.mdx @@ -217,7 +217,7 @@ See [PGVector's documentation](https://github.com/pgvector/pgvector) for more de | `port` | `int \| None` | No | 5432 | | | `db` | `str \| None` | No | postgres | | | `user` | `str \| None` | No | postgres | | -| `password` | `str \| None` | No | mysecretpassword | | +| `password` | `pydantic.types.SecretStr \| None` | No | mysecretpassword | | | `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig, annotation=NoneType, required=False, default='sqlite', discriminator='type'` | No | | Config for KV store backend (SQLite only for now) | ## Sample Configuration diff --git a/llama_stack/providers/inline/scoring/braintrust/__init__.py b/llama_stack/providers/inline/scoring/braintrust/__init__.py index 3b492ae3f..2f3dce966 100644 --- a/llama_stack/providers/inline/scoring/braintrust/__init__.py +++ b/llama_stack/providers/inline/scoring/braintrust/__init__.py @@ -5,7 +5,7 @@ # the root directory of this source tree. from typing import Any -from pydantic import BaseModel +from pydantic import BaseModel, SecretStr from llama_stack.core.datatypes import Api @@ -13,7 +13,7 @@ from .config import BraintrustScoringConfig class BraintrustProviderDataValidator(BaseModel): - openai_api_key: str + openai_api_key: SecretStr async def get_provider_impl( diff --git a/llama_stack/providers/inline/scoring/braintrust/braintrust.py b/llama_stack/providers/inline/scoring/braintrust/braintrust.py index 14810f706..7f2b1e205 100644 --- a/llama_stack/providers/inline/scoring/braintrust/braintrust.py +++ b/llama_stack/providers/inline/scoring/braintrust/braintrust.py @@ -17,7 +17,7 @@ from autoevals.ragas import ( ContextRelevancy, Faithfulness, ) -from pydantic import BaseModel +from pydantic import BaseModel, SecretStr from llama_stack.apis.datasetio import DatasetIO from llama_stack.apis.datasets import Datasets @@ -152,9 +152,9 @@ class BraintrustScoringImpl( raise ValueError( 'Pass OpenAI API Key in the header X-LlamaStack-Provider-Data as { "openai_api_key": }' ) - self.config.openai_api_key = provider_data.openai_api_key + self.config.openai_api_key = SecretStr(provider_data.openai_api_key) - os.environ["OPENAI_API_KEY"] = self.config.openai_api_key + os.environ["OPENAI_API_KEY"] = self.config.openai_api_key.get_secret_value() async def score_batch( self, diff --git a/llama_stack/providers/inline/scoring/braintrust/config.py b/llama_stack/providers/inline/scoring/braintrust/config.py index 057f0ba5d..088674afd 100644 --- a/llama_stack/providers/inline/scoring/braintrust/config.py +++ b/llama_stack/providers/inline/scoring/braintrust/config.py @@ -5,11 +5,11 @@ # the root directory of this source tree. from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr class BraintrustScoringConfig(BaseModel): - openai_api_key: str | None = Field( + openai_api_key: SecretStr | None = Field( default=None, description="The OpenAI API Key", ) diff --git a/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py b/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py index 78e49af94..8beb3a841 100644 --- a/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py +++ b/llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py @@ -64,7 +64,9 @@ class ConsoleSpanProcessor(SpanProcessor): for key, value in event.attributes.items(): if key.startswith("__") or key in ["message", "severity"]: continue - logger.info(f"[dim]{key}[/dim]: {value}") + + str_value = str(value) + logger.info(f"[dim]{key}[/dim]: {str_value}") def shutdown(self) -> None: """Shutdown the processor.""" diff --git a/llama_stack/providers/remote/files/s3/config.py b/llama_stack/providers/remote/files/s3/config.py index da20d8668..5cddd2b38 100644 --- a/llama_stack/providers/remote/files/s3/config.py +++ b/llama_stack/providers/remote/files/s3/config.py @@ -6,7 +6,7 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig @@ -17,7 +17,7 @@ class S3FilesImplConfig(BaseModel): bucket_name: str = Field(description="S3 bucket name to store files") region: str = Field(default="us-east-1", description="AWS region where the bucket is located") aws_access_key_id: str | None = Field(default=None, description="AWS access key ID (optional if using IAM roles)") - aws_secret_access_key: str | None = Field( + aws_secret_access_key: SecretStr | None = Field( default=None, description="AWS secret access key (optional if using IAM roles)" ) endpoint_url: str | None = Field(default=None, description="Custom S3 endpoint URL (for MinIO, LocalStack, etc.)") diff --git a/llama_stack/providers/remote/files/s3/files.py b/llama_stack/providers/remote/files/s3/files.py index 8ea96af9e..61b82105e 100644 --- a/llama_stack/providers/remote/files/s3/files.py +++ b/llama_stack/providers/remote/files/s3/files.py @@ -46,7 +46,7 @@ def _create_s3_client(config: S3FilesImplConfig) -> boto3.client: s3_config.update( { "aws_access_key_id": config.aws_access_key_id, - "aws_secret_access_key": config.aws_secret_access_key, + "aws_secret_access_key": config.aws_secret_access_key.get_secret_value(), } ) diff --git a/llama_stack/providers/remote/inference/anthropic/anthropic.py b/llama_stack/providers/remote/inference/anthropic/anthropic.py index cdde4a411..41f65ade7 100644 --- a/llama_stack/providers/remote/inference/anthropic/anthropic.py +++ b/llama_stack/providers/remote/inference/anthropic/anthropic.py @@ -4,6 +4,7 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. + from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin @@ -27,7 +28,7 @@ class AnthropicInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin): LiteLLMOpenAIMixin.__init__( self, litellm_provider_name="anthropic", - api_key_from_config=config.api_key, + api_key_from_config=config.api_key.get_secret_value() if config.api_key else None, provider_data_api_key_field="anthropic_api_key", ) self.config = config diff --git a/llama_stack/providers/remote/inference/anthropic/config.py b/llama_stack/providers/remote/inference/anthropic/config.py index a74b97a9e..ae3034219 100644 --- a/llama_stack/providers/remote/inference/anthropic/config.py +++ b/llama_stack/providers/remote/inference/anthropic/config.py @@ -6,13 +6,13 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.schema_utils import json_schema_type class AnthropicProviderDataValidator(BaseModel): - anthropic_api_key: str | None = Field( + anthropic_api_key: SecretStr | None = Field( default=None, description="API key for Anthropic models", ) @@ -20,7 +20,7 @@ class AnthropicProviderDataValidator(BaseModel): @json_schema_type class AnthropicConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( default=None, description="API key for Anthropic models", ) diff --git a/llama_stack/providers/remote/inference/gemini/config.py b/llama_stack/providers/remote/inference/gemini/config.py index c897777f7..27897965c 100644 --- a/llama_stack/providers/remote/inference/gemini/config.py +++ b/llama_stack/providers/remote/inference/gemini/config.py @@ -6,13 +6,13 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.schema_utils import json_schema_type class GeminiProviderDataValidator(BaseModel): - gemini_api_key: str | None = Field( + gemini_api_key: SecretStr | None = Field( default=None, description="API key for Gemini models", ) @@ -20,7 +20,7 @@ class GeminiProviderDataValidator(BaseModel): @json_schema_type class GeminiConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( default=None, description="API key for Gemini models", ) diff --git a/llama_stack/providers/remote/inference/gemini/gemini.py b/llama_stack/providers/remote/inference/gemini/gemini.py index 30ceedff0..b02d0d611 100644 --- a/llama_stack/providers/remote/inference/gemini/gemini.py +++ b/llama_stack/providers/remote/inference/gemini/gemini.py @@ -4,6 +4,7 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. + from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin @@ -19,7 +20,7 @@ class GeminiInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin): LiteLLMOpenAIMixin.__init__( self, litellm_provider_name="gemini", - api_key_from_config=config.api_key, + api_key_from_config=config.api_key.get_secret_value() if config.api_key else None, provider_data_api_key_field="gemini_api_key", ) self.config = config diff --git a/llama_stack/providers/remote/inference/groq/config.py b/llama_stack/providers/remote/inference/groq/config.py index 67e9fa358..1011e8da9 100644 --- a/llama_stack/providers/remote/inference/groq/config.py +++ b/llama_stack/providers/remote/inference/groq/config.py @@ -6,13 +6,13 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.schema_utils import json_schema_type class GroqProviderDataValidator(BaseModel): - groq_api_key: str | None = Field( + groq_api_key: SecretStr | None = Field( default=None, description="API key for Groq models", ) @@ -20,7 +20,7 @@ class GroqProviderDataValidator(BaseModel): @json_schema_type class GroqConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( # The Groq client library loads the GROQ_API_KEY environment variable by default default=None, description="The Groq API key", diff --git a/llama_stack/providers/remote/inference/llama_openai_compat/config.py b/llama_stack/providers/remote/inference/llama_openai_compat/config.py index 57bc7240d..064cc860d 100644 --- a/llama_stack/providers/remote/inference/llama_openai_compat/config.py +++ b/llama_stack/providers/remote/inference/llama_openai_compat/config.py @@ -6,13 +6,13 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.schema_utils import json_schema_type class LlamaProviderDataValidator(BaseModel): - llama_api_key: str | None = Field( + llama_api_key: SecretStr | None = Field( default=None, description="API key for api.llama models", ) @@ -20,7 +20,7 @@ class LlamaProviderDataValidator(BaseModel): @json_schema_type class LlamaCompatConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( default=None, description="The Llama API key", ) diff --git a/llama_stack/providers/remote/inference/openai/config.py b/llama_stack/providers/remote/inference/openai/config.py index ad25cdfa5..7de7bcc32 100644 --- a/llama_stack/providers/remote/inference/openai/config.py +++ b/llama_stack/providers/remote/inference/openai/config.py @@ -6,13 +6,13 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.schema_utils import json_schema_type class OpenAIProviderDataValidator(BaseModel): - openai_api_key: str | None = Field( + openai_api_key: SecretStr | None = Field( default=None, description="API key for OpenAI models", ) @@ -20,7 +20,7 @@ class OpenAIProviderDataValidator(BaseModel): @json_schema_type class OpenAIConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( default=None, description="API key for OpenAI models", ) diff --git a/llama_stack/providers/remote/inference/runpod/config.py b/llama_stack/providers/remote/inference/runpod/config.py index 7bc9e8485..60086acef 100644 --- a/llama_stack/providers/remote/inference/runpod/config.py +++ b/llama_stack/providers/remote/inference/runpod/config.py @@ -6,7 +6,7 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.schema_utils import json_schema_type @@ -17,7 +17,7 @@ class RunpodImplConfig(BaseModel): default=None, description="The URL for the Runpod model serving endpoint", ) - api_token: str | None = Field( + api_token: SecretStr | None = Field( default=None, description="The API token", ) diff --git a/llama_stack/providers/remote/inference/runpod/runpod.py b/llama_stack/providers/remote/inference/runpod/runpod.py index ff2fe6401..5cfa05fca 100644 --- a/llama_stack/providers/remote/inference/runpod/runpod.py +++ b/llama_stack/providers/remote/inference/runpod/runpod.py @@ -103,7 +103,10 @@ class RunpodInferenceAdapter( tool_config=tool_config, ) - client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) + client = OpenAI( + base_url=self.config.url, + api_key=self.config.api_token.get_secret_value() if self.config.api_token else None, + ) if stream: return self._stream_chat_completion(request, client) else: diff --git a/llama_stack/providers/remote/inference/vertexai/vertexai.py b/llama_stack/providers/remote/inference/vertexai/vertexai.py index 770d21a2a..733f2ab4a 100644 --- a/llama_stack/providers/remote/inference/vertexai/vertexai.py +++ b/llama_stack/providers/remote/inference/vertexai/vertexai.py @@ -8,6 +8,7 @@ from typing import Any import google.auth.transport.requests from google.auth import default +from pydantic import SecretStr from llama_stack.apis.inference import ChatCompletionRequest from llama_stack.providers.utils.inference.litellm_openai_mixin import ( @@ -43,7 +44,7 @@ class VertexAIInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin): except Exception: # If we can't get credentials, return empty string to let LiteLLM handle it # This allows the LiteLLM mixin to work with ADC directly - return "" + return SecretStr("") def get_base_url(self) -> str: """ diff --git a/llama_stack/providers/remote/inference/vllm/config.py b/llama_stack/providers/remote/inference/vllm/config.py index a5bf0e4bc..978427930 100644 --- a/llama_stack/providers/remote/inference/vllm/config.py +++ b/llama_stack/providers/remote/inference/vllm/config.py @@ -6,7 +6,7 @@ from pathlib import Path -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, SecretStr, field_validator from llama_stack.schema_utils import json_schema_type @@ -21,8 +21,8 @@ class VLLMInferenceAdapterConfig(BaseModel): default=4096, description="Maximum number of tokens to generate.", ) - api_token: str | None = Field( - default="fake", + api_token: SecretStr | None = Field( + default=SecretStr("fake"), description="The API token", ) tls_verify: bool | str = Field( diff --git a/llama_stack/providers/remote/inference/vllm/vllm.py b/llama_stack/providers/remote/inference/vllm/vllm.py index 8fbb4b815..2e2e1ea5c 100644 --- a/llama_stack/providers/remote/inference/vllm/vllm.py +++ b/llama_stack/providers/remote/inference/vllm/vllm.py @@ -294,7 +294,7 @@ class VLLMInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin, Inference, ModelsPro self, model_entries=build_hf_repo_model_entries(), litellm_provider_name="vllm", - api_key_from_config=config.api_token, + api_key_from_config=config.api_token.get_secret_value(), provider_data_api_key_field="vllm_api_token", openai_compat_api_base=config.url, ) diff --git a/llama_stack/providers/remote/tool_runtime/bing_search/bing_search.py b/llama_stack/providers/remote/tool_runtime/bing_search/bing_search.py index e40903969..227deaecd 100644 --- a/llama_stack/providers/remote/tool_runtime/bing_search/bing_search.py +++ b/llama_stack/providers/remote/tool_runtime/bing_search/bing_search.py @@ -40,7 +40,7 @@ class BingSearchToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsReq def _get_api_key(self) -> str: if self.config.api_key: - return self.config.api_key + return self.config.api_key.get_secret_value() provider_data = self.get_request_provider_data() if provider_data is None or not provider_data.bing_search_api_key: diff --git a/llama_stack/providers/remote/tool_runtime/bing_search/config.py b/llama_stack/providers/remote/tool_runtime/bing_search/config.py index 30269dbc1..3c4852fba 100644 --- a/llama_stack/providers/remote/tool_runtime/bing_search/config.py +++ b/llama_stack/providers/remote/tool_runtime/bing_search/config.py @@ -6,13 +6,13 @@ from typing import Any -from pydantic import BaseModel +from pydantic import BaseModel, SecretStr class BingSearchToolConfig(BaseModel): """Configuration for Bing Search Tool Runtime""" - api_key: str | None = None + api_key: SecretStr | None = None top_k: int = 3 @classmethod diff --git a/llama_stack/providers/remote/tool_runtime/brave_search/brave_search.py b/llama_stack/providers/remote/tool_runtime/brave_search/brave_search.py index ba3b910d5..a34e94ad3 100644 --- a/llama_stack/providers/remote/tool_runtime/brave_search/brave_search.py +++ b/llama_stack/providers/remote/tool_runtime/brave_search/brave_search.py @@ -39,7 +39,7 @@ class BraveSearchToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsRe def _get_api_key(self) -> str: if self.config.api_key: - return self.config.api_key + return self.config.api_key.get_secret_value() provider_data = self.get_request_provider_data() if provider_data is None or not provider_data.brave_search_api_key: diff --git a/llama_stack/providers/remote/tool_runtime/brave_search/config.py b/llama_stack/providers/remote/tool_runtime/brave_search/config.py index f02967ce8..41f6c9bbc 100644 --- a/llama_stack/providers/remote/tool_runtime/brave_search/config.py +++ b/llama_stack/providers/remote/tool_runtime/brave_search/config.py @@ -6,11 +6,11 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr class BraveSearchToolConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( default=None, description="The Brave Search API Key", ) diff --git a/llama_stack/providers/remote/tool_runtime/tavily_search/config.py b/llama_stack/providers/remote/tool_runtime/tavily_search/config.py index ca4e615db..f81b3455a 100644 --- a/llama_stack/providers/remote/tool_runtime/tavily_search/config.py +++ b/llama_stack/providers/remote/tool_runtime/tavily_search/config.py @@ -6,11 +6,11 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr class TavilySearchToolConfig(BaseModel): - api_key: str | None = Field( + api_key: SecretStr | None = Field( default=None, description="The Tavily Search API Key", ) diff --git a/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py b/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py index 976ec9c57..9048cd864 100644 --- a/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py +++ b/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py @@ -39,7 +39,7 @@ class TavilySearchToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsR def _get_api_key(self) -> str: if self.config.api_key: - return self.config.api_key + return self.config.api_key.get_secret_value() provider_data = self.get_request_provider_data() if provider_data is None or not provider_data.tavily_search_api_key: diff --git a/llama_stack/providers/remote/tool_runtime/wolfram_alpha/wolfram_alpha.py b/llama_stack/providers/remote/tool_runtime/wolfram_alpha/wolfram_alpha.py index f12a44958..1ba7639d5 100644 --- a/llama_stack/providers/remote/tool_runtime/wolfram_alpha/wolfram_alpha.py +++ b/llama_stack/providers/remote/tool_runtime/wolfram_alpha/wolfram_alpha.py @@ -40,7 +40,7 @@ class WolframAlphaToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsR def _get_api_key(self) -> str: if self.config.api_key: - return self.config.api_key + return self.config.api_key.get_secret_value() provider_data = self.get_request_provider_data() if provider_data is None or not provider_data.wolfram_alpha_api_key: diff --git a/llama_stack/providers/remote/vector_io/pgvector/config.py b/llama_stack/providers/remote/vector_io/pgvector/config.py index 334cbe5be..abfa0eacf 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/config.py +++ b/llama_stack/providers/remote/vector_io/pgvector/config.py @@ -6,7 +6,7 @@ from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.providers.utils.kvstore.config import ( KVStoreConfig, @@ -21,7 +21,7 @@ class PGVectorVectorIOConfig(BaseModel): port: int | None = Field(default=5432) db: str | None = Field(default="postgres") user: str | None = Field(default="postgres") - password: str | None = Field(default="mysecretpassword") + password: SecretStr | None = Field(default="mysecretpassword") kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) @classmethod diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index 1c140e782..fae3c7200 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -366,7 +366,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco port=self.config.port, database=self.config.db, user=self.config.user, - password=self.config.password, + password=self.config.password.get_secret_value(), ) self.conn.autocommit = True with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur: diff --git a/llama_stack/providers/utils/bedrock/client.py b/llama_stack/providers/utils/bedrock/client.py index b3c8629e0..16f630326 100644 --- a/llama_stack/providers/utils/bedrock/client.py +++ b/llama_stack/providers/utils/bedrock/client.py @@ -50,8 +50,8 @@ def create_bedrock_client(config: BedrockBaseConfig, service_name: str = "bedroc session_args = { "aws_access_key_id": config.aws_access_key_id, - "aws_secret_access_key": config.aws_secret_access_key, - "aws_session_token": config.aws_session_token, + "aws_secret_access_key": config.aws_secret_access_key.get_secret_value(), + "aws_session_token": config.aws_session_token.get_secret_value(), "region_name": config.region_name, "profile_name": config.profile_name, "session_ttl": config.session_ttl, diff --git a/llama_stack/providers/utils/bedrock/config.py b/llama_stack/providers/utils/bedrock/config.py index 2745c88cb..b48f0ea69 100644 --- a/llama_stack/providers/utils/bedrock/config.py +++ b/llama_stack/providers/utils/bedrock/config.py @@ -6,7 +6,7 @@ import os -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr class BedrockBaseConfig(BaseModel): @@ -14,12 +14,12 @@ class BedrockBaseConfig(BaseModel): default_factory=lambda: os.getenv("AWS_ACCESS_KEY_ID"), description="The AWS access key to use. Default use environment variable: AWS_ACCESS_KEY_ID", ) - aws_secret_access_key: str | None = Field( - default_factory=lambda: os.getenv("AWS_SECRET_ACCESS_KEY"), + aws_secret_access_key: SecretStr | None = Field( + default_factory=lambda: SecretStr(val) if (val := os.getenv("AWS_SECRET_ACCESS_KEY")) else None, description="The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY", ) - aws_session_token: str | None = Field( - default_factory=lambda: os.getenv("AWS_SESSION_TOKEN"), + aws_session_token: SecretStr | None = Field( + default_factory=lambda: SecretStr(val) if (val := os.getenv("AWS_SESSION_TOKEN")) else None, description="The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN", ) region_name: str | None = Field( diff --git a/llama_stack/providers/utils/inference/litellm_openai_mixin.py b/llama_stack/providers/utils/inference/litellm_openai_mixin.py index b1e38f323..fc39852ae 100644 --- a/llama_stack/providers/utils/inference/litellm_openai_mixin.py +++ b/llama_stack/providers/utils/inference/litellm_openai_mixin.py @@ -8,6 +8,7 @@ from collections.abc import AsyncGenerator, AsyncIterator from typing import Any import litellm +from pydantic import SecretStr from llama_stack.apis.common.content_types import ( InterleavedContent, @@ -68,7 +69,7 @@ class LiteLLMOpenAIMixin( def __init__( self, litellm_provider_name: str, - api_key_from_config: str | None, + api_key_from_config: SecretStr | None, provider_data_api_key_field: str, model_entries: list[ProviderModelEntry] | None = None, openai_compat_api_base: str | None = None, @@ -247,14 +248,14 @@ class LiteLLMOpenAIMixin( return { "model": request.model, - "api_key": self.get_api_key(), + "api_key": self.get_api_key().get_secret_value(), "api_base": self.api_base, **input_dict, "stream": request.stream, **get_sampling_options(request.sampling_params), } - def get_api_key(self) -> str: + def get_api_key(self) -> SecretStr: provider_data = self.get_request_provider_data() key_field = self.provider_data_api_key_field if provider_data and getattr(provider_data, key_field, None): @@ -305,7 +306,7 @@ class LiteLLMOpenAIMixin( response = litellm.embedding( model=self.get_litellm_model_name(model_obj.provider_resource_id), input=input_list, - api_key=self.get_api_key(), + api_key=self.get_api_key().get_secret_value(), api_base=self.api_base, dimensions=dimensions, ) @@ -368,7 +369,7 @@ class LiteLLMOpenAIMixin( user=user, guided_choice=guided_choice, prompt_logprobs=prompt_logprobs, - api_key=self.get_api_key(), + api_key=self.get_api_key().get_secret_value(), api_base=self.api_base, ) return await litellm.atext_completion(**params) @@ -424,7 +425,7 @@ class LiteLLMOpenAIMixin( top_logprobs=top_logprobs, top_p=top_p, user=user, - api_key=self.get_api_key(), + api_key=self.get_api_key().get_secret_value(), api_base=self.api_base, ) return await litellm.acompletion(**params) diff --git a/llama_stack/providers/utils/inference/openai_mixin.py b/llama_stack/providers/utils/inference/openai_mixin.py index 7da97e6b1..7fbd62ef6 100644 --- a/llama_stack/providers/utils/inference/openai_mixin.py +++ b/llama_stack/providers/utils/inference/openai_mixin.py @@ -11,6 +11,7 @@ from collections.abc import AsyncIterator from typing import Any from openai import NOT_GIVEN, AsyncOpenAI +from pydantic import SecretStr from llama_stack.apis.inference import ( Model, @@ -70,14 +71,14 @@ class OpenAIMixin(ModelRegistryHelper, ABC): allowed_models: list[str] = [] @abstractmethod - def get_api_key(self) -> str: + def get_api_key(self) -> SecretStr: """ Get the API key. This method must be implemented by child classes to provide the API key for authenticating with the OpenAI API or compatible endpoints. - :return: The API key as a string + :return: The API key as a SecretStr """ pass diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index 7b6a79350..1cd90d308 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -8,7 +8,7 @@ import re from enum import Enum from typing import Annotated, Literal -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, SecretStr, field_validator from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR @@ -74,7 +74,7 @@ class PostgresKVStoreConfig(CommonConfig): port: int = 5432 db: str = "llamastack" user: str - password: str | None = None + password: SecretStr | None = None ssl_mode: str | None = None ca_cert_path: str | None = None table_name: str = "llamastack_kvstore" @@ -118,7 +118,7 @@ class MongoDBKVStoreConfig(CommonConfig): port: int = 27017 db: str = "llamastack" user: str | None = None - password: str | None = None + password: SecretStr | None = None collection_name: str = "llamastack_kvstore" @classmethod diff --git a/llama_stack/providers/utils/kvstore/mongodb/mongodb.py b/llama_stack/providers/utils/kvstore/mongodb/mongodb.py index 4d60949c1..83ce6da60 100644 --- a/llama_stack/providers/utils/kvstore/mongodb/mongodb.py +++ b/llama_stack/providers/utils/kvstore/mongodb/mongodb.py @@ -34,7 +34,7 @@ class MongoDBKVStoreImpl(KVStore): "host": self.config.host, "port": self.config.port, "username": self.config.user, - "password": self.config.password, + "password": self.config.password.get_secret_value(), } conn_creds = {k: v for k, v in conn_creds.items() if v is not None} self.conn = AsyncMongoClient(**conn_creds) diff --git a/llama_stack/providers/utils/kvstore/postgres/postgres.py b/llama_stack/providers/utils/kvstore/postgres/postgres.py index 56d6dbb48..778aa04be 100644 --- a/llama_stack/providers/utils/kvstore/postgres/postgres.py +++ b/llama_stack/providers/utils/kvstore/postgres/postgres.py @@ -30,7 +30,7 @@ class PostgresKVStoreImpl(KVStore): port=self.config.port, database=self.config.db, user=self.config.user, - password=self.config.password, + password=self.config.password.get_secret_value(), sslmode=self.config.ssl_mode, sslrootcert=self.config.ca_cert_path, ) diff --git a/llama_stack/providers/utils/sqlstore/sqlstore.py b/llama_stack/providers/utils/sqlstore/sqlstore.py index fc44402ae..3f6bedc7e 100644 --- a/llama_stack/providers/utils/sqlstore/sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/sqlstore.py @@ -9,7 +9,7 @@ from enum import StrEnum from pathlib import Path from typing import Annotated, Literal -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR @@ -63,11 +63,11 @@ class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig): port: int = 5432 db: str = "llamastack" user: str - password: str | None = None + password: SecretStr | None = None @property def engine_str(self) -> str: - return f"postgresql+asyncpg://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}" + return f"postgresql+asyncpg://{self.user}:{self.password.get_secret_value() if self.password else ''}@{self.host}:{self.port}/{self.db}" @classmethod def pip_packages(cls) -> list[str]: diff --git a/tests/unit/providers/inference/test_inference_client_caching.py b/tests/unit/providers/inference/test_inference_client_caching.py index f4b3201e9..974d55ade 100644 --- a/tests/unit/providers/inference/test_inference_client_caching.py +++ b/tests/unit/providers/inference/test_inference_client_caching.py @@ -33,7 +33,7 @@ def test_groq_provider_openai_client_caching(): with request_provider_data_context( {"x-llamastack-provider-data": json.dumps({inference_adapter.provider_data_api_key_field: api_key})} ): - assert inference_adapter.client.api_key == api_key + assert inference_adapter.client.api_key.get_secret_value() == api_key def test_openai_provider_openai_client_caching(): @@ -52,7 +52,7 @@ def test_openai_provider_openai_client_caching(): {"x-llamastack-provider-data": json.dumps({inference_adapter.provider_data_api_key_field: api_key})} ): openai_client = inference_adapter.client - assert openai_client.api_key == api_key + assert openai_client.api_key.get_secret_value() == api_key def test_together_provider_openai_client_caching(): @@ -86,4 +86,4 @@ def test_llama_compat_provider_openai_client_caching(): for api_key in ["test1", "test2"]: with request_provider_data_context({"x-llamastack-provider-data": json.dumps({"llama_api_key": api_key})}): - assert inference_adapter.client.api_key == api_key + assert inference_adapter.client.api_key.get_secret_value() == api_key diff --git a/tests/unit/providers/inference/test_litellm_openai_mixin.py b/tests/unit/providers/inference/test_litellm_openai_mixin.py index dc17e6abf..cf7623dd1 100644 --- a/tests/unit/providers/inference/test_litellm_openai_mixin.py +++ b/tests/unit/providers/inference/test_litellm_openai_mixin.py @@ -8,7 +8,7 @@ import json from unittest.mock import MagicMock import pytest -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SecretStr from llama_stack.core.request_headers import request_provider_data_context from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin @@ -16,11 +16,11 @@ from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOp # Test fixtures and helper classes class TestConfig(BaseModel): - api_key: str | None = Field(default=None) + api_key: SecretStr | None = Field(default=None) class TestProviderDataValidator(BaseModel): - test_api_key: str | None = Field(default=None) + test_api_key: SecretStr | None = Field(default=None) class TestLiteLLMAdapter(LiteLLMOpenAIMixin): @@ -36,7 +36,7 @@ class TestLiteLLMAdapter(LiteLLMOpenAIMixin): @pytest.fixture def adapter_with_config_key(): """Fixture to create adapter with API key in config""" - config = TestConfig(api_key="config-api-key") + config = TestConfig(api_key=SecretStr("config-api-key")) adapter = TestLiteLLMAdapter(config) adapter.__provider_spec__ = MagicMock() adapter.__provider_spec__.provider_data_validator = ( @@ -59,7 +59,7 @@ def adapter_without_config_key(): def test_api_key_from_config_when_no_provider_data(adapter_with_config_key): """Test that adapter uses config API key when no provider data is available""" - api_key = adapter_with_config_key.get_api_key() + api_key = adapter_with_config_key.get_api_key().get_secret_value() assert api_key == "config-api-key" @@ -68,28 +68,28 @@ def test_provider_data_takes_priority_over_config(adapter_with_config_key): with request_provider_data_context( {"x-llamastack-provider-data": json.dumps({"test_api_key": "provider-data-key"})} ): - api_key = adapter_with_config_key.get_api_key() + api_key = adapter_with_config_key.get_api_key().get_secret_value() assert api_key == "provider-data-key" def test_fallback_to_config_when_provider_data_missing_key(adapter_with_config_key): """Test fallback to config when provider data doesn't have the required key""" with request_provider_data_context({"x-llamastack-provider-data": json.dumps({"wrong_key": "some-value"})}): - api_key = adapter_with_config_key.get_api_key() + api_key = adapter_with_config_key.get_api_key().get_secret_value() assert api_key == "config-api-key" def test_error_when_no_api_key_available(adapter_without_config_key): """Test that ValueError is raised when neither config nor provider data have API key""" with pytest.raises(ValueError, match="API key is not set"): - adapter_without_config_key.get_api_key() + adapter_without_config_key.get_api_key().get_secret_value() def test_error_when_provider_data_has_wrong_key(adapter_without_config_key): """Test that ValueError is raised when provider data exists but doesn't have required key""" with request_provider_data_context({"x-llamastack-provider-data": json.dumps({"wrong_key": "some-value"})}): with pytest.raises(ValueError, match="API key is not set"): - adapter_without_config_key.get_api_key() + adapter_without_config_key.get_api_key().get_secret_value() def test_provider_data_works_when_config_is_none(adapter_without_config_key): @@ -97,14 +97,14 @@ def test_provider_data_works_when_config_is_none(adapter_without_config_key): with request_provider_data_context( {"x-llamastack-provider-data": json.dumps({"test_api_key": "provider-only-key"})} ): - api_key = adapter_without_config_key.get_api_key() + api_key = adapter_without_config_key.get_api_key().get_secret_value() assert api_key == "provider-only-key" def test_error_message_includes_correct_field_names(adapter_without_config_key): """Test that error message includes correct field name and header information""" try: - adapter_without_config_key.get_api_key() + adapter_without_config_key.get_api_key().get_secret_value() raise AssertionError("Should have raised ValueError") except ValueError as e: assert "test_api_key" in str(e) # Should mention the correct field name diff --git a/tests/unit/providers/inference/test_openai_base_url_config.py b/tests/unit/providers/inference/test_openai_base_url_config.py index 903772f0c..7bc908069 100644 --- a/tests/unit/providers/inference/test_openai_base_url_config.py +++ b/tests/unit/providers/inference/test_openai_base_url_config.py @@ -7,6 +7,8 @@ import os from unittest.mock import MagicMock, patch +from pydantic import SecretStr + from llama_stack.core.stack import replace_env_vars from llama_stack.providers.remote.inference.openai.config import OpenAIConfig from llama_stack.providers.remote.inference.openai.openai import OpenAIInferenceAdapter @@ -59,14 +61,14 @@ class TestOpenAIBaseURLConfig: adapter = OpenAIInferenceAdapter(config) # Mock the get_api_key method since it's delegated to LiteLLMOpenAIMixin - adapter.get_api_key = MagicMock(return_value="test-key") + adapter.get_api_key = MagicMock(return_value=SecretStr("test-key")) # Access the client property to trigger AsyncOpenAI initialization _ = adapter.client # Verify AsyncOpenAI was called with the correct base_url mock_openai_class.assert_called_once_with( - api_key="test-key", + api_key=SecretStr("test-key"), base_url=custom_url, ) @@ -78,7 +80,7 @@ class TestOpenAIBaseURLConfig: adapter = OpenAIInferenceAdapter(config) # Mock the get_api_key method - adapter.get_api_key = MagicMock(return_value="test-key") + adapter.get_api_key = MagicMock(return_value=SecretStr("test-key")) # Mock a model object that will be returned by models.list() mock_model = MagicMock() @@ -101,7 +103,7 @@ class TestOpenAIBaseURLConfig: # Verify the client was created with the custom URL mock_openai_class.assert_called_with( - api_key="test-key", + api_key=SecretStr("test-key"), base_url=custom_url, ) @@ -119,7 +121,7 @@ class TestOpenAIBaseURLConfig: adapter = OpenAIInferenceAdapter(config) # Mock the get_api_key method - adapter.get_api_key = MagicMock(return_value="test-key") + adapter.get_api_key = MagicMock(return_value=SecretStr("test-key")) # Mock a model object that will be returned by models.list() mock_model = MagicMock() @@ -142,6 +144,6 @@ class TestOpenAIBaseURLConfig: # Verify the client was created with the environment variable URL mock_openai_class.assert_called_with( - api_key="test-key", + api_key=SecretStr("test-key"), base_url="https://proxy.openai.com/v1", )