This commit is contained in:
Sébastien Han 2025-10-01 15:47:54 +02:00 committed by GitHub
commit 79ced0c85b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
94 changed files with 341 additions and 209 deletions

View file

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

View file

@ -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 = Field(
default_factory=lambda: SecretStr(os.getenv("AWS_SECRET_ACCESS_KEY", "")),
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 = Field(
default_factory=lambda: SecretStr(os.getenv("AWS_SESSION_TOKEN", "")),
description="The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN",
)
region_name: str | None = Field(

View file

@ -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,
@ -61,7 +62,7 @@ class LiteLLMOpenAIMixin(
def __init__(
self,
litellm_provider_name: str,
api_key_from_config: str | None,
api_key_from_config: SecretStr,
provider_data_api_key_field: str,
model_entries: list[ProviderModelEntry] | None = None,
openai_compat_api_base: str | None = None,
@ -240,14 +241,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):
@ -280,7 +281,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,
)
@ -343,7 +344,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)
@ -407,7 +408,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)

View file

@ -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
@ -113,7 +114,7 @@ class OpenAIMixin(ModelRegistryHelper, ABC):
implemented by child classes.
"""
return AsyncOpenAI(
api_key=self.get_api_key(),
api_key=self.get_api_key().get_secret_value(),
base_url=self.get_base_url(),
**self.get_extra_client_params(),
)

View file

@ -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 = SecretStr("")
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 = SecretStr("")
collection_name: str = "llamastack_kvstore"
@classmethod

View file

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

View file

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

View file

@ -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 = SecretStr("")
@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()}@{self.host}:{self.port}/{self.db}"
@classmethod
def pip_packages(cls) -> list[str]: