mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(spend-logs): turn off prompt storage per key or team (#50135)
* Add per-key and team prompt storage restrictions for spend logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Apply spend-log prompt restrictions to passthrough requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Apply team prompt storage restrictions to queued requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend-logs): honor tenant prompt storage in the collector and on "false" strings The spend-event collector decided whether to ship request and response bodies from the global setting alone, so an opted-out key or team still sent prompts through the collector transport. It now uses the same key/team check as the spend log writer. A key or team setting of the string "false" was ignored and prompts kept being stored. It now counts as false, the same way the global setting parses strings. Splits the 42-case tenant storage test into one test for the policy and one for the pass-through logging contexts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(spend-logs): carry key/team prompt storage into failure, batch and side-call rows Failure rows, executed batch rows, semantic index embeddings, websearch interception calls and errors rejected before request setup built their logging metadata without the key and team metadata, so a restricted key or team still had prompts stored on those rows. They now stamp it through one helper, LiteLLMProxyRequestSetup.get_tenant_logging_metadata. The dashboard toggle now shows a saved string "false" as off, matching the server, and its label and help text come from the translation file. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Krrish Dholakia <krrishdholakia@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia
GitHub
parent
a3bf1f4399
commit
b3c4a1c3e4
25 changed files with 629 additions and 33 deletions
|
|
@ -1702,6 +1702,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
return {
|
||||
**user_api_key_metadata,
|
||||
**LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_auth),
|
||||
**parent_correlation.as_search_metadata(),
|
||||
"model_group": search_tool_name,
|
||||
"user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_auth),
|
||||
|
|
|
|||
|
|
@ -658,6 +658,7 @@ class LiteLLMExecutedBatchRunner:
|
|||
def _row_metadata(self, run: _BatchRun) -> dict[str, object]: # mutable-ok: router updates metadata in place
|
||||
return {
|
||||
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(run.user_api_key_dict),
|
||||
**LiteLLMProxyRequestSetup.get_tenant_logging_metadata(run.user_api_key_dict),
|
||||
"user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(run.user_api_key_dict),
|
||||
"user_api_end_user_max_budget": run.user_api_key_dict.end_user_max_budget,
|
||||
"tags": list(run.request_tags),
|
||||
|
|
|
|||
|
|
@ -83,6 +83,7 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj
|
|||
|
||||
return {
|
||||
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict),
|
||||
**LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict),
|
||||
"user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -67,7 +67,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 #
|
|||
deployment_params_provider,
|
||||
get_request_model_access_groups,
|
||||
sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export
|
||||
should_store_prompts_and_responses_in_spend_logs,
|
||||
should_store_prompts_for_call,
|
||||
)
|
||||
from litellm.proxy.utils import (
|
||||
REJECTED_BEFORE_REQUEST_SETUP,
|
||||
|
|
@ -425,7 +425,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
response_obj,
|
||||
start_time,
|
||||
end_time,
|
||||
store_bodies=should_store_prompts_and_responses_in_spend_logs(),
|
||||
store_bodies=should_store_prompts_for_call(kwargs),
|
||||
)
|
||||
if isinstance(event, SpendEventBuildError):
|
||||
verbose_proxy_logger.warning("collector: tracking cost in-process, event not buildable: %s", event.reason)
|
||||
|
|
@ -489,6 +489,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
|
||||
)
|
||||
_metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
|
||||
_metadata.update(LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict))
|
||||
_metadata["status"] = "failure"
|
||||
_error_information = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=original_exception,
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from fastapi import HTTPException, Request, status
|
|||
from pydantic import TypeAdapter
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
from starlette.datastructures import Headers
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
|
|
@ -1643,7 +1644,19 @@ def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
|
|||
}
|
||||
|
||||
|
||||
class TenantLoggingMetadata(TypedDict):
|
||||
user_api_key_metadata: ReadOnly[Mapping[str, object] | None]
|
||||
user_api_key_team_metadata: ReadOnly[Mapping[str, object] | None]
|
||||
|
||||
|
||||
class LiteLLMProxyRequestSetup:
|
||||
@staticmethod
|
||||
def get_tenant_logging_metadata(user_api_key_dict: UserAPIKeyAuth) -> TenantLoggingMetadata:
|
||||
return {
|
||||
"user_api_key_metadata": strip_callback_config(user_api_key_dict.metadata),
|
||||
"user_api_key_team_metadata": strip_callback_config(user_api_key_dict.team_metadata),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _get_timeout_from_request(headers: dict) -> float | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -301,7 +301,11 @@ def get_passthrough_router_request_metadata(
|
|||
)
|
||||
if request is not None:
|
||||
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_data)
|
||||
return _METADATA_ADAPTER.validate_python(request_data["litellm_metadata"])
|
||||
metadata: Final = _METADATA_ADAPTER.validate_python(request_data["litellm_metadata"])
|
||||
return {
|
||||
**metadata,
|
||||
**LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict),
|
||||
}
|
||||
|
||||
|
||||
async def llm_passthrough_factory_proxy_route(
|
||||
|
|
|
|||
|
|
@ -114,6 +114,7 @@ from litellm.proxy.common_request_processing import (
|
|||
open_sse_before_first_byte,
|
||||
resolve_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import strip_callback_config
|
||||
from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id
|
||||
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
|
||||
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -730,6 +731,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
_metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
|
||||
_metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation
|
||||
_metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups
|
||||
_metadata["user_api_key_metadata"] = _metadata["user_api_key_auth_metadata"]
|
||||
_metadata["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata)
|
||||
# The per-model budget counters are keyed off these. get_sanitized_user_information_from_key
|
||||
# returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this
|
||||
# the post-call increment finds nothing and every passthrough request goes untracked and
|
||||
|
|
|
|||
|
|
@ -173,7 +173,6 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
is_sensitive_callback_key,
|
||||
normalize_callback_names,
|
||||
process_callback,
|
||||
strip_callback_config,
|
||||
)
|
||||
from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 # legacy module exports
|
||||
_realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -19265,7 +19264,7 @@ async def async_queue_request(
|
|||
logged_api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
|
||||
data["metadata"]["user_api_key"] = logged_api_key
|
||||
data["metadata"]["user_api_key_hash"] = logged_api_key
|
||||
data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
|
||||
data["metadata"].update(LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict))
|
||||
_headers: Final = safe_get_request_headers(request).copy()
|
||||
_headers.pop("authorization", None) # do not store the original `sk-..` api key in the db
|
||||
data["metadata"]["headers"] = _headers
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ from litellm.litellm_core_utils.classifier_logging import (
|
|||
without_classifier_audit,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_str_mapping,
|
||||
get_litellm_metadata_from_kwargs,
|
||||
proxy_stamped_kubernetes_pod_routing,
|
||||
proxy_stamped_used_client_oauth_token,
|
||||
|
|
@ -87,6 +88,7 @@ from litellm.responses.litellm_completion_transformation.compaction_items import
|
|||
summary_only_compaction_content,
|
||||
with_compaction_summary,
|
||||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, SilentExperimentFlag
|
||||
from litellm.types.utils import (
|
||||
PROMPT_CARRYING_GUARDRAIL_FIELDS,
|
||||
|
|
@ -113,6 +115,12 @@ if TYPE_CHECKING:
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
_TENANT_PROMPT_STORAGE_FIELDS: Final = ("user_api_key_metadata", "user_api_key_team_metadata")
|
||||
_SPEND_LOGS_PRIVACY_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(
|
||||
Mapping[str, object] | None
|
||||
)
|
||||
|
||||
|
||||
def _get_max_string_length_prompt_in_db() -> int:
|
||||
"""
|
||||
Resolve prompt truncation threshold at runtime so values loaded later via
|
||||
|
|
@ -336,9 +344,11 @@ def _get_spend_logs_metadata(
|
|||
clean_metadata["batch_failed_requests"] = batch_failed_requests
|
||||
clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
|
||||
clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload(
|
||||
vector_store_request_metadata
|
||||
vector_store_request_metadata, metadata=_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata)
|
||||
)
|
||||
clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs(
|
||||
guardrail_information, metadata=_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata)
|
||||
)
|
||||
clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs(guardrail_information)
|
||||
clean_metadata["usage_object"] = usage_object
|
||||
clean_metadata["model_map_information"] = model_map_information
|
||||
clean_metadata["cold_storage_object_key"] = cold_storage_object_key
|
||||
|
|
@ -906,7 +916,9 @@ def get_logging_payload(
|
|||
if drop_requester_ip_address:
|
||||
clean_metadata["requester_ip_address"] = None
|
||||
clean_metadata["error_information"] = sanitize_error_information_for_spend_logs(
|
||||
clean_metadata["error_information"], original_exception=failure_exception
|
||||
clean_metadata["error_information"],
|
||||
original_exception=failure_exception,
|
||||
metadata=_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata),
|
||||
)
|
||||
|
||||
if litellm.cache is None:
|
||||
|
|
@ -1248,7 +1260,7 @@ def _get_messages_for_spend_logs_payload(
|
|||
standard_logging_payload: StandardLoggingPayload | None,
|
||||
metadata: dict | None = None,
|
||||
) -> str:
|
||||
if should_store_prompts_and_responses_in_spend_logs():
|
||||
if should_store_prompts_and_responses_in_spend_logs(_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata)):
|
||||
if standard_logging_payload is not None:
|
||||
call_type: Final = standard_logging_payload.get("call_type", "")
|
||||
if call_type == "_arealtime":
|
||||
|
|
@ -1291,8 +1303,8 @@ def stores_latest_turn_only() -> bool:
|
|||
return general_settings_view().get("spend_logs_prompt_storage") == _LATEST_TURN_STORAGE
|
||||
|
||||
|
||||
def stores_full_prompts_in_spend_logs() -> bool:
|
||||
return should_store_prompts_and_responses_in_spend_logs() and not stores_latest_turn_only()
|
||||
def stores_full_prompts_in_spend_logs(metadata: Mapping[str, object] | None = None) -> bool:
|
||||
return should_store_prompts_and_responses_in_spend_logs(metadata) and not stores_latest_turn_only()
|
||||
|
||||
|
||||
def _block_type(block: object) -> str | None:
|
||||
|
|
@ -1651,6 +1663,7 @@ def _redact_prompt_leaks_in_error_string(text: str) -> str:
|
|||
|
||||
def _sanitize_guardrail_information_for_spend_logs(
|
||||
guardrail_information: list[StandardLoggingGuardrailInformation] | None,
|
||||
metadata: Mapping[str, object] | None = None,
|
||||
) -> list[StandardLoggingGuardrailInformation] | None:
|
||||
"""
|
||||
When ``store_prompts_in_spend_logs`` is False, redact prompt-carrying fields
|
||||
|
|
@ -1671,7 +1684,7 @@ def _sanitize_guardrail_information_for_spend_logs(
|
|||
here to match OTEL's defensive read pattern; otherwise iteration would
|
||||
yield the dict's keys and crash the whole spend-log write.
|
||||
"""
|
||||
if guardrail_information is None or stores_full_prompts_in_spend_logs():
|
||||
if guardrail_information is None or stores_full_prompts_in_spend_logs(metadata):
|
||||
return guardrail_information
|
||||
entries: Final = [guardrail_information] if isinstance(guardrail_information, dict) else guardrail_information
|
||||
return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)]
|
||||
|
|
@ -1721,6 +1734,7 @@ def _redact_prompt_fields_in_guardrail_entry(
|
|||
def sanitize_error_information_for_spend_logs(
|
||||
error_information: StandardLoggingPayloadErrorInformation | None,
|
||||
original_exception: BaseException | None = None,
|
||||
metadata: Mapping[str, object] | None = None,
|
||||
) -> StandardLoggingPayloadErrorInformation | None:
|
||||
"""
|
||||
Sanitize ``error_information`` before it lands in ``LiteLLM_SpendLogs.metadata``.
|
||||
|
|
@ -1749,7 +1763,7 @@ def sanitize_error_information_for_spend_logs(
|
|||
)
|
||||
sanitized = cast(dict, {**persisted})
|
||||
|
||||
if not stores_full_prompts_in_spend_logs():
|
||||
if not stores_full_prompts_in_spend_logs(metadata):
|
||||
for field in ("error_message", "traceback"):
|
||||
value = sanitized.get(field)
|
||||
if isinstance(value, str):
|
||||
|
|
@ -2057,7 +2071,7 @@ def _get_proxy_server_request_for_spend_logs_payload(
|
|||
|
||||
If turn_off_message_logging is enabled, redact messages in the request body.
|
||||
"""
|
||||
if should_store_prompts_and_responses_in_spend_logs():
|
||||
if should_store_prompts_and_responses_in_spend_logs(_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata)):
|
||||
_proxy_server_request: Final = cast(dict | None, litellm_params.get("proxy_server_request", EMPTY_MAPPING))
|
||||
if _proxy_server_request is not None:
|
||||
_request_body = _proxy_server_request.get("body", EMPTY_MAPPING) or EMPTY_MAPPING
|
||||
|
|
@ -2098,11 +2112,12 @@ def _get_proxy_server_request_for_spend_logs_payload(
|
|||
|
||||
def _get_vector_store_request_for_spend_logs_payload(
|
||||
vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None,
|
||||
metadata: Mapping[str, object] | None = None,
|
||||
) -> list[StandardLoggingVectorStoreRequest] | None:
|
||||
"""
|
||||
If user does not want to store prompts and responses, then remove the content from the vector store request metadata
|
||||
"""
|
||||
if stores_full_prompts_in_spend_logs():
|
||||
if stores_full_prompts_in_spend_logs(metadata):
|
||||
return vector_store_request_metadata
|
||||
|
||||
if vector_store_request_metadata is None:
|
||||
|
|
@ -2136,7 +2151,9 @@ def _get_response_for_spend_logs_payload(
|
|||
) -> str:
|
||||
if payload is None:
|
||||
return "{}"
|
||||
if should_store_prompts_and_responses_in_spend_logs():
|
||||
if should_store_prompts_and_responses_in_spend_logs(
|
||||
_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(get_litellm_metadata_from_kwargs(kwargs or {}))
|
||||
):
|
||||
response_obj: object = payload.get("response")
|
||||
if response_obj is None:
|
||||
return "{}"
|
||||
|
|
@ -2186,10 +2203,27 @@ def _get_response_for_spend_logs_payload(
|
|||
return "{}"
|
||||
|
||||
|
||||
def should_store_prompts_and_responses_in_spend_logs() -> bool:
|
||||
def _disables_prompt_storage(tenant_metadata: object) -> bool:
|
||||
settings: Final = as_str_mapping(tenant_metadata)
|
||||
value: Final = settings.get("store_prompts_in_spend_logs") if settings is not None else None
|
||||
return value is False or (isinstance(value, str) and str_to_bool(value) is False)
|
||||
|
||||
|
||||
def should_store_prompts_for_call(kwargs: Mapping[str, object]) -> bool:
|
||||
return should_store_prompts_and_responses_in_spend_logs(
|
||||
_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(get_litellm_metadata_from_kwargs(dict(kwargs)))
|
||||
)
|
||||
|
||||
|
||||
def should_store_prompts_and_responses_in_spend_logs(metadata: Mapping[str, object] | None = None) -> bool:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
if metadata is not None and any(
|
||||
_disables_prompt_storage(metadata.get(field)) for field in _TENANT_PROMPT_STORAGE_FIELDS
|
||||
):
|
||||
return False
|
||||
|
||||
# Check general_settings (from DB or proxy_config.yaml)
|
||||
store_prompts_value: Final = general_settings.get("store_prompts_in_spend_logs")
|
||||
|
||||
|
|
|
|||
|
|
@ -4258,7 +4258,12 @@ class ProxyLogging:
|
|||
)
|
||||
request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller
|
||||
request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP})
|
||||
normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata)
|
||||
normalized_metadata_slot(request_data, "metadata").update(
|
||||
{
|
||||
**user_api_key_logged_metadata,
|
||||
**LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict),
|
||||
}
|
||||
)
|
||||
|
||||
if litellm_logging_obj is not None:
|
||||
## UPDATE LOGGING INPUT
|
||||
|
|
|
|||
|
|
@ -726,6 +726,33 @@ async def test_create_dispatches_each_row_with_the_key_team_and_project_tags_as_
|
|||
assert inherited == [("key-tag", "shared", "®ion:eu", "project-tag")] * 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key_metadata", "team_metadata", "stored"),
|
||||
[
|
||||
({"store_prompts_in_spend_logs": False}, {}, False),
|
||||
({}, {"store_prompts_in_spend_logs": False}, False),
|
||||
({}, {}, True),
|
||||
],
|
||||
)
|
||||
async def test_create_dispatches_rows_that_keep_the_key_and_team_prompt_storage_policy(
|
||||
monkeypatch: pytest.MonkeyPatch, key_metadata: dict[str, object], team_metadata: dict[str, object], stored: bool
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import should_store_prompts_and_responses_in_spend_logs
|
||||
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True})
|
||||
harness: Final = make_runner(
|
||||
user=UserAPIKeyAuth(api_key="sk-batch-key", metadata=key_metadata, team_metadata=team_metadata)
|
||||
)
|
||||
await harness.create_and_finish()
|
||||
|
||||
decisions: Final = [
|
||||
should_store_prompts_and_responses_in_spend_logs(call.kwargs["metadata"])
|
||||
for call in harness.router.acompletion.await_args_list
|
||||
]
|
||||
assert decisions == [stored] * 2
|
||||
|
||||
|
||||
async def test_create_dispatches_cli_session_rows_under_the_per_user_alias_not_the_login_token() -> None:
|
||||
session = UserAPIKeyAuth(
|
||||
api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa",
|
||||
|
|
|
|||
|
|
@ -498,6 +498,50 @@ async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_lit
|
|||
assert stored_request["litellm_metadata"]["headers"] == request_headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("key_metadata", "team_metadata", "code_word_is_stored"),
|
||||
[
|
||||
({"store_prompts_in_spend_logs": False}, {}, False),
|
||||
({}, {"store_prompts_in_spend_logs": False}, False),
|
||||
({"store_prompts_in_spend_logs": True}, {"store_prompts_in_spend_logs": False}, False),
|
||||
({}, {}, True),
|
||||
],
|
||||
)
|
||||
async def test_async_post_call_failure_hook_honors_key_and_team_prompt_storage(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
key_metadata: dict[str, object],
|
||||
team_metadata: dict[str, object],
|
||||
code_word_is_stored: bool,
|
||||
):
|
||||
monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true")
|
||||
writer: Final = MagicMock(spec=DBSpendUpdateWriter)
|
||||
writer.update_database = AsyncMock()
|
||||
request_input: Final = "The code word is VERMILION-42. Reply OK."
|
||||
|
||||
await ProxyDBLogger(spend_writer=lambda: writer).async_post_call_failure_hook(
|
||||
request_data={
|
||||
"model": "gpt-5.6",
|
||||
"input": request_input,
|
||||
"litellm_metadata": {},
|
||||
"proxy_server_request": {
|
||||
"url": "http://localhost/v1/responses",
|
||||
"body": {"model": "gpt-5.6", "input": request_input},
|
||||
},
|
||||
},
|
||||
original_exception=Exception("temperature must be at most 2"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", metadata=key_metadata, team_metadata=team_metadata),
|
||||
)
|
||||
|
||||
row: Final = get_logging_payload(
|
||||
kwargs=writer.update_database.call_args.kwargs["kwargs"],
|
||||
response_obj={},
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
assert ("VERMILION-42" in row["proxy_server_request"]) is code_word_is_stored
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"used_client_oauth_token, custom_llm_provider, expected",
|
||||
|
|
@ -3398,6 +3442,43 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski
|
|||
assert event.response_cost == 0.0125
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("tenant_metadata", "stored"),
|
||||
[
|
||||
({"user_api_key_metadata": {"store_prompts_in_spend_logs": False}}, False),
|
||||
({"user_api_key_team_metadata": {"store_prompts_in_spend_logs": False}}, False),
|
||||
({"user_api_key_team_metadata": {"store_prompts_in_spend_logs": "false"}}, False),
|
||||
({"user_api_key_metadata": {}, "user_api_key_team_metadata": {}}, True),
|
||||
],
|
||||
)
|
||||
async def test_spend_event_leaves_out_bodies_for_tenants_that_turned_off_prompt_storage(
|
||||
socket_path: Path, monkeypatch: pytest.MonkeyPatch, tenant_metadata: dict[str, object], stored: bool
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True})
|
||||
handler: Final = _RecordingHandler()
|
||||
consumer: Final = SpendEventConsumer(handler)
|
||||
address: Final = UnixAddress(path=str(socket_path))
|
||||
server: Final = await consumer.serve(address)
|
||||
producer: Final = SpendEventProducer(
|
||||
address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback
|
||||
)
|
||||
kwargs: Final[dict[str, object]] = _offload_kwargs()
|
||||
kwargs["litellm_params"]["metadata"].update(tenant_metadata)
|
||||
|
||||
await ProxyDBLogger(producer).async_log_success_event(kwargs, _offload_response(), datetime.now(), datetime.now())
|
||||
await producer.close(drain_timeout=5.0)
|
||||
server.close()
|
||||
assert await consumer.drain(timeout=5.0) == 0
|
||||
|
||||
event: Final = decode_spend_event(handler.lines[0])
|
||||
assert not isinstance(event, SpendEventDecodeError)
|
||||
assert (event.litellm_params["proxy_server_request"] is not None) is stored
|
||||
assert ("messages" in event.standard_logging_object) is stored
|
||||
|
||||
|
||||
def _completed_batch() -> LiteLLMBatch:
|
||||
return LiteLLMBatch(
|
||||
id="batch_abc",
|
||||
|
|
|
|||
|
|
@ -8,14 +8,19 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import Callable
|
||||
from contextlib import AbstractContextManager
|
||||
from datetime import datetime, timezone
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm import Router
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
|
||||
from .conftest import normalize # type: ignore[import-not-found]
|
||||
|
||||
|
|
@ -133,6 +138,65 @@ def _queue_router(
|
|||
return router
|
||||
|
||||
|
||||
class _QueueRequestLogger(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.requests: Final = SimpleQueue[dict[str, object]]()
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
|
||||
self.requests.put(kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("storage_mode", [None, "latest_turn"])
|
||||
@pytest.mark.parametrize(
|
||||
("key_setting", "team_setting", "stored"),
|
||||
[(None, False, False), (True, False, False), (False, None, False), (None, None, True), (True, True, True)],
|
||||
)
|
||||
def test_queued_requests_apply_authenticated_prompt_storage_policy(
|
||||
client: TestClient,
|
||||
auth_as: Callable[..., AbstractContextManager[None]],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
storage_mode: str | None,
|
||||
key_setting: bool | None,
|
||||
team_setting: bool | None,
|
||||
stored: bool,
|
||||
) -> None:
|
||||
_queue_router(monkeypatch, default_priority=0)
|
||||
logger: Final = _QueueRequestLogger()
|
||||
monkeypatch.setattr(litellm, "callbacks", [logger])
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"general_settings",
|
||||
{"store_prompts_in_spend_logs": True, "spend_logs_prompt_storage": storage_mode},
|
||||
)
|
||||
with auth_as(
|
||||
metadata={} if key_setting is None else {"store_prompts_in_spend_logs": key_setting},
|
||||
team_metadata={} if team_setting is None else {"store_prompts_in_spend_logs": team_setting},
|
||||
):
|
||||
response: Final = client.post(
|
||||
"/queue/chat/completions",
|
||||
json={
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "queue-private-prompt"}],
|
||||
"metadata": {
|
||||
"user_api_key_metadata": {"store_prompts_in_spend_logs": not stored},
|
||||
"user_api_key_team_metadata": {"store_prompts_in_spend_logs": not stored},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["choices"][0]["message"]["content"] == "queued"
|
||||
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs=logger.requests.get_nowait(),
|
||||
response_obj=litellm.ModelResponse.model_validate(response.json()),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
assert ("queue-private-prompt" in payload["proxy_server_request"]) is stored
|
||||
assert ("queued" in payload["response"]) is stored
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("default_priority", "priority_fields", "expected_priority"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing import Any, Final, cast
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
|
|
@ -32,10 +33,14 @@ from litellm.constants import (
|
|||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import SpendLogsMetadata, SpendLogsMetadataFields, SpendLogsPayload, UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import get_passthrough_router_request_metadata
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import PassthroughStandardLoggingPayload
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
_extract_usage_for_ocr_call,
|
||||
|
|
@ -211,6 +216,178 @@ def test_classifier_audit_spend_storage_obeys_privacy_and_truncation(monkeypatch
|
|||
assert audit["classifier_input"]["system"] == "rubric" * 1000
|
||||
|
||||
|
||||
_TENANT_PROMPT: Final = "tenant-private-prompt"
|
||||
_TENANT_RESPONSE: Final = "tenant-private-response"
|
||||
|
||||
|
||||
def _tenant_settings(setting: bool | str | None) -> dict[str, object]:
|
||||
return {} if setting is None else {"store_prompts_in_spend_logs": setting}
|
||||
|
||||
|
||||
def _spend_log_for(litellm_params: Mapping[str, object]) -> tuple[SpendLogsPayload, dict[str, object]]:
|
||||
now: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
messages: Final = [{"role": "user", "content": _TENANT_PROMPT}]
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "gpt-4o-mini",
|
||||
"litellm_params": {**litellm_params, "proxy_server_request": {"body": {"messages": messages}}},
|
||||
"standard_logging_object": {
|
||||
"call_type": "_arealtime",
|
||||
"model_map_information": None,
|
||||
"messages": messages,
|
||||
"response": {"choices": [{"message": {"content": _TENANT_RESPONSE}}]},
|
||||
"metadata": {},
|
||||
},
|
||||
}
|
||||
response: Final = litellm.ModelResponse(
|
||||
id="tenant-storage",
|
||||
choices=[{"message": {"role": "assistant", "content": _TENANT_RESPONSE}}],
|
||||
usage=litellm.Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5),
|
||||
)
|
||||
return get_logging_payload(kwargs=kwargs, response_obj=response, start_time=now, end_time=now), kwargs
|
||||
|
||||
|
||||
def _stores_tenant_content(payload: SpendLogsPayload) -> tuple[bool, bool, bool]:
|
||||
return (
|
||||
_TENANT_PROMPT in payload["messages"],
|
||||
_TENANT_PROMPT in payload["proxy_server_request"],
|
||||
_TENANT_RESPONSE in payload["response"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("storage_mode", [None, "latest_turn"])
|
||||
@pytest.mark.parametrize(
|
||||
("key_setting", "team_setting", "global_store", "stored"),
|
||||
[
|
||||
(False, None, True, False),
|
||||
(None, False, True, False),
|
||||
(True, False, True, False),
|
||||
(False, True, True, False),
|
||||
(None, "false", True, False),
|
||||
(None, None, True, True),
|
||||
(True, True, True, True),
|
||||
(True, True, False, False),
|
||||
],
|
||||
)
|
||||
def test_tenant_prompt_storage_policy_decides_what_spend_logs_keep(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
key_setting: bool | str | None,
|
||||
team_setting: bool | str | None,
|
||||
global_store: bool,
|
||||
stored: bool,
|
||||
storage_mode: str | None,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"general_settings",
|
||||
{"store_prompts_in_spend_logs": global_store, "spend_logs_prompt_storage": storage_mode},
|
||||
)
|
||||
monkeypatch.delenv("STORE_PROMPTS_IN_SPEND_LOGS", raising=False)
|
||||
payload, kwargs = _spend_log_for(
|
||||
{
|
||||
"metadata": {
|
||||
"user_api_key_metadata": _tenant_settings(key_setting),
|
||||
"user_api_key_team_metadata": _tenant_settings(team_setting),
|
||||
"store_prompts_in_spend_logs": True,
|
||||
"headers": {"litellm-disable-message-redaction": "true"},
|
||||
}
|
||||
}
|
||||
)
|
||||
assert _stores_tenant_content(payload) == (stored, stored, stored)
|
||||
assert kwargs["standard_logging_object"]["messages"] == [{"role": "user", "content": _TENANT_PROMPT}]
|
||||
assert kwargs["standard_logging_object"]["response"] == {"choices": [{"message": {"content": _TENANT_RESPONSE}}]}
|
||||
|
||||
|
||||
def _passthrough_logging_params(auth: UserAPIKeyAuth, request: Request) -> Mapping[str, object]:
|
||||
context: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint(
|
||||
request=request,
|
||||
user_api_key_dict=auth,
|
||||
passthrough_logging_payload=PassthroughStandardLoggingPayload(url="https://api.openai.com/v1/chat/completions"),
|
||||
logging_obj=LiteLLMLoggingObj(
|
||||
model="test-model",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="pass_through_endpoint",
|
||||
start_time=datetime.datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
litellm_call_id="tenant-storage",
|
||||
function_id="tenant-storage",
|
||||
),
|
||||
_parsed_body={
|
||||
"messages": [{"role": "user", "content": _TENANT_PROMPT}],
|
||||
"metadata": {
|
||||
"user_api_key_metadata": {"store_prompts_in_spend_logs": True},
|
||||
"user_api_key_team_metadata": {"store_prompts_in_spend_logs": True},
|
||||
},
|
||||
},
|
||||
)
|
||||
return TypeAdapter[Mapping[str, object]](Mapping[str, object]).validate_python(context["litellm_params"])
|
||||
|
||||
|
||||
def _router_passthrough_logging_params(auth: UserAPIKeyAuth, request: Request) -> Mapping[str, object]:
|
||||
return {"litellm_metadata": dict(get_passthrough_router_request_metadata(auth, request))}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("logging_params_for", [_passthrough_logging_params, _router_passthrough_logging_params])
|
||||
@pytest.mark.parametrize(
|
||||
("key_setting", "team_setting", "stored"),
|
||||
[(False, None, False), (None, False, False), (True, False, False), (None, None, True)],
|
||||
)
|
||||
def test_passthrough_spend_logs_follow_the_authenticated_tenant_storage_policy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
logging_params_for: Callable[[UserAPIKeyAuth, Request], Mapping[str, object]],
|
||||
key_setting: bool | None,
|
||||
team_setting: bool | None,
|
||||
stored: bool,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True})
|
||||
monkeypatch.delenv("STORE_PROMPTS_IN_SPEND_LOGS", raising=False)
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="test-storage-key", metadata=_tenant_settings(key_setting), team_metadata=_tenant_settings(team_setting)
|
||||
)
|
||||
request: Final = Request(
|
||||
{"type": "http", "method": "POST", "path": "/openai/v1/chat/completions", "headers": [], "query_string": b""}
|
||||
)
|
||||
payload, _ = _spend_log_for(logging_params_for(auth, request))
|
||||
assert _stores_tenant_content(payload) == (stored, stored, stored)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tenant_field", ["user_api_key_metadata", "user_api_key_team_metadata"])
|
||||
def test_tenant_prompt_restriction_redacts_ancillary_fields_without_mutating_callbacks(
|
||||
monkeypatch: pytest.MonkeyPatch, tenant_field: str
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True})
|
||||
metadata: Final = {tenant_field: {"store_prompts_in_spend_logs": False}}
|
||||
guardrails: Final = [{"guardrail_name": "test", "guardrail_request": "private-guardrail"}]
|
||||
vectors: Final = [
|
||||
{"vector_store_search_response": {"data": [{"content": [{"text": "private-vector", "type": "text"}]}]}}
|
||||
]
|
||||
error: Final = {
|
||||
"error_message": 'bad request {"messages": [{"role": "user", "content": "private-error"}]}',
|
||||
"error_class": "BadRequestError",
|
||||
}
|
||||
cleaned: Final = _get_spend_logs_metadata(
|
||||
metadata=metadata, guardrail_information=guardrails, vector_store_request_metadata=vectors
|
||||
)
|
||||
assert "private-guardrail" not in json.dumps(cleaned["guardrail_information"])
|
||||
assert "private-vector" not in json.dumps(cleaned["vector_store_request_metadata"])
|
||||
assert guardrails == [{"guardrail_name": "test", "guardrail_request": "private-guardrail"}]
|
||||
assert vectors[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == "private-vector"
|
||||
assert "private-error" not in json.dumps(sanitize_error_information_for_spend_logs(error, metadata=metadata))
|
||||
assert "private-error" in error["error_message"]
|
||||
|
||||
|
||||
def test_request_body_metadata_cannot_disable_spend_prompt_storage(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True})
|
||||
assert should_store_prompts_and_responses_in_spend_logs({"store_prompts_in_spend_logs": False}) is True
|
||||
|
||||
|
||||
def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_tokens():
|
||||
additional_usage_values = _get_additional_usage_values_for_usage(
|
||||
litellm.Usage(
|
||||
|
|
|
|||
|
|
@ -47,6 +47,20 @@ afterEach(() => {
|
|||
});
|
||||
|
||||
describe("Teams creation integration", () => {
|
||||
it("sends the prompt storage restriction in team metadata", async () => {
|
||||
renderWithProviders(<Teams accessToken="sk-session-test" userID="test-admin" userRole="Admin" premiumUser />);
|
||||
fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]);
|
||||
fireEvent.change(await screen.findByLabelText(/team name/i), { target: { value: "Private Team" } });
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Store prompts in spend logs" }));
|
||||
fireEvent.click(screen.getAllByRole("button", { name: /create team/i }).at(-1)!);
|
||||
await waitFor(() =>
|
||||
expect(vi.mocked(fetch).mock.calls.some(([url]) => String(url).includes("/team/new"))).toBe(true),
|
||||
);
|
||||
const request = vi.mocked(fetch).mock.calls.find(([url]) => String(url).includes("/team/new"));
|
||||
const payload = JSON.parse(String(request?.[1]?.body)) as { metadata: Record<string, unknown> };
|
||||
expect(payload.metadata.store_prompts_in_spend_logs).toBe(false);
|
||||
});
|
||||
|
||||
it("sends a configured default member per-model budget", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
renderWithProviders(<Teams accessToken="sk-session-test" userID="test-admin" userRole="Admin" premiumUser />);
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ import DeleteResourceModal from "./common_components/DeleteResourceModal";
|
|||
import { teamCreateCall } from "./networking";
|
||||
import { normalizeTeamModelSelection } from "./team/teamModelAccess";
|
||||
import { ModelSelect } from "./ModelSelect/ModelSelect";
|
||||
import { StorePromptsToggle } from "./common_components/StorePromptsToggle";
|
||||
|
||||
const SUPPRESSED_BY_DESCRIPTION = "";
|
||||
|
||||
|
|
@ -877,6 +878,19 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
|
|||
</FormField>
|
||||
<Field>
|
||||
<FieldLabel>{t("teams.page.create.fields.metadata.label")}</FieldLabel>
|
||||
<FormField control={form.control} name="metadata">
|
||||
{({ value, onChange }) => (
|
||||
<StorePromptsToggle
|
||||
value={metadataPairsToObject(value).store_prompts_in_spend_logs}
|
||||
onChange={(store) =>
|
||||
onChange([
|
||||
...(value ?? []).filter((pair) => pair.key !== "store_prompts_in_spend_logs"),
|
||||
{ key: "store_prompts_in_spend_logs", value: JSON.stringify(store) },
|
||||
])
|
||||
}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
<MetadataKeyValueFields
|
||||
control={form.control}
|
||||
getValues={form.getValues}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,34 @@
|
|||
import { fireEvent, render, screen } from "@testing-library/react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { StorePromptsMetadataToggle } from "./StorePromptsToggle";
|
||||
|
||||
describe("StorePromptsMetadataToggle", () => {
|
||||
it("leaves an unset default untouched and writes false while preserving metadata", () => {
|
||||
const onChange = vi.fn();
|
||||
render(<StorePromptsMetadataToggle value={'{"note":"keep"}'} onChange={onChange} />);
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Store prompts in spend logs" }));
|
||||
expect(JSON.parse(onChange.mock.calls[0][0])).toEqual({ note: "keep", store_prompts_in_spend_logs: false });
|
||||
});
|
||||
|
||||
it("can remove an existing restriction", () => {
|
||||
const onChange = vi.fn();
|
||||
render(<StorePromptsMetadataToggle value={'{"store_prompts_in_spend_logs":false}'} onChange={onChange} />);
|
||||
expect(screen.getByRole("switch")).not.toBeChecked();
|
||||
fireEvent.click(screen.getByRole("switch"));
|
||||
expect(JSON.parse(onChange.mock.calls[0][0])).toEqual({ store_prompts_in_spend_logs: true });
|
||||
});
|
||||
|
||||
it("shows a saved string false restriction as off", () => {
|
||||
render(<StorePromptsMetadataToggle value={'{"store_prompts_in_spend_logs":"False"}'} onChange={vi.fn()} />);
|
||||
expect(screen.getByRole("switch", { name: "Store prompts in spend logs" })).not.toBeChecked();
|
||||
});
|
||||
|
||||
it("does not overwrite invalid metadata", () => {
|
||||
const onChange = vi.fn();
|
||||
render(<StorePromptsMetadataToggle value="{" onChange={onChange} />);
|
||||
expect(screen.getByRole("switch")).toHaveAttribute("aria-disabled", "true");
|
||||
fireEvent.click(screen.getByRole("switch"));
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import { useId } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { Field, FieldDescription, FieldLabel } from "@/components/ui/field";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
|
||||
interface StorePromptsToggleProps {
|
||||
value?: unknown;
|
||||
onChange: (value: boolean) => void;
|
||||
disabled?: boolean;
|
||||
}
|
||||
|
||||
function disablesPromptStorage(value: unknown): boolean {
|
||||
return value === false || (typeof value === "string" && value.trim().toLowerCase() === "false");
|
||||
}
|
||||
|
||||
export function StorePromptsToggle({ value, onChange, disabled }: StorePromptsToggleProps) {
|
||||
const id = useId();
|
||||
const { t } = useTranslation();
|
||||
return (
|
||||
<Field>
|
||||
<FieldLabel htmlFor={id}>{t("storePromptsToggle.label")}</FieldLabel>
|
||||
<Switch id={id} checked={!disablesPromptStorage(value)} onCheckedChange={onChange} disabled={disabled} />
|
||||
<FieldDescription>{t("storePromptsToggle.description")}</FieldDescription>
|
||||
</Field>
|
||||
);
|
||||
}
|
||||
|
||||
export function StorePromptsMetadataToggle({
|
||||
value,
|
||||
onChange,
|
||||
}: {
|
||||
value: string | undefined;
|
||||
onChange: (value: string) => void;
|
||||
}) {
|
||||
const metadata = readMetadata(value);
|
||||
return (
|
||||
<StorePromptsToggle
|
||||
value={metadata?.store_prompts_in_spend_logs}
|
||||
disabled={metadata === null}
|
||||
onChange={(store) => onChange(JSON.stringify({ ...metadata, store_prompts_in_spend_logs: store }, null, 2))}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
function readMetadata(value: string | undefined): Record<string, unknown> | null {
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(value || "{}");
|
||||
return parsed !== null && typeof parsed === "object" && !Array.isArray(parsed)
|
||||
? (parsed as Record<string, unknown>)
|
||||
: null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
|
@ -22,6 +22,13 @@ const baseInput: KeyCreateInput = {
|
|||
const build = (formValues: Record<string, unknown>, overrides: Partial<KeyCreateInput> = {}): KeyPayloadResult =>
|
||||
buildKeyCreatePayload({ ...baseInput, ...overrides, formValues });
|
||||
|
||||
it.each([false, true])("preserves the prompt storage restriction %s in key metadata", (store) => {
|
||||
const result = build({ metadata: JSON.stringify({ note: "keep", store_prompts_in_spend_logs: store }) });
|
||||
expect(result.kind).toBe("ok");
|
||||
if (result.kind !== "ok") throw new Error(result.kind);
|
||||
expect(JSON.parse(String(result.payload.metadata))).toEqual({ note: "keep", store_prompts_in_spend_logs: store });
|
||||
});
|
||||
|
||||
const payloadOf = (result: KeyPayloadResult): Record<string, unknown> => {
|
||||
expect(result.kind).toBe("ok");
|
||||
if (result.kind !== "ok") throw new Error("unreachable");
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ import BudgetDurationDropdown from "../common_components/budget_duration_dropdow
|
|||
import SchemaFormFields from "../common_components/check_openapi_schema";
|
||||
import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings";
|
||||
import ModelAliasManager from "../common_components/ModelAliasManager";
|
||||
import { StorePromptsMetadataToggle } from "../common_components/StorePromptsToggle";
|
||||
import {
|
||||
MountedFormField,
|
||||
MountedFormProvider,
|
||||
|
|
@ -1571,12 +1572,18 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
className="mt-4"
|
||||
>
|
||||
{(control) => (
|
||||
<Textarea
|
||||
{...control}
|
||||
value={(control.value as string | undefined) ?? ""}
|
||||
rows={4}
|
||||
placeholder={t("createKey.metadata.placeholder")}
|
||||
/>
|
||||
<>
|
||||
<StorePromptsMetadataToggle
|
||||
value={control.value as string | undefined}
|
||||
onChange={control.onChange}
|
||||
/>
|
||||
<Textarea
|
||||
{...control}
|
||||
value={(control.value as string | undefined) ?? ""}
|
||||
rows={4}
|
||||
placeholder={t("createKey.metadata.placeholder")}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
</MountedFormField>
|
||||
<MountedFormField
|
||||
|
|
|
|||
|
|
@ -832,6 +832,21 @@ describe("TeamInfoView - model settings", () => {
|
|||
seedDefaultMocks();
|
||||
});
|
||||
|
||||
it("submits the prompt storage restriction in team metadata", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: [] }));
|
||||
renderWithProviders(<TeamInfoView {...props} userModels={["gpt-4"]} />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Store prompts in spend logs" }));
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).toEqual(
|
||||
expect.objectContaining({ metadata: expect.objectContaining({ store_prompts_in_spend_logs: false }) }),
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -114,6 +114,7 @@ import SearchToolSelector from "../search_tools/SearchToolSelector";
|
|||
import SkillSelector from "../skills/SkillSelector";
|
||||
import EditLoggingSettings from "./EditLoggingSettings";
|
||||
import RouterSettingsAccordion, { RouterSettingsAccordionRef } from "../common_components/RouterSettingsAccordion";
|
||||
import { StorePromptsToggle } from "../common_components/StorePromptsToggle";
|
||||
import MemberModal from "./EditMembership";
|
||||
import MemberPermissions from "./member_permissions";
|
||||
import MyUserTab from "./MyUserTab";
|
||||
|
|
@ -1976,6 +1977,19 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
|
||||
<Field>
|
||||
<FieldLabel>{t("teams.teamInfo.editForm.metadata")}</FieldLabel>
|
||||
<FormField control={form.control} name="metadata">
|
||||
{({ value, onChange }) => (
|
||||
<StorePromptsToggle
|
||||
value={metadataPairsToObject(value).store_prompts_in_spend_logs}
|
||||
onChange={(store) =>
|
||||
onChange([
|
||||
...(value ?? []).filter((pair) => pair.key !== "store_prompts_in_spend_logs"),
|
||||
{ key: "store_prompts_in_spend_logs", value: JSON.stringify(store) },
|
||||
])
|
||||
}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
<MetadataKeyValueFields
|
||||
control={form.control}
|
||||
getValues={form.getValues}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import { FormField } from "@/components/shared/form/FormField";
|
|||
import { toast } from "@/lib/toast";
|
||||
import AgentSelector from "../agent_management/AgentSelector";
|
||||
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
|
||||
import { StorePromptsMetadataToggle } from "../common_components/StorePromptsToggle";
|
||||
import NumericalInput from "../shared/numerical_input";
|
||||
import SkillSelector from "../skills/SkillSelector";
|
||||
import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers";
|
||||
|
|
@ -195,15 +196,18 @@ export const KeyMetadataField = ({ form }: { form: KeyEditForm }) => {
|
|||
description={t("keyDetail.keyEditViewControls.metadata.description")}
|
||||
>
|
||||
{(field) => (
|
||||
<Textarea
|
||||
{...field}
|
||||
value={(field.value as string | undefined) ?? ""}
|
||||
rows={10}
|
||||
onBlur={() => {
|
||||
field.onBlur();
|
||||
moveMetadataTagsToTagsField(form, t);
|
||||
}}
|
||||
/>
|
||||
<>
|
||||
<StorePromptsMetadataToggle value={field.value as string | undefined} onChange={field.onChange} />
|
||||
<Textarea
|
||||
{...field}
|
||||
value={(field.value as string | undefined) ?? ""}
|
||||
rows={10}
|
||||
onBlur={() => {
|
||||
field.onBlur();
|
||||
moveMetadataTagsToTagsField(form, t);
|
||||
}}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
</FormField>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -358,6 +358,24 @@ describe("KeyEditView", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("submits the spend prompt storage restriction in metadata", async () => {
|
||||
const onSubmit = vi.fn().mockResolvedValue(undefined);
|
||||
renderWithProviders(
|
||||
<KeyEditView
|
||||
keyData={MOCK_KEY_DATA}
|
||||
onCancel={() => {}}
|
||||
onSubmit={onSubmit}
|
||||
accessToken=""
|
||||
userID=""
|
||||
userRole="Admin"
|
||||
/>,
|
||||
);
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Store prompts in spend logs" }));
|
||||
fireEvent.click(screen.getByText("Save Changes"));
|
||||
await waitFor(() => expect(onSubmit).toHaveBeenCalled());
|
||||
expect(JSON.parse(onSubmit.mock.calls[0][0].metadata).store_prompts_in_spend_logs).toBe(false);
|
||||
});
|
||||
|
||||
it("should render tags", async () => {
|
||||
renderWithProviders(
|
||||
<KeyEditView
|
||||
|
|
|
|||
|
|
@ -3081,6 +3081,10 @@
|
|||
"empty": "No teams found",
|
||||
"loading": "Loading teams…"
|
||||
},
|
||||
"storePromptsToggle": {
|
||||
"label": "Store prompts in spend logs",
|
||||
"description": "Turn off to exclude prompts and responses from spend logs. Global and team restrictions still apply."
|
||||
},
|
||||
"accessGroupSelector": {
|
||||
"placeholder": "Select access groups",
|
||||
"label": "Access Group",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue