From 56e429e33d036548e75ce4ded5e9a782c97b9684 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 1 Oct 2025 23:38:53 +0530 Subject: [PATCH] refactor code for better handling cost --- litellm/proxy/common_request_processing.py | 198 +++++++++++++-------- 1 file changed, 119 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e7bad1e19dd..091077b6e7e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -46,6 +46,7 @@ if TYPE_CHECKING: else: ProxyConfig = Any from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: @@ -760,85 +761,8 @@ class ProxyBaseLLMRequestProcessing: str_so_far += response_str # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider - # Handle both dict SSE events and pre-formatted string SSE lines - if getattr(litellm, "include_cost_in_streaming_usage", False) is True: - try: - def _inject_cost_into_usage_dict(obj: dict) -> Optional[dict]: - if ( - obj.get("type") == "message_delta" - and isinstance(obj.get("usage"), dict) - ): - _usage = obj["usage"] - prompt_tokens = int(_usage.get("input_tokens", 0) or 0) - completion_tokens = int(_usage.get("output_tokens", 0) or 0) - total_tokens = int( - _usage.get("total_tokens", prompt_tokens + completion_tokens) - or (prompt_tokens + completion_tokens) - ) - - _mr = ModelResponse( - usage=Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=total_tokens, - ) - ) - model_name = request_data.get("model", "") - try: - cost_val = litellm.completion_cost( - completion_response=_mr, - model=model_name, - ) - except Exception: - cost_val = None - if cost_val is not None: - obj.setdefault("usage", {})["cost"] = cost_val - return obj - return None - - def _inject_cost_into_sse_frame_str(frame_str: str) -> Optional[str]: - # frame_str may contain multiple lines like 'event: ...\ndata: {...}\n\n' - # We only modify the JSON in the 'data:' line - try: - # Split preserving lines - lines = frame_str.split("\n") - for idx, ln in enumerate(lines): - stripped_ln = ln.strip() - if stripped_ln.startswith("data:"): - json_part = stripped_ln.split("data:", 1)[1].strip() - if json_part and json_part != "[DONE]": - obj = json.loads(json_part) - maybe_modified = _inject_cost_into_usage_dict(obj) - if maybe_modified is not None: - # Replace just this line with updated JSON using safe_dumps - lines[idx] = f"data: {safe_dumps(maybe_modified)}" - return "\n".join(lines) - return None - except Exception: - return None - - if isinstance(chunk, dict): - maybe_modified = _inject_cost_into_usage_dict(chunk) - if maybe_modified is not None: - chunk = maybe_modified - elif isinstance(chunk, (bytes, bytearray)): - # Decode to str, inject, and rebuild as bytes - try: - s = chunk.decode("utf-8", errors="ignore") - maybe_mod = _inject_cost_into_sse_frame_str(s) - if maybe_mod is not None: - chunk = (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") - except Exception: - pass - elif isinstance(chunk, str): - # Try to parse SSE frame and inject cost into the data line - maybe_mod = _inject_cost_into_sse_frame_str(chunk) - if maybe_mod is not None: - # Ensure trailing frame separator - chunk = maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") - except Exception: - # Never break streaming on optional cost injection - pass + model_name = request_data.get("model", "") + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) # Format chunk using helper function yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) @@ -871,3 +795,119 @@ class ProxyBaseLLMRequestProcessing: ) error_returned = json.dumps({"error": proxy_exception.to_dict()}) yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n" + + @staticmethod + def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: + """ + Process a streaming chunk and inject cost information if enabled. + + Args: + chunk: The streaming chunk (dict, str, bytes, or bytearray) + model_name: Model name for cost calculation + + Returns: + The processed chunk with cost information injected if applicable + """ + if not getattr(litellm, "include_cost_in_streaming_usage", False): + return chunk + + try: + if isinstance(chunk, dict): + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + if maybe_modified is not None: + return maybe_modified + elif isinstance(chunk, (bytes, bytearray)): + # Decode to str, inject, and rebuild as bytes + try: + s = chunk.decode("utf-8", errors="ignore") + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + if maybe_mod is not None: + return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + except Exception: + pass + elif isinstance(chunk, str): + # Try to parse SSE frame and inject cost into the data line + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + if maybe_mod is not None: + # Ensure trailing frame separator + return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") + except Exception: + # Never break streaming on optional cost injection + pass + + return chunk + + @staticmethod + def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> Optional[str]: + """ + Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. + + Args: + frame_str: SSE frame string that may contain multiple lines + model_name: Model name for cost calculation + + Returns: + Modified SSE frame string with cost injected, or None if no modification needed + """ + try: + # Split preserving lines + lines = frame_str.split("\n") + for idx, ln in enumerate(lines): + stripped_ln = ln.strip() + if stripped_ln.startswith("data:"): + json_part = stripped_ln.split("data:", 1)[1].strip() + if json_part and json_part != "[DONE]": + obj = json.loads(json_part) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) + return None + except Exception: + return None + + @staticmethod + def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]: + """ + Inject cost information into a usage dictionary for message_delta events. + + Args: + obj: Dictionary containing the SSE event data + model_name: Model name for cost calculation + + Returns: + Modified dictionary with cost injected, or None if no modification needed + """ + if ( + obj.get("type") == "message_delta" + and isinstance(obj.get("usage"), dict) + ): + _usage = obj["usage"] + prompt_tokens = int(_usage.get("input_tokens", 0) or 0) + completion_tokens = int(_usage.get("output_tokens", 0) or 0) + total_tokens = int( + _usage.get("total_tokens", prompt_tokens + completion_tokens) + or (prompt_tokens + completion_tokens) + ) + + _mr = ModelResponse( + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + ) + + try: + cost_val = litellm.completion_cost( + completion_response=_mr, + model=model_name, + ) + except Exception: + cost_val = None + + if cost_val is not None: + obj.setdefault("usage", {})["cost"] = cost_val + return obj + return None \ No newline at end of file