Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/app/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def _create_postgres_engine(
) -> Engine:
"""Create PostgreSQL database engine."""
postgres_url = (
f"postgresql://{config.user}:{config.password}@"
f"postgresql://{config.user}:{config.password.get_secret_value()}@"
f"{config.host}:{config.port}/{config.db}"
f"?sslmode={config.ssl_mode}&gssencmode={config.gss_encmode}"
)
Comment on lines +62 to 65
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Build the Postgres URL with SQLAlchemy URL.create to correctly escape credentials.

String interpolation will break with special characters in user/password and risks subtle connection bugs. Use URL.create and fix the options flag spacing.

Apply:

-from sqlalchemy import create_engine, text
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import URL
@@
-    postgres_url = (
-        f"postgresql://{config.user}:{config.password.get_secret_value()}@"
-        f"{config.host}:{config.port}/{config.db}"
-        f"?sslmode={config.ssl_mode}&gssencmode={config.gss_encmode}"
-    )
+    postgres_url = URL.create(
+        drivername="postgresql",
+        username=config.user,
+        password=config.password.get_secret_value(),
+        host=config.host,
+        port=int(config.port),
+        database=config.db,
+        query={"sslmode": config.ssl_mode, "gssencmode": config.gss_encmode},
+    )
@@
-    if is_custom_schema:
-        connect_args["options"] = f"-csearch_path={config.namespace}"
+    if is_custom_schema:
+        connect_args["options"] = f"-c search_path={config.namespace}"
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
f"postgresql://{config.user}:{config.password.get_secret_value()}@"
f"{config.host}:{config.port}/{config.db}"
f"?sslmode={config.ssl_mode}&gssencmode={config.gss_encmode}"
)
++ b/src/app/database.py
@@ -1,3 +1,4 @@
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
# … earlier code …
- postgres_url = (
- f"postgresql://{config.user}:{config.password.get_secret_value()}@"
- f"{config.host}:{config.port}/{config.db}"
- f"?sslmode={config.ssl_mode}&gssencmode={config.gss_encmode}"
postgres_url = URL.create(
drivername="postgresql",
username=config.user,
password=config.password.get_secret_value(),
host=config.host,
port=int(config.port),
database=config.db,
query={
"sslmode": config.ssl_mode,
"gssencmode": config.gss_encmode,
},
)
- if is_custom_schema:
if is_custom_schema:
connect_args["options"] = f"-c search_path={config.namespace}"
# … later code …
🤖 Prompt for AI Agents
In src/app/database.py around lines 62 to 65, the Postgres connection string is
being built via naive f-string interpolation which will mis-handle special
characters in the user/password and has incorrect query option spacing; replace
the f-string construction with SQLAlchemy's URL.create to properly escape
credentials and assemble the components (drivername 'postgresql', username from
config.user, password from config.password.get_secret_value(), host, port,
database name) and pass the query options as a dict (e.g., {'sslmode':
config.ssl_mode, 'gssencmode': config.gss_encmode}) so the URL is correctly
encoded and the options flag spacing is fixed.

Expand Down
2 changes: 1 addition & 1 deletion src/app/endpoints/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
},
"llama_stack": {
"url": "http://localhost:8321",
"api_key": "xyzzy",
"api_key": "*****",
"use_as_library_client": False,
"library_client_config_path": None,
},
Expand Down
6 changes: 2 additions & 4 deletions src/app/endpoints/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,10 +176,8 @@ async def query_endpoint_handler(
# Enforce RBAC: optionally disallow overriding model/provider in requests
validate_model_provider_override(query_request, request.state.authorized_actions)

# log Llama Stack configuration, but without sensitive information
llama_stack_config = configuration.llama_stack_configuration.model_copy()
llama_stack_config.api_key = "********"
logger.info("Llama stack config: %s", llama_stack_config)
# log Llama Stack configuration
logger.info("Llama stack config: %s", configuration.llama_stack_configuration)

Comment on lines +179 to 181
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Avoid logging secret-bearing config at INFO; sanitize and downgrade to DEBUG.

Even with SecretStr masking, logging the full config object at INFO on every request is risky and noisy. Emit only a sanitized JSON view and use DEBUG.

-    # log Llama Stack configuration
-    logger.info("Llama stack config: %s", configuration.llama_stack_configuration)
+    # log Llama Stack configuration (sanitized)
+    logger.debug(
+        "Llama stack config: %s",
+        configuration.llama_stack_configuration.model_dump(
+            mode="json", exclude={"api_key"}, exclude_none=True
+        ),
+    )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# log Llama Stack configuration
logger.info("Llama stack config: %s", configuration.llama_stack_configuration)
# log Llama Stack configuration (sanitized)
logger.debug(
"Llama stack config: %s",
configuration.llama_stack_configuration.model_dump(
mode="json", exclude={"api_key"}, exclude_none=True
),
)
🤖 Prompt for AI Agents
In src/app/endpoints/query.py around lines 179-181, the code currently logs the
entire Llama stack configuration at INFO which may expose secrets; change this
to log a sanitized representation at DEBUG instead: build a JSON-safe dict from
configuration that excludes or redacts secret-bearing fields (exclude fields
named like '*secret*' or instances of SecretStr/SecretBytes, or use pydantic's
.dict(exclude=...) if applicable), then logger.debug the sanitized JSON view and
remove or downgrade the existing logger.info call so no full config is logged at
INFO.

user_id, _, token = auth

Expand Down
6 changes: 2 additions & 4 deletions src/app/endpoints/streaming_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -552,10 +552,8 @@ async def streaming_query_endpoint_handler( # pylint: disable=too-many-locals
# Enforce RBAC: optionally disallow overriding model/provider in requests
validate_model_provider_override(query_request, request.state.authorized_actions)

# log Llama Stack configuration, but without sensitive information
llama_stack_config = configuration.llama_stack_configuration.model_copy()
llama_stack_config.api_key = "********"
logger.info("Llama stack config: %s", llama_stack_config)
# log Llama Stack configuration
logger.info("Llama stack config: %s", configuration.llama_stack_configuration)

Comment on lines +555 to 557
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Same logging concern: sanitize and use DEBUG.

Mirror the sanitized, lower-verbosity logging to avoid leaking sensitive config details.

-    # log Llama Stack configuration
-    logger.info("Llama stack config: %s", configuration.llama_stack_configuration)
+    # log Llama Stack configuration (sanitized)
+    logger.debug(
+        "Llama stack config: %s",
+        configuration.llama_stack_configuration.model_dump(
+            mode="json", exclude={"api_key"}, exclude_none=True
+        ),
+    )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# log Llama Stack configuration
logger.info("Llama stack config: %s", configuration.llama_stack_configuration)
# log Llama Stack configuration (sanitized)
logger.debug(
"Llama stack config: %s",
configuration.llama_stack_configuration.model_dump(
mode="json",
exclude={"api_key"},
exclude_none=True,
),
)

user_id, _user_name, token = auth

Expand Down
7 changes: 6 additions & 1 deletion src/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,12 @@ async def load(self, llama_stack_config: LlamaStackConfiguration) -> None:
else:
logger.info("Using Llama stack running as a service")
self._lsc = AsyncLlamaStackClient(
base_url=llama_stack_config.url, api_key=llama_stack_config.api_key
base_url=llama_stack_config.url,
api_key=(
llama_stack_config.api_key.get_secret_value()
if llama_stack_config.api_key is not None
else None
),
)

def get_client(self) -> AsyncLlamaStackClient:
Expand Down
5 changes: 3 additions & 2 deletions src/models/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
FilePath,
AnyHttpUrl,
PositiveInt,
SecretStr,
)
from typing_extensions import Self, Literal

Expand Down Expand Up @@ -80,7 +81,7 @@ class PostgreSQLDatabaseConfiguration(ConfigurationBase):
port: PositiveInt = 5432
db: str
user: str
password: str
password: SecretStr
namespace: Optional[str] = "lightspeed-stack"
ssl_mode: str = constants.POSTGRES_DEFAULT_SSL_MODE
gss_encmode: str = constants.POSTGRES_DEFAULT_GSS_ENCMODE
Expand Down Expand Up @@ -167,7 +168,7 @@ class LlamaStackConfiguration(ConfigurationBase):
"""Llama stack configuration."""

url: Optional[str] = None
api_key: Optional[str] = None
api_key: Optional[SecretStr] = None
use_as_library_client: Optional[bool] = None
library_client_config_path: Optional[str] = None

Expand Down
2 changes: 1 addition & 1 deletion tests/integration/test_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ def test_loading_proper_configuration(configuration_filename: str) -> None:
ls_config = cfg.llama_stack_configuration
assert ls_config.use_as_library_client is False
assert ls_config.url == "http://localhost:8321"
assert ls_config.api_key == "xyzzy"
assert ls_config.api_key.get_secret_value() == "xyzzy"

# check 'user_data_collection' section
udc_config = cfg.user_data_collection_configuration
Expand Down
30 changes: 26 additions & 4 deletions tests/unit/models/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -537,6 +537,18 @@ def test_dump_configuration(tmp_path) -> None:
user_data_collection=UserDataCollection(
feedback_enabled=False, feedback_storage=None
),
database=DatabaseConfiguration(
sqlite=None,
postgres=PostgreSQLDatabaseConfiguration(
db="lightspeed_stack",
user="ls_user",
password="ls_password",
port=5432,
ca_cert_path=None,
ssl_mode="require",
gss_encmode="disable",
),
),
mcp_servers=[],
customization=None,
inference=InferenceConfiguration(
Expand Down Expand Up @@ -601,8 +613,8 @@ def test_dump_configuration(tmp_path) -> None:
},
"llama_stack": {
"url": None,
"api_key": "whatever",
"use_as_library_client": True,
"api_key": "**********",
"library_client_config_path": "tests/configuration/run.yaml",
},
"user_data_collection": {
Expand All @@ -625,8 +637,18 @@ def test_dump_configuration(tmp_path) -> None:
"default_model": "default_model",
},
"database": {
"sqlite": {"db_path": "/tmp/lightspeed-stack.db"},
"postgres": None,
"sqlite": None,
"postgres": {
"host": "localhost",
"port": 5432,
"db": "lightspeed_stack",
"user": "ls_user",
"password": "**********",
"ssl_mode": "require",
"gss_encmode": "disable",
"namespace": "lightspeed-stack",
"ca_cert_path": None,
},
},
"authorization": None,
}
Expand Down Expand Up @@ -980,7 +1002,7 @@ def test_postgresql_database_configuration() -> None:
assert c.port == 5432
assert c.db == "db"
assert c.user == "user"
assert c.password == "password"
assert c.password.get_secret_value() == "password"
assert c.ssl_mode == POSTGRES_DEFAULT_SSL_MODE
assert c.gss_encmode == POSTGRES_DEFAULT_GSS_ENCMODE
assert c.namespace == "lightspeed-stack"
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ def test_init_from_dict() -> None:
assert cfg.configuration.name == "foo"

# check for llama_stack_configuration subsection
assert cfg.llama_stack_configuration.api_key == "xyzzy"
assert cfg.llama_stack_configuration.api_key.get_secret_value() == "xyzzy"
assert cfg.llama_stack_configuration.url == "http://x.y.com:1234"
assert cfg.llama_stack_configuration.use_as_library_client is False

Expand Down