diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs index e5d4ce3a708..d95a50a5c68 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs @@ -367,6 +367,7 @@ mod tests { end_time: 1.0, stream: false, metadata: StandardLoggingMetadata::default(), + hidden_params: Default::default(), messages: None, }); diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs index 792717dacfc..d39f7934787 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs @@ -211,6 +211,7 @@ mod tests { user_api_key_team_id: Some("team".to_string()), ..Default::default() }, + hidden_params: Default::default(), messages: Some(json!([{"role": "user", "content": "read this"}])), } } diff --git a/litellm-rust/crates/ai-gateway/src/integrations/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/types.rs index 34dce93d8e0..0bcaf096eb8 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/types.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/types.rs @@ -58,6 +58,8 @@ pub struct StandardLoggingPayload { pub metadata: StandardLoggingMetadata, + pub hidden_params: HashMap, + /// Optional; stored as request input on the spend log row. #[serde(skip_serializing_if = "Option::is_none")] pub messages: Option, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 8f5fd0d38af..4df26f336d3 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -5,6 +5,7 @@ use std::pin::Pin; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrAuthStrategy; +use litellm_core::ocr::types::litellm_rust_hidden_params; use litellm_core::CoreResult; use serde_json::{json, Map, Value}; @@ -168,6 +169,7 @@ impl OcrLifecycleHooks { user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), ..Default::default() }, + hidden_params: litellm_rust_hidden_params().into_iter().collect(), messages: None, } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 699a0a4066c..299d3bdbcc2 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -351,6 +351,7 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { .expect("ocr request succeeds"); assert_eq!(response["pages"][0]["markdown"], "ok"); + assert_eq!(response["_hidden_params"]["litellm_rust"], true); assert_eq!( guardrail.events(), vec!["async_pre_call_hook", "async_moderation_hook"] diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs index c32e727de54..9ebfa3a7d70 100644 --- a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -178,6 +178,7 @@ impl RealTimeStreaming { user_api_key_team_id: self.metadata.user_api_key_team_id.clone(), ..Default::default() }, + hidden_params: Default::default(), messages: None, } } diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 1a72b8f1d66..ba4d08206da 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,5 +1,5 @@ use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{json, Map, Value}; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct OcrRequestData { @@ -24,6 +24,32 @@ impl OcrResponseData { "document_annotation": self.document_annotation, "usage_info": self.usage_info, "object": self.object, + "_hidden_params": Value::Object(litellm_rust_hidden_params()), }) } } + +pub fn litellm_rust_hidden_params() -> Map { + let mut hidden_params = Map::new(); + hidden_params.insert("litellm_rust".to_string(), json!(true)); + hidden_params +} + +#[cfg(test)] +mod tests { + use super::OcrResponseData; + + #[test] + fn ocr_response_data_marks_litellm_rust_hidden_params() { + let response = OcrResponseData { + pages: vec![], + model: "mistral-ocr-latest".to_string(), + document_annotation: None, + usage_info: None, + object: "ocr".to_string(), + } + .into_json(); + + assert_eq!(response["_hidden_params"]["litellm_rust"], true); + } +} diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e860f48a7bd..8300d460b07 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5342,6 +5342,7 @@ class StandardLoggingPayloadSetup: batch_models=None, litellm_model_name=None, usage_object=None, + litellm_rust=None, ) if hidden_params is not None: for key in StandardLoggingHiddenParams.__annotations__.keys(): diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index a2946c62506..3668c64defa 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -77,6 +77,13 @@ class OCRResponse(LiteLLMPydanticObjectBase): # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) + def model_post_init(self, __context: Any) -> None: + extra_fields = getattr(self, "__pydantic_extra__", None) + if isinstance(extra_fields, dict): + hidden_params = extra_fields.pop("_hidden_params", None) + if isinstance(hidden_params, dict): + self._hidden_params.update(hidden_params) + class OCRRequestData(LiteLLMPydanticObjectBase): """OCR request data structure.""" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index be5d7a2db78..4e0b2e318f5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3197,6 +3197,7 @@ class SpendLogsMetadata(TypedDict): str ] # S3/GCS object key for cold storage retrieval litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds + litellm_rust: Optional[bool] # True when the request used the LiteLLM Rust path attempted_retries: Optional[ int ] # Number of retries attempted (0 = first attempt succeeded) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f6e0303e349..25078d368b3 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -78,6 +78,7 @@ def _get_spend_logs_metadata( model_map_information: Optional[StandardLoggingModelInformation] = None, cold_storage_object_key: Optional[str] = None, litellm_overhead_time_ms: Optional[float] = None, + litellm_rust: Optional[bool] = None, cost_breakdown: Optional[CostBreakdown] = None, litellm_call_id: Optional[str] = None, ) -> SpendLogsMetadata: @@ -107,6 +108,7 @@ def _get_spend_logs_metadata( eval_information=None, cold_storage_object_key=cold_storage_object_key, litellm_overhead_time_ms=None, + litellm_rust=litellm_rust, attempted_retries=None, max_retries=None, cost_breakdown=None, @@ -134,6 +136,7 @@ def _get_spend_logs_metadata( clean_metadata["model_map_information"] = model_map_information clean_metadata["cold_storage_object_key"] = cold_storage_object_key clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms + clean_metadata["litellm_rust"] = litellm_rust clean_metadata["cost_breakdown"] = cost_breakdown clean_metadata["litellm_call_id"] = litellm_call_id @@ -381,6 +384,11 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs else None ), litellm_overhead_time_ms=litellm_overhead_time_ms, + litellm_rust=( + standard_logging_payload.get("hidden_params", {}).get("litellm_rust", None) + if standard_logging_payload is not None + else None + ), cost_breakdown=( standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c3f99b1d18b..26d6058c7b5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2728,6 +2728,7 @@ class StandardLoggingHiddenParams(TypedDict): batch_models: Optional[List[str]] litellm_model_name: Optional[str] # the model name sent to the provider by litellm usage_object: Optional[dict] + litellm_rust: Optional[bool] class StandardLoggingModelInformation(TypedDict): diff --git a/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py b/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py index bd7a7712572..8eee32a73b8 100644 --- a/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py +++ b/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py @@ -60,14 +60,11 @@ class OCRCustomGuardrail(CustomGuardrail): self, user_api_key_dict, cache, data: dict[str, Any], call_type: str ) -> None: self.calls.append({"data": data, "call_type": call_type}) - data["document"] = { - **data["document"], - "guardrail_executed": True, - } data["optional_params"] = { **data["optional_params"], "include_image_base64": True, } + data["guardrail_executed"] = True async def async_log_success_event( self, @@ -114,6 +111,7 @@ async def test_rust_ocr_executes_custom_logger_from_callback_manager(): ) assert len(response.pages) > 0 + assert response._hidden_params["litellm_rust"] is True assert custom_logger.response_obj["object"] == "ocr" assert len(custom_logger.response_obj["value"]["pages"]) > 0 assert isinstance(custom_logger.start_time, datetime) @@ -128,6 +126,7 @@ async def test_rust_ocr_executes_custom_logger_from_callback_manager(): assert logged_payload["custom_llm_provider"] == "mistral" assert logged_payload["call_type"] == "ocr" assert logged_payload["response_cost"] == 0.0 + assert logged_payload["hidden_params"]["litellm_rust"] is True @pytest.mark.asyncio @@ -144,8 +143,9 @@ async def test_rust_ocr_executes_custom_guardrail_from_callback_manager(): ) assert len(response.pages) > 0 + assert response._hidden_params["litellm_rust"] is True assert custom_guardrail.calls[0]["call_type"] == "ocr" - assert custom_guardrail.calls[0]["data"]["document"]["guardrail_executed"] is True + assert custom_guardrail.calls[0]["data"]["guardrail_executed"] is True assert ( custom_guardrail.calls[0]["data"]["optional_params"]["include_image_base64"] is True