phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 04:04:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: prevent telemetry from leaking sensitive info

Browse source 
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 <seb@redhat.com>
This commit is contained in:
Sébastien Han 2025-08-08 15:54:45 +02:00
parent 8dc9fd6844
commit c4cb6aa8d9
No known key found for this signature in database
53 changed files with 121 additions and 109 deletions
Download patch file Download diff file Expand all files Collapse all files

6
tests/unit/providers/inference/test_inference_client_caching.py
View file

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

22
tests/unit/providers/inference/test_litellm_openai_mixin.py
View file

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

14
tests/unit/providers/inference/test_openai_base_url_config.py
View file

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