From 97dbd2dfbf3de0386efbe77057bbf75bc1aa1336 Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 22:16:19 +0530 Subject: [PATCH 01/33] fix(gemini): preserve candidates with finishReason and no content (#40477) --- litellm/litellm_core_utils/core_helpers.py | 1 + .../adapters/transformation.py | 2 + .../vertex_and_google_ai_studio_gemini.py | 30 ++-- .../transformation.py | 20 ++- ...test_vertex_and_google_ai_studio_gemini.py | 146 ++++++++++++++++++ 5 files changed, 184 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index aa7d6ca1699..66180b165f8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -224,6 +224,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "IMAGE_PROHIBITED_CONTENT": "content_filter", "TOO_MANY_TOOL_CALLS": "stop", "MALFORMED_RESPONSE": "stop", + "NO_IMAGE": "content_filter", # Zhipu GLM "network_error": "stop", "sensitive": "content_filter", diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 8ff9f2e0679..f4cc569bcef 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1367,6 +1367,8 @@ class LiteLLMAnthropicMessagesAdapter: return "max_tokens" elif openai_finish_reason == "tool_calls": return "tool_use" + elif openai_finish_reason in ["content_filter", "refusal"]: + return "refusal" return "end_turn" @staticmethod diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d113b2b4f6b..01d1f063b1b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1340,6 +1340,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT", "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", + "NO_IMAGE", } ) @@ -2224,22 +2225,23 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): grounding_metadata: Final[list[dict]] = [] url_context_metadata: Final[list[dict]] = [] - image_response: list[ImageURLListItem] | None = None safety_ratings: Final[list] = [] citation_metadata: Final[list] = [] - chat_completion_message: Final[ChatCompletionResponseMessage] = {"role": "assistant"} - chat_completion_logprobs: ChoiceLogprobs | None = None - tools: list[ChatCompletionToolCallChunk] | None = [] - functions: ChatCompletionToolCallFunctionChunk | None = None - thinking_blocks: list[ChatCompletionThinkingBlock] | None = None - reasoning_content: str | None = None - thought_signatures: Sequence[str] | None = None - server_side_tool_invocations: list[dict[str, object]] | None = None for idx, candidate in enumerate(_candidates): - if "content" not in candidate: + if "content" not in candidate and "finishReason" not in candidate: continue + image_response: list[ImageURLListItem] | None = None + chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} + chat_completion_logprobs: ChoiceLogprobs | None = None + tools: list[ChatCompletionToolCallChunk] | None = [] + functions: ChatCompletionToolCallFunctionChunk | None = None + thinking_blocks: list[ChatCompletionThinkingBlock] | None = None + reasoning_content: str | None = None + thought_signatures: Sequence[str] | None = None + server_side_tool_invocations: list[dict[str, object]] | None = None + # Extract metadata using helper function ( candidate_grounding_metadata, @@ -2253,7 +2255,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): safety_ratings.extend(candidate_safety_ratings) citation_metadata.extend(candidate_citation_metadata) - if "parts" in candidate["content"]: + if "content" in candidate and candidate["content"] and "parts" in candidate["content"]: ( content, reasoning_content, @@ -2348,6 +2350,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_invocation_fields["server_side_tool_invocations"] = server_side_tool_invocations chat_completion_message["provider_specific_fields"] = tool_invocation_fields + if candidate.get("finishReason"): + finish_reason_fields = chat_completion_message.get("provider_specific_fields") or {} + finish_reason_fields["native_finish_reason"] = candidate.get("finishReason") + chat_completion_message["provider_specific_fields"] = finish_reason_fields + if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, @@ -2368,6 +2375,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): message=chat_completion_message, logprobs=chat_completion_logprobs, enhancements=None, + provider_specific_fields=chat_completion_message.get("provider_specific_fields"), ) model_response.choices.append(choice) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..13119085e46 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2272,13 +2272,27 @@ class LiteLLMCompletionResponsesConfig: if choices and len(choices) > 0: finish_reason = choices[0].finish_reason + status: Final[ResponsesAPIStatus] = ( + LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + finish_reason + ) + ) + incomplete_details = getattr(chat_completion_response, "incomplete_details", None) + if incomplete_details is None and status == "incomplete": + from openai.types.responses.response import IncompleteDetails + + if finish_reason == "length": + incomplete_details = IncompleteDetails(reason="max_output_tokens") + elif finish_reason in ["content_filter", "refusal"]: + incomplete_details = IncompleteDetails(reason="content_filter") + responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse( id=chat_completion_response.id, created_at=chat_completion_response.created, model=chat_completion_response.model, object="response", error=getattr(chat_completion_response, "error", None), - incomplete_details=getattr(chat_completion_response, "incomplete_details", None), + incomplete_details=incomplete_details, instructions=getattr(chat_completion_response, "instructions", None), metadata=getattr(chat_completion_response, "metadata", {}), output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output( @@ -2296,9 +2310,7 @@ class LiteLLMCompletionResponsesConfig: max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), previous_response_id=getattr(chat_completion_response, "previous_response_id", None), reasoning=None, - status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - finish_reason - ), + status=status, text={}, truncation=getattr(chat_completion_response, "truncation", None), usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 101f6e6fa5d..a048ad4f171 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5836,3 +5836,149 @@ def test_supported_reasoning_efforts_still_map(model): drop_params=False, ) assert "thinkingConfig" in result + + +def test_gemini_candidate_with_finish_reason_no_content_chat_completion(): + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + model_response = ModelResponse() + logging_obj = MagicMock() + raw_response = MagicMock() + raw_response.headers = {} + + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=model_response, + model="gemini-2.5-flash-image", + logging_obj=logging_obj, + raw_response=raw_response, + ) + assert len(resp.choices) == 1 + assert resp.choices[0].finish_reason == "content_filter" + assert resp.choices[0].message.content is None + assert resp.choices[0].provider_specific_fields["native_finish_reason"] == "NO_IMAGE" + + +def test_gemini_candidate_with_finish_reason_no_content_anthropic_messages(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_resp = adapter.translate_openai_response_to_anthropic( + response=resp, + tool_name_mapping={}, + ) + assert anthropic_resp["stop_reason"] == "refusal" + assert anthropic_resp["content"] == [] + + +def test_gemini_candidate_with_finish_reason_no_content_responses_api(): + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + responses_resp = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Generate picture", + responses_api_request={}, + chat_completion_response=resp, + ) + assert responses_resp.status == "incomplete" + assert responses_resp.incomplete_details is not None + assert responses_resp.incomplete_details.reason == "content_filter" + + +def test_gemini_candidate_other_finish_reasons_no_content(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + max_tokens_response = { + "candidates": [{"finishReason": "MAX_TOKENS", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 50, "totalTokenCount": 60}, + } + resp_length = config._transform_google_generate_content_to_openai_model_response( + completion_response=max_tokens_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + assert len(resp_length.choices) == 1 + assert resp_length.choices[0].finish_reason == "length" + assert resp_length.choices[0].provider_specific_fields["native_finish_reason"] == "MAX_TOKENS" + + anthropic_length = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( + response=resp_length, + tool_name_mapping={}, + ) + assert anthropic_length["stop_reason"] == "max_tokens" + + responses_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="thinking request", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert responses_length.status == "incomplete" + assert responses_length.incomplete_details.reason == "max_output_tokens" + From cd66b34b45dace036126f4ea809df132407d907c Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 22:48:08 +0530 Subject: [PATCH 02/33] style(responses): apply ruff formatting to transformation.py --- .../litellm_completion_transformation/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 13119085e46..10e56e85bff 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2273,9 +2273,7 @@ class LiteLLMCompletionResponsesConfig: finish_reason = choices[0].finish_reason status: Final[ResponsesAPIStatus] = ( - LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - finish_reason - ) + LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(finish_reason) ) incomplete_details = getattr(chat_completion_response, "incomplete_details", None) if incomplete_details is None and status == "incomplete": From e54399eff7a13e649b6a353486922166b1bf5688 Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 23:07:31 +0530 Subject: [PATCH 03/33] test: add direct coverage for content_filter and refusal in anthropic and responses adapters --- ...al_pass_through_adapters_transformation.py | 40 ++++++++++++ .../test_litellm_completion_responses.py | 62 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 03b9840b1c3..0e09e27f4db 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -102,6 +102,46 @@ def test_translate_chat_length_takes_precedence_over_refusal(): assert result.get("stop_details") is None +def test_translate_chat_content_filter_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-content-filter", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="content_filter", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + +def test_translate_chat_refusal_finish_reason_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-refusal-reason", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="refusal", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + def test_translate_streaming_openai_chunk_to_anthropic_content_block(): choices = [ StreamingChoices( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 46249e50572..f2ebbf316c3 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4246,3 +4246,65 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +def test_transform_chat_completion_response_incomplete_details(): + from openai.types.responses.response import IncompleteDetails + + resp_length = ModelResponse( + id="resp-length", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + result_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert result_length.status == "incomplete" + assert result_length.incomplete_details is not None + assert result_length.incomplete_details.reason == "max_output_tokens" + + resp_filter = ModelResponse( + id="resp-filter", + choices=[Choices(index=0, finish_reason="content_filter", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_filter = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_filter, + ) + assert result_filter.status == "incomplete" + assert result_filter.incomplete_details is not None + assert result_filter.incomplete_details.reason == "content_filter" + + resp_refusal = ModelResponse( + id="resp-refusal", + choices=[Choices(index=0, finish_reason="refusal", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_refusal = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_refusal, + ) + assert result_refusal.status == "incomplete" + assert result_refusal.incomplete_details is not None + assert result_refusal.incomplete_details.reason == "content_filter" + + existing_details = IncompleteDetails(reason="content_filter") + resp_existing = ModelResponse( + id="resp-existing", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + resp_existing.incomplete_details = existing_details + result_existing = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_existing, + ) + assert result_existing.status == "incomplete" + assert result_existing.incomplete_details == existing_details + From 9fa85c5da5a6a9497a0df12cbc02866497af353d Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 08:39:01 +0000 Subject: [PATCH 04/33] fix(proxy): name the blocking guardrail in x-litellm-applied-guardrails When a guardrail hook raises, the common ProxyLogging dispatch (sequential and parallel pre_call, pipeline block, during_call and post_call metrics wrapper, streaming iterator wrapper) now records that guardrail in applied_guardrails before re-raising, and pre_call_hook folds request-declared guardrails in on its raising path. Buffered streams rebuild their response headers after the first chunk so a post_call block reached while buffering carries the blocker too Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 11 ++-- litellm/proxy/utils.py | 58 +++++++++++++++---- .../proxy/test_common_request_processing.py | 57 ++++++++++++++++++ .../proxy_logging/test_during_call_hook.py | 18 ++++++ .../proxy_logging/test_guardrail_pipeline.py | 8 ++- .../test_post_call_success_hook.py | 24 ++++++++ .../utils/proxy_logging/test_pre_call_hook.py | 40 +++++++++++++ .../proxy_logging/test_streaming_hooks.py | 38 +++++++++++- 8 files changed, 234 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2f39e6c71bc..9a038ca79ba 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2576,11 +2576,14 @@ class ProxyBaseLLMRequestProcessing: ) async def refresh_stream_headers() -> Mapping[str, str]: - """`custom_headers` rebuilt for whichever deployment served the stream.""" - if not getattr(response, "fallback_headers_adopted", False): - return custom_headers + """`custom_headers` rebuilt once the first chunk is buffered, from `self.data` as the + guardrails left it and for whichever deployment served the stream.""" return self._stream_response_headers( - hidden_params=get_hidden_params_dict(response), + hidden_params=( + get_hidden_params_dict(response) + if getattr(response, "fallback_headers_adopted", False) + else hidden_params + ), user_api_key_dict=user_api_key_dict, logging_obj=logging_obj, version=version, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8225fef3492..fea6ce20b61 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -123,6 +123,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( @@ -437,6 +438,12 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: detail.setdefault("guardrail_mode", event_hook) +def _record_raising_guardrail(request_data: Mapping[str, object], callback: object) -> None: + guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) + if isinstance(request_data, dict) and isinstance(guardrail_name, str): + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=guardrail_name) + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -1795,13 +1802,19 @@ class ProxyLogging: ) if expected_if_unmutated is not None: callback.mark_pre_call_hook_ran(expected_if_unmutated) - result: Final = await self._process_guardrail_callback( - callback=callback, - data=input_data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - event_type=GuardrailEventHooks.pre_call, - ) + try: + result: Final = await self._process_guardrail_callback( + callback=callback, + data=input_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + event_type=GuardrailEventHooks.pre_call, + ) + except SensitiveDataRouteException: + raise + except Exception: + _record_raising_guardrail(data, callback) + raise if ( scans_raw_request and expected_if_unmutated is not None @@ -2031,6 +2044,7 @@ class ProxyLogging: callback: Final = PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) if callback is not None: _enrich_http_exception_with_guardrail_context(original_exception, callback) + _record_raising_guardrail(data, callback) raise original_exception step_results_serializable: Final = [ @@ -2296,8 +2310,10 @@ class ProxyLogging: if data is not None: self._process_guardrail_metadata(data) return data - except Exception as e: - raise e + except Exception: + if data is not None: + self._process_guardrail_metadata(data) + raise async def _run_parallel_pre_call_guardrails( self, @@ -2355,6 +2371,8 @@ class ProxyLogging: # live kwargs. if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: callback.mark_pre_call_hook_ran(data) + if isinstance(result, BaseException) and not isinstance(result, SensitiveDataRouteException): + _record_raising_guardrail(data, callback) raised: Final = tuple(result for result in results if isinstance(result, BaseException)) blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) if blocking is not None: @@ -2433,7 +2451,12 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics(callback: object, coro: Awaitable[_T], hook_type: str) -> _T: + async def _run_guardrail_with_metrics( + callback: object, + coro: Awaitable[_T], + hook_type: str, + request_data: Mapping[str, object], + ) -> _T: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and @@ -2453,6 +2476,7 @@ class ProxyLogging: status = "error" error_type = type(e).__name__ _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise finally: ProxyLogging._emit_guardrail_metrics( @@ -2465,7 +2489,9 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: object, gen: AsyncGenerator[_T, None] + callback: object, + gen: AsyncGenerator[_T, None], + request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: """ Yield from `gen`; if iteration raises an HTTPException with dict detail, @@ -2480,6 +2506,7 @@ class ProxyLogging: yield chunk except Exception as e: _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise # Cache for callback-capability detection. Keyed on a signature of @@ -2714,6 +2741,7 @@ class ProxyLogging: call_type=call_type, ), "during_call", + request_data=data, ) return await self._run_guardrail_with_metrics( @@ -2724,6 +2752,7 @@ class ProxyLogging: call_type=call_type, ), "during_call", + request_data=data, ) async def failed_tracking_alert( @@ -3242,6 +3271,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) else: guardrail_response = await self._run_guardrail_with_metrics( @@ -3252,6 +3282,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) if guardrail_response is not None: @@ -3315,6 +3346,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) else: await self._run_guardrail_with_metrics( @@ -3325,6 +3357,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) results: Final = await asyncio.gather( @@ -3388,6 +3421,7 @@ class ProxyLogging: request_data=request_data, ), "post_mcp_call", + request_data=request_data, ) return response @@ -3637,6 +3671,7 @@ class ProxyLogging: response=current_response, request_data=request_data, ), + request_data=request_data, ) else: # kind == "apply_guardrail": route through unified_guardrail @@ -3649,6 +3684,7 @@ class ProxyLogging: guardrail_to_apply=resolved_callback, buffer_until_moderated_default=(kind == "override"), ), + request_data=request_data, ) pipeline_translation: Final = ( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 4ac687625c2..028ea29c093 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -40,9 +40,12 @@ from litellm.proxy.common_request_processing import ( _parse_event_data_for_error, _resolve_per_request_model_group_alias, _should_return_raw_model_name, + _sse_error_frames, _UpstreamClosingStreamingResponse, create_response, + sse_error_payload, ) +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -8952,6 +8955,60 @@ class TestStreamingResponseHeadersFollowFallback: assert "llm_provider-stale-marker" not in result.headers assert result.headers["x-callback-header"] == "kept" + @pytest.mark.asyncio + async def test_streaming_block_headers_name_the_blocking_guardrail(self, monkeypatch): + processor_data: dict[str, object] = {"model": "oa", "stream": True, "metadata": {}} + + def select_data_generator(**kwargs): + async def generator(): + add_guardrail_to_applied_guardrails_header(processor_data, "stream-blocker") + _, error_obj = sse_error_payload(HTTPException(status_code=400, detail="blocked")) + for frame in _sse_error_frames(error_obj): + yield frame + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-7144-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + processor_data["litellm_logging_obj"] = logging_obj + processor = ProxyBaseLLMRequestProcessing(data=processor_data) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + async def fake_route_request(**kwargs): + async def call(): + return SimpleNamespace(_hidden_params={}, fallback_headers_adopted=False) + + return call() + + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, JSONResponse) + assert result.status_code == 400 + assert result.headers["x-litellm-applied-guardrails"] == "stream-blocker" + class _MessagesFallbackStream: def __init__(self) -> None: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py index 3c5d879c2dc..46f39ef6fb7 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py @@ -6,6 +6,7 @@ from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -84,3 +85,20 @@ async def test_during_call_hook_guardrail_error_raises(proxy_logging, make_user_ user_api_key_dict=make_user_api_key_auth(), call_type="completion", ) + + +@pytest.mark.asyncio +async def test_during_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + g = _make_guardrail("blocker") + g.async_moderation_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked")) + monkeypatch.setattr(litellm, "callbacks", [_make_guardrail("passer"), g]) + data = {"model": "m", "metadata": {}} + with pytest.raises(HTTPException): + await proxy_logging.during_call_hook( + data=data, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + assert "blocker" in data["metadata"]["applied_guardrails"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 077bf5a313e..cd2b7a278bb 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -527,17 +527,19 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): result.step_results = [MagicMock(guardrail_name="g")] result.original_exception = original + data: dict[str, object] = {"model": "m"} saved = litellm.callbacks litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") finally: litellm.callbacks = saved assert info.value is original assert info.value.detail["guardrail_name"] == "g" assert info.value.detail["guardrail_mode"] == GuardrailEventHooks.pre_call + assert data["metadata"] == {"applied_guardrails": ["g"]} def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route(): @@ -617,7 +619,7 @@ async def test_run_guardrail_with_metrics_passes_result_and_records_success(monk monkeypatch.setattr(litellm, "callbacks", [prom]) out = await ProxyLogging._run_guardrail_with_metrics( - callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call" + callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call", request_data={} ) assert out == {"a": 1, "b": 2, "c": 3} @@ -643,7 +645,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call", request_data={}) assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py index 715d66db181..53d8948869f 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -6,6 +6,7 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -96,3 +97,26 @@ async def test_post_call_success_hook_guardrail_returns_modified_response( data={}, response={"orig": True}, user_api_key_dict=make_user_api_key_auth() ) assert out == modified + + +@pytest.mark.asyncio +@pytest.mark.parametrize("run_in_parallel", [False, True], ids=["sequential", "parallel"]) +async def test_post_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch, run_in_parallel +): + def _passer_that_records(data, user_api_key_dict, response): + data["metadata"]["applied_guardrails"] = ["passer"] + + passer = _make_guardrail("passer") + passer.async_post_call_success_hook = AsyncMock(side_effect=_passer_that_records) + passer.run_in_parallel = run_in_parallel + blocker = _make_guardrail("blocker") + blocker.async_post_call_success_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked")) + blocker.run_in_parallel = run_in_parallel + monkeypatch.setattr(litellm, "callbacks", [passer, blocker]) + data = {"model": "m", "metadata": {}} + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=MagicMock(), user_api_key_dict=make_user_api_key_auth() + ) + assert data["metadata"]["applied_guardrails"] == ["passer", "blocker"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 6e5cb7fcae3..dbc6fba4ab1 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -905,3 +905,43 @@ async def test_scan_raw_request_warns_on_in_place_mutation_returning_none( ) mock_logger.warning.assert_called_once() assert "scan_raw_request" in str(mock_logger.warning.call_args) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "blocker_kwargs", + [ + pytest.param({}, id="sequential"), + pytest.param({"scan_raw_request": True}, id="scan_raw_request"), + pytest.param({"run_in_parallel": True}, id="parallel"), + ], +) +async def test_pre_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch, blocker_kwargs +): + monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(**blocker_kwargs)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data = _secret_request() + with pytest.raises(HTTPException, match="blocked"): + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["blocker"] + + +@pytest.mark.asyncio +async def test_pre_call_block_keeps_request_declared_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(default_on=False)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data = {**_secret_request(), "metadata": {"guardrails": ["blocker", "declared-post-call"]}} + with pytest.raises(HTTPException, match="blocked"): + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["blocker", "declared-post-call"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ebc831b4102..52586ed2174 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -20,6 +20,7 @@ from fastapi import HTTPException import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( @@ -27,6 +28,7 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Usage @@ -175,7 +177,7 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro yield ch cb = MagicMock(guardrail_name="g", event_hook="pre_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen()) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen(), request_data={}) out = [ch async for ch in wrapped] snapshot = { "chunks": out, @@ -201,7 +203,7 @@ async def test_wrap_streaming_iterator_with_enrichment_enriches_http_exception_r raise HTTPException(status_code=400, detail=detail) cb = MagicMock(guardrail_name="presidio", event_hook="post_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen()) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen(), request_data={}) with pytest.raises(HTTPException): async for _ in wrapped: pass @@ -696,3 +698,35 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log data={}, user_api_key_dict=make_user_api_key_auth(), response=response ) assert out == {} + + +@pytest.mark.asyncio +async def test_stream_guardrail_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class _StreamBlocker(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="stream-blocker", event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for _ in response: + raise HTTPException(status_code=400, detail={"error": "blocked"}) + yield # pragma: no cover + + monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker()]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + async def upstream(): + yield "chunk" + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(HTTPException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] From 73dea4c5676218005a42bdbe841c76ce6eb42707 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 18:58:21 +0000 Subject: [PATCH 05/33] fix(proxy): attribute only the raising layer in stream and pipeline blocks The streaming wrapper caught every exception crossing its boundary and named its own callback, so a block by an inner guardrail or a provider stream failure also named every outer guardrail. The wrapper now runs the hook over an upstream boundary that remembers the exception it raised, and skips attribution when the same exception passes through Pipeline blocks converted from SensitiveDataRouteException or ModifyResponseException into a generic guardrail_pipeline_error now still record the blocking step's guardrail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 129 ++++++++++++----- .../proxy_logging/test_guardrail_pipeline.py | 26 +++- .../proxy_logging/test_streaming_hooks.py | 132 +++++++++++++++--- 3 files changed, 223 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fea6ce20b61..21e6ed1d24f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -12,13 +12,36 @@ import sys import threading import time import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Mapping, + Sequence, +) from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText +from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Final, + Generic, + Literal, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from typing_extensions import ReadOnly, TypedDict @@ -444,6 +467,33 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=guardrail_name) +class _UpstreamStreamBoundary(Generic[_T]): + """Remembers the exception the upstream iterator raised, so the wrapper around a + streaming hook can tell a pass-through failure from one the hook raised itself.""" + + __slots__ = ("_upstream", "failure") + + def __init__(self, upstream: AsyncIterable[_T]) -> None: + self._upstream: Final = upstream.__aiter__() + self.failure: BaseException | None = None + + def __aiter__(self) -> "_UpstreamStreamBoundary[_T]": + return self + + async def __anext__(self) -> _T: + try: + return await self._upstream.__anext__() + except StopAsyncIteration: + raise + except Exception as e: + self.failure = e + raise + + +class _StreamIteratorHook(Protocol[_T]): + def __call__(self, *, response: AsyncIterator[_T]) -> AsyncGenerator[_T, None]: ... + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -2037,14 +2087,18 @@ class ProxyLogging: _merge_pipeline_metadata_writes(data, result.modified_data) if result.terminal_action == "block": + blocking_step: Final = result.step_results[-1] if result.step_results else None + callback: Final = ( + PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) + if blocking_step is not None + else None + ) + if callback is not None: + _record_raising_guardrail(data, callback) original_exception: Final = result.original_exception if original_exception is not None and not _exception_changes_request_flow(original_exception): - blocking_step: Final = result.step_results[-1] if result.step_results else None - if blocking_step is not None: - callback: Final = PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) - if callback is not None: - _enrich_http_exception_with_guardrail_context(original_exception, callback) - _record_raising_guardrail(data, callback) + if callback is not None: + _enrich_http_exception_with_guardrail_context(original_exception, callback) raise original_exception step_results_serializable: Final = [ @@ -2490,23 +2544,26 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( callback: object, - gen: AsyncGenerator[_T, None], + response: AsyncIterable[_T], + hook: _StreamIteratorHook[_T], request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: """ - Yield from `gen`; if iteration raises an HTTPException with dict detail, - enrich the detail with the originating callback's `guardrail_name` and - `guardrail_mode` before re-raising. Used to wrap each layer of the - async_post_call_streaming_iterator_hook chain so the enrichment is - attributed to the callback that produced the chunk pipeline at that - point in the chain. + Run `hook` over `response` and yield its chunks. If the hook itself raises, + enrich an HTTPException's dict detail with the callback's `guardrail_name` + and `guardrail_mode` and record the callback in `applied_guardrails` before + re-raising. Failures raised by `response` (the provider stream or an inner + layer of the async_post_call_streaming_iterator_hook chain) pass through + untouched, so only the layer that actually raised is attributed. """ + upstream: Final = _UpstreamStreamBoundary(response) try: - async for chunk in gen: + async for chunk in hook(response=upstream): yield chunk except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) - _record_raising_guardrail(request_data, callback) + if e is not upstream.failure: + _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise # Cache for callback-capability detection. Keyed on a signature of @@ -3663,29 +3720,27 @@ class ProxyLogging: ) else kind ) - if effective_kind == "override": - current_response = self._wrap_streaming_iterator_with_enrichment( - resolved_callback, - resolved_callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=current_response, - request_data=request_data, - ), + hook: _StreamIteratorHook[object] = ( + partial( + resolved_callback.async_post_call_streaming_iterator_hook, + user_api_key_dict=user_api_key_dict, request_data=request_data, ) - else: - # kind == "apply_guardrail": route through unified_guardrail - current_response = self._wrap_streaming_iterator_with_enrichment( - resolved_callback, - unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - request_data=request_data, - response=current_response, - guardrail_to_apply=resolved_callback, - buffer_until_moderated_default=(kind == "override"), - ), + if effective_kind == "override" + else partial( + unified_guardrail.async_post_call_streaming_iterator_hook, + user_api_key_dict=user_api_key_dict, request_data=request_data, + guardrail_to_apply=resolved_callback, + buffer_until_moderated_default=(kind == "override"), ) + ) + current_response = self._wrap_streaming_iterator_with_enrichment( + resolved_callback, + current_response, + hook, + request_data=request_data, + ) pipeline_translation: Final = ( resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index cd2b7a278bb..5f3c09d9195 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -551,14 +551,23 @@ def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route(): session_id="sess-1", guardrail_name="pii-router", ) + cb = _make_guardrail() + cb.guardrail_name = "pii-router" result = MagicMock() result.terminal_action = "block" result.step_results = [MagicMock(guardrail_name="pii-router")] result.original_exception = original - with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + data: dict[str, object] = {"model": "m"} + saved = litellm.callbacks + litellm.callbacks = [cb] + try: + with pytest.raises(HTTPException) as info: + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") + finally: + litellm.callbacks = saved assert info.value.status_code == 400 assert info.value.detail["error"]["type"] == "guardrail_pipeline_error" + assert data["metadata"] == {"applied_guardrails": ["pii-router"]} def test_handle_pipeline_result_block_does_not_reraise_modify_response(): @@ -571,14 +580,23 @@ def test_handle_pipeline_result_block_does_not_reraise_modify_response(): request_data={"model": "m"}, guardrail_name="masker", ) + cb = _make_guardrail() + cb.guardrail_name = "masker" result = MagicMock() result.terminal_action = "block" result.step_results = [MagicMock(guardrail_name="masker")] result.original_exception = original - with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + data: dict[str, object] = {"model": "m"} + saved = litellm.callbacks + litellm.callbacks = [cb] + try: + with pytest.raises(HTTPException) as info: + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") + finally: + litellm.callbacks = saved assert info.value.status_code == 400 assert info.value.detail["error"]["type"] == "guardrail_pipeline_error" + assert data["metadata"] == {"applied_guardrails": ["masker"]} def test_handle_pipeline_result_modify_response_raises_modify_exception(): diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 52586ed2174..6fb000b4fa7 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -170,6 +170,15 @@ def test_init_response_taking_too_long_task_no_slack_instance_no_error_raises(pr # --------------------------------------------------------------------------- +async def _passthrough_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + + +async def _one_chunk() -> AsyncGenerator[object, None]: + yield "chunk" + + @pytest.mark.asyncio async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(proxy_logging): async def gen(): @@ -177,7 +186,9 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro yield ch cb = MagicMock(guardrail_name="g", event_hook="pre_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen(), request_data={}) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=gen(), hook=_passthrough_hook, request_data={} + ) out = [ch async for ch in wrapped] snapshot = { "chunks": out, @@ -197,18 +208,43 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro async def test_wrap_streaming_iterator_with_enrichment_enriches_http_exception_raises(proxy_logging): detail = {"error": "blocked"} - async def boom_gen(): + async def boom_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]: if False: yield # pragma: no cover raise HTTPException(status_code=400, detail=detail) cb = MagicMock(guardrail_name="presidio", event_hook="post_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen(), request_data={}) + request_data: dict[str, object] = {} + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=_one_chunk(), hook=boom_hook, request_data=request_data + ) with pytest.raises(HTTPException): async for _ in wrapped: pass assert detail["guardrail_name"] == "presidio" assert detail["guardrail_mode"] == "post_call" + assert request_data["metadata"]["applied_guardrails"] == ["presidio"] + + +@pytest.mark.asyncio +async def test_wrap_streaming_iterator_leaves_upstream_http_exception_unattributed(proxy_logging): + detail = {"error": "upstream rejected the stream"} + + async def failing_upstream() -> AsyncGenerator[object, None]: + if False: + yield # pragma: no cover + raise HTTPException(status_code=502, detail=detail) + + cb = MagicMock(guardrail_name="presidio", event_hook="post_call") + request_data: dict[str, object] = {} + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=failing_upstream(), hook=_passthrough_hook, request_data=request_data + ) + with pytest.raises(HTTPException): + async for _ in wrapped: + pass + assert detail == {"error": "upstream rejected the stream"} + assert request_data == {} # --------------------------------------------------------------------------- @@ -700,33 +736,83 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log assert out == {} +class _StreamBlocker(CustomGuardrail): + def __init__(self, guardrail_name: str = "stream-blocker") -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for _ in response: + raise HTTPException(status_code=400, detail={"error": "blocked"}) + yield # pragma: no cover + + +class _StreamPasser(CustomGuardrail): + def __init__(self, guardrail_name: str = "stream-passer") -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + + +async def _drain_stream_chain( + proxy_logging: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, + upstream: AsyncIterator[object], + request_data: dict[str, object], +) -> None: + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ): + pass + + +async def _failing_provider_stream() -> AsyncGenerator[object, None]: + yield "chunk" + raise RuntimeError("provider connection dropped") + + @pytest.mark.asyncio async def test_stream_guardrail_block_names_the_blocking_guardrail_in_applied_guardrails( proxy_logging, make_user_api_key_auth, monkeypatch ): - class _StreamBlocker(CustomGuardrail): - def __init__(self) -> None: - super().__init__(guardrail_name="stream-blocker", event_hook=GuardrailEventHooks.post_call, default_on=True) - - async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] - ) -> AsyncGenerator[object, None]: - async for _ in response: - raise HTTPException(status_code=400, detail={"error": "blocked"}) - yield # pragma: no cover - monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker()]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) - async def upstream(): - yield "chunk" - request_data: dict[str, object] = {"metadata": {}} with pytest.raises(HTTPException): - async for _ in proxy_logging.async_post_call_streaming_iterator_hook( - response=upstream(), - user_api_key_dict=make_user_api_key_auth(), - request_data=request_data, - ): - pass + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _one_chunk(), request_data) assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] + + +@pytest.mark.asyncio +async def test_stream_block_by_inner_guardrail_does_not_name_the_outer_layers( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker(), _StreamPasser("outer-a"), _StreamPasser("outer-b")]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(HTTPException) as info: + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _one_chunk(), request_data) + assert info.value.detail["guardrail_name"] == "stream-blocker" + assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] + + +@pytest.mark.asyncio +async def test_stream_provider_failure_is_not_attributed_to_any_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamPasser("outer-a"), _StreamPasser("outer-b")]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(RuntimeError, match="provider connection dropped"): + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _failing_provider_stream(), request_data) + assert request_data["metadata"] == {} From 0f3556b1bf505864b8860b41e5d917f38b7e5468 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 19:17:49 +0000 Subject: [PATCH 06/33] refactor(proxy): drop explanatory docstrings from the stream attribution helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 3 +-- litellm/proxy/utils.py | 11 ----------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9a038ca79ba..849211a2dae 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2576,8 +2576,7 @@ class ProxyBaseLLMRequestProcessing: ) async def refresh_stream_headers() -> Mapping[str, str]: - """`custom_headers` rebuilt once the first chunk is buffered, from `self.data` as the - guardrails left it and for whichever deployment served the stream.""" + """`custom_headers` rebuilt for whichever deployment served the stream.""" return self._stream_response_headers( hidden_params=( get_hidden_params_dict(response) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 21e6ed1d24f..4182a864e4c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -468,9 +468,6 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje class _UpstreamStreamBoundary(Generic[_T]): - """Remembers the exception the upstream iterator raised, so the wrapper around a - streaming hook can tell a pass-through failure from one the hook raised itself.""" - __slots__ = ("_upstream", "failure") def __init__(self, upstream: AsyncIterable[_T]) -> None: @@ -2548,14 +2545,6 @@ class ProxyLogging: hook: _StreamIteratorHook[_T], request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: - """ - Run `hook` over `response` and yield its chunks. If the hook itself raises, - enrich an HTTPException's dict detail with the callback's `guardrail_name` - and `guardrail_mode` and record the callback in `applied_guardrails` before - re-raising. Failures raised by `response` (the provider stream or an inner - layer of the async_post_call_streaming_iterator_hook chain) pass through - untouched, so only the layer that actually raised is attributed. - """ upstream: Final = _UpstreamStreamBoundary(response) try: async for chunk in hook(response=upstream): From 318b027782b0b6ab1aec7103a619d308661894e6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 08:26:39 +0000 Subject: [PATCH 07/33] fix(otel): keep caller traceparent and tracestate on pass-through relays Pass-through requests inject the proxy span into upstream headers since #40669, which replaced an explicit x-pass-traceparent with an unrelated trace and dropped its x-pass-tracestate. Keep the caller's context when the carrier already names a different trace, and keep the proxy child span for same-trace or missing headers. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 19 +++++- .../otel/test_otel_v2_components.py | 63 ++++++++++++++----- .../test_pass_through_endpoints.py | 59 +++++++++++++++++ 3 files changed, 125 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 1282e654365..19243d64c64 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -325,12 +325,27 @@ def _outgoing_trace_context(parent_span: object) -> Context | None: return None +def _propagated_context(headers: Mapping[str, str], request_context: Context) -> Context: + """``request_context`` when it continues the trace ``headers`` already name, else the + caller's own context, so an explicit upstream ``traceparent`` (``x-pass-traceparent``) + is never swapped for an unrelated trace and its ``tracestate`` survives.""" + caller: Final = extract_traceparent(headers) + if caller is None: + return request_context + caller_span: Final = get_current_span(caller).get_span_context() + request_span: Final = get_current_span(request_context).get_span_context() + if not caller_span.is_valid or caller_span.trace_id == request_span.trace_id: + return request_context + return caller + + def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]: """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span. Parent preference: ``parent_span`` (the request span auth stashed on the key), then the anchored request root span, then the ambient active span. Only trace context is - injected, never Baggage. Unchanged when no valid span exists anywhere. + injected, never Baggage. Unchanged when no valid span exists anywhere. A ``traceparent`` + already in ``headers`` from a different trace is forwarded as-is instead of replaced. """ context: Final = _outgoing_trace_context(parent_span) if context is None: @@ -338,7 +353,7 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS } - _PROPAGATOR.inject(carrier, context=context) + _PROPAGATOR.inject(carrier, context=_propagated_context(headers, context)) return carrier diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 72f6213e880..0c95049ce05 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -477,19 +477,21 @@ def _test_tracer(): return provider.get_tracer("test") +_CALLER_TRACEPARENT = "00-11111111111111111111111111111111-2222222222222222-01" + + def test_inject_trace_context_prefers_request_root_span(): def run(): tracer = _test_tracer() - with tracer.start_as_current_span("root") as root: + inbound = TraceContextTextMapPropagator().extract({"traceparent": _CALLER_TRACEPARENT}) + with tracer.start_as_current_span("root", context=inbound) as root: ctx_mod.set_request_root_span(root) - result = ctx_mod.inject_trace_context( - {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} - ) + result = ctx_mod.inject_trace_context({"traceparent": _CALLER_TRACEPARENT}) propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) return result, root, propagated result, root, propagated = ContextVarContext().run(run) - assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01" + assert result["traceparent"] != _CALLER_TRACEPARENT assert propagated.get_span_context().trace_id == root.get_span_context().trace_id assert propagated.get_span_context().span_id == root.get_span_context().span_id @@ -507,24 +509,57 @@ def test_inject_trace_context_uses_ambient_span_without_request_root(): assert propagated.get_span_context().span_id == ambient.get_span_context().span_id -def test_inject_trace_context_replaces_stale_trace_headers(): +def test_inject_trace_context_replaces_same_trace_headers_with_request_span(): def run(): tracer = _test_tracer() - with tracer.start_as_current_span("ambient") as ambient: - headers = { - "Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01", - "Tracestate": "vendor=old", - "x-keep": "1", - } + headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"} + inbound = TraceContextTextMapPropagator().extract({key.lower(): value for key, value in headers.items()}) + with tracer.start_as_current_span("ambient", context=inbound) as ambient: result = ctx_mod.inject_trace_context(headers) propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) return result, ambient, propagated result, ambient, propagated = ContextVarContext().run(run) assert sum(key.lower() == "traceparent" for key in result) == 1 - assert not any(key.lower() == "tracestate" for key in result) + assert sum(key.lower() == "tracestate" for key in result) == 1 assert result["x-keep"] == "1" - assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + assert result["tracestate"] == "vendor=caller" + assert propagated.get_span_context().span_id == ambient.get_span_context().span_id + + +def test_inject_trace_context_keeps_caller_traceparent_from_another_trace(): + def run(): + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient") as ambient: + ctx_mod.set_request_root_span(ambient) + headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"} + result = ctx_mod.inject_trace_context(headers, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, parent, propagated + + result, parent, propagated = ContextVarContext().run(run) + assert result["traceparent"] == _CALLER_TRACEPARENT + assert result["tracestate"] == "vendor=caller" + assert result["x-keep"] == "1" + assert sum(key.lower() == "traceparent" for key in result) == 1 + assert sum(key.lower() == "tracestate" for key in result) == 1 + assert propagated.get_span_context().trace_id != parent.get_span_context().trace_id + + +def test_inject_trace_context_replaces_malformed_caller_traceparent(): + def run(): + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient"): + headers = {"traceparent": "not-a-traceparent", "tracestate": "vendor=caller"} + result = ctx_mod.inject_trace_context(headers, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, parent, propagated + + result, parent, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().span_id == parent.get_span_context().span_id + assert "tracestate" not in result def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e3e7ad618e0..df4842a7887 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4328,6 +4328,65 @@ async def test_pass_through_request_propagates_active_trace_context(span_source: assert propagated.get_span_context().span_id == span.get_span_context().span_id +async def _relay_with_trace_headers(inbound_headers: dict[str, str], forward_headers: bool): + from opentelemetry.sdk.trace import TracerProvider + + captured: dict[str, httpx.Headers] = {} + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + captured["headers"] = upstream_request.headers + return httpx.Response(200, json={"ok": True}, request=upstream_request) + + fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None) + tracer = TracerProvider().get_tracer("test") + try: + with ExitStack() as stack: + _enter_relay_logging_mocks(stack, {}) + span = tracer.start_span("litellm_request") + stack.callback(span.end) + request = _relay_client_request(method="POST") + request.headers = Headers(inbound_headers) + response = await pass_through_request( + request=request, + target="http://internal-api.test/v1/generate", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span), + forward_headers=forward_headers, + ) + finally: + cleanup() + await fake_client.aclose() + assert response.status_code == 200 + return captured["headers"], span + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forward_headers", [False, True]) +async def test_pass_through_request_keeps_x_pass_trace_headers_when_otel_span_is_active(forward_headers: bool): + caller_traceparent = "00-11111111111111111111111111111111-2222222222222222-01" + + upstream_headers, span = await _relay_with_trace_headers( + {"x-pass-traceparent": caller_traceparent, "x-pass-tracestate": "vendor=caller"}, + forward_headers=forward_headers, + ) + + assert upstream_headers["traceparent"] == caller_traceparent + assert upstream_headers["tracestate"] == "vendor=caller" + assert format(span.get_span_context().trace_id, "032x") not in upstream_headers["traceparent"] + + +@pytest.mark.asyncio +async def test_pass_through_request_without_caller_trace_headers_still_propagates_proxy_span(): + from opentelemetry.trace import get_current_span + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + + upstream_headers, span = await _relay_with_trace_headers({"x-pass-anthropic-beta": "beta-1"}, forward_headers=False) + + propagated = get_current_span(TraceContextTextMapPropagator().extract(upstream_headers)) + assert propagated.get_span_context().span_id == span.get_span_context().span_id + assert upstream_headers["anthropic-beta"] == "beta-1" + + @pytest.mark.asyncio async def test_pass_through_request_relays_non_json_body_without_buffering(): """ From d1563e0b552e79eea3b8c1a84bbd7808d467cae0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:09:13 -0700 Subject: [PATCH 08/33] feat: honor eager_input_streaming on Bedrock and Anthropic Claude tools --- litellm/llms/anthropic/chat/transformation.py | 19 +++- litellm/llms/anthropic/common_utils.py | 27 ++++- .../adapters/transformation.py | 11 +- .../bedrock/chat/converse_transformation.py | 12 +- .../anthropic_claude3_transformation.py | 12 +- litellm/llms/bedrock/common_utils.py | 14 ++- .../anthropic_claude3_transformation.py | 9 ++ litellm/proxy/_lazy_openapi_snapshot.json | 10 +- litellm/types/llms/anthropic.py | 3 + litellm/types/llms/openai.py | 2 + .../test_anthropic_chat_transformation.py | 71 ++++++++++++ ...al_pass_through_adapters_transformation.py | 42 +++++++ ...ations_anthropic_claude3_transformation.py | 53 +++++++++ .../chat/test_converse_transformation.py | 106 ++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 62 ++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 16 files changed, 444 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0f99441a115..1f90d375bc2 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -94,6 +94,7 @@ from litellm.utils import ( from ..common_utils import ( AnthropicError, AnthropicModelInfo, + eager_input_streaming_flag, process_anthropic_headers, strip_advisor_blocks_from_messages, ) @@ -732,10 +733,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): input_anthropic_schema: Final = sanitize_input_schema_for_anthropic(_input_schema) - _tool: Final = AnthropicMessagesTool( - name=tool["function"]["name"], - input_schema=input_anthropic_schema, - type="custom", + _eager_input_streaming: Final = eager_input_streaming_flag(tool) + _tool: Final = ( + AnthropicMessagesTool( + name=tool["function"]["name"], + input_schema=input_anthropic_schema, + type="custom", + ) + if _eager_input_streaming is None + else AnthropicMessagesTool( + name=tool["function"]["name"], + input_schema=input_anthropic_schema, + type="custom", + eager_input_streaming=_eager_input_streaming, + ) ) _description: Final = tool["function"].get("description") diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d35a9372058..2de9ab41d95 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -10,7 +10,7 @@ from types import MappingProxyType from typing import Any, Final, Literal import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError import litellm from litellm.constants import ( @@ -19,6 +19,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, is_encrypted_reasoning_block, @@ -231,6 +232,27 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup return headers, api_key +class _EagerInputStreamingFunction(BaseModel): + eager_input_streaming: StrictBool | None = None + + +class _EagerInputStreamingTool(BaseModel): + eager_input_streaming: StrictBool | None = None + function: _EagerInputStreamingFunction | None = None + + +def eager_input_streaming_flag(tool: object) -> bool | None: + try: + parsed: Final = _EagerInputStreamingTool.model_validate(tool) + except ValidationError as error: + if isinstance(tool, Mapping): + raise UnsupportedParamsError(message="eager_input_streaming must be a boolean") from error + return None + if parsed.eager_input_streaming is not None: + return parsed.eager_input_streaming + return parsed.function.eager_input_streaming if parsed.function is not None else None + + class AnthropicError(BaseLLMException): def __init__( self, @@ -373,6 +395,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): return False + def is_eager_input_streaming_used(self, tools: Sequence[object] | None) -> bool: + return any(eager_input_streaming_flag(tool) is True for tool in tools or ()) + @staticmethod def _supports_sampling_params(model: str) -> bool: """Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 7eb56ae55d3..d1047c8d86c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -111,6 +111,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_budget, ) from litellm.llms.anthropic.common_utils import ( + eager_input_streaming_flag, is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, strip_encrypted_reasoning_blocks_from_anthropic_messages, @@ -770,6 +771,7 @@ class LiteLLMAnthropicMessagesAdapter: "cache_control", "strict", "type", + "eager_input_streaming", ] for idx, tool in enumerate(tools): @@ -808,7 +810,14 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - tool_param = ChatCompletionToolParam(type="function", function=function_chunk) + eager_input_streaming = eager_input_streaming_flag(tool) + tool_param = ( + ChatCompletionToolParam(type="function", function=function_chunk) + if eager_input_streaming is None + else ChatCompletionToolParam( + type="function", function=function_chunk, eager_input_streaming=eager_input_streaming + ) + ) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1a4f27e2ddd..87e28054817 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -48,6 +48,7 @@ from litellm.llms.bedrock.request_metadata import ( merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) +from litellm.types.llms.anthropic import ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER from litellm.types.llms.bedrock import * from litellm.types.llms.openai import ( AllMessageValues, @@ -1518,10 +1519,7 @@ class AmazonConverseConfig(BaseConfig): bedrock_tools: list[ToolBlock] = [] # Collect anthropic_beta values from user headers - anthropic_beta_list: Final = [] - if headers: - user_betas: Final = get_anthropic_beta_from_headers(headers) - anthropic_beta_list.extend(user_betas) + anthropic_beta_list: Final = list(get_anthropic_beta_from_headers(headers or {})) # Separate pre-formatted Bedrock tools (e.g. systemTool from web_search_options) # from OpenAI-format tools that need transformation via _bedrock_tools_pt @@ -1634,6 +1632,12 @@ class AmazonConverseConfig(BaseConfig): if ANTHROPIC_EFFORT_BETA_HEADER not in anthropic_beta_list: anthropic_beta_list.append(ANTHROPIC_EFFORT_BETA_HEADER) + if ( + AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) + and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in anthropic_beta_list + ): + anthropic_beta_list.append(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + # Bedrock Converse: compact_20260112 edits only (+ beta header). AmazonConverseConfig._filter_context_management_for_bedrock_converse( additional_request_params, anthropic_beta_list diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 72bc43ba938..1326dc22ca0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -24,8 +24,12 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + tools_without_eager_input_streaming, +) +from litellm.types.llms.anthropic import ( + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, + ANTHROPIC_TOOL_SEARCH_BETA_HEADER, ) -from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -237,6 +241,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) normalize_custom_field_on_tools(anthropic_request) normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request) + outbound_tools: Final = tools_without_eager_input_streaming(anthropic_request) + if outbound_tools is not None: + anthropic_request["tools"] = outbound_tools return anthropic_request def _compute_bedrock_invoke_beta_headers( @@ -269,6 +276,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if bedrock_supports_tool_search(model): beta_set.add("tool-search-tool-2025-10-19") + if self.is_eager_input_streaming_used(tools): + beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + auto_beta_list: Final = filter_and_transform_beta_headers( beta_headers=list(beta_set - user_beta_set), provider="bedrock", diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 62b485588ac..7e24292a87e 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -9,7 +9,7 @@ import functools import json import os import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict if TYPE_CHECKING: @@ -18,6 +18,7 @@ if TYPE_CHECKING: from litellm.types.llms.bedrock import BedrockCreateBatchRequest import httpx +from pydantic import TypeAdapter, ValidationError import litellm from litellm import verbose_logger @@ -330,6 +331,17 @@ def normalize_custom_field_on_tools(request_body: dict) -> None: tool["defer_loading"] = deferred +_TOOL_DICTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def tools_without_eager_input_streaming(request_body: Mapping[str, object]) -> Sequence[object] | None: + try: + tools: Final = _TOOL_DICTS_ADAPTER.validate_python(request_body.get("tools")) + except ValidationError: + return None + return [{key: value for key, value in tool.items() if key != "eager_input_streaming"} for tool in tools] + + def normalize_json_schema_custom_types_to_object(schema: dict) -> None: """ In-place: replace JSON Schema ``type: \"custom\"`` with ``\"object\"`` (iterative walk). diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 4aa2afdbc78..d2be1ad9156 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -39,6 +39,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + tools_without_eager_input_streaming, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -46,6 +47,7 @@ from litellm.llms.bedrock.request_metadata import ( ) from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, ) from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest @@ -525,6 +527,9 @@ class AmazonAnthropicClaudeMessagesConfig( if injected_thinking_for_clear_thinking: beta_set.add("interleaved-thinking-2025-05-14") + if anthropic_model_info.is_eager_input_streaming_used(tools): + beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + self._filter_context_management_for_bedrock_invoke( anthropic_messages_request=anthropic_messages_request, beta_set=beta_set, @@ -719,6 +724,10 @@ class AmazonAnthropicClaudeMessagesConfig( if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + outbound_tools: Final = tools_without_eager_input_streaming(anthropic_messages_request) + if outbound_tools is not None: + anthropic_messages_request["tools"] = outbound_tools + remaining_output_config: Final = anthropic_messages_request.get("output_config") if ( litellm.drop_params is True diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5270df7b468..64daa695316 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19370,7 +19370,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { @@ -32899,6 +32899,10 @@ "cache_control": { "$ref": "#/components/schemas/ChatCompletionCachedContent" }, + "eager_input_streaming": { + "title": "Eager Input Streaming", + "type": "boolean" + }, "function": { "$ref": "#/components/schemas/ChatCompletionToolParamFunctionChunk" }, @@ -32928,6 +32932,10 @@ "title": "Description", "type": "string" }, + "eager_input_streaming": { + "title": "Eager Input Streaming", + "type": "boolean" + }, "name": { "title": "Name", "type": "string" diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index bcdee86360e..bcd24695f25 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -56,6 +56,7 @@ class AnthropicMessagesTool(TypedDict, total=False): defer_loading: bool allowed_callers: list[str] | None input_examples: list[dict[str, Any]] | None + eager_input_streaming: ReadOnly[bool] class AnthropicComputerTool(TypedDict, total=False): @@ -755,6 +756,8 @@ ANTHROPIC_TOOL_SEARCH_BETA_HEADER: Final = "advanced-tool-use-2025-11-20" # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" +ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER: Final = "fine-grained-tool-streaming-2025-05-14" + # OAuth constants ANTHROPIC_OAUTH_TOKEN_PREFIX: Final = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER: Final = "oauth-2025-04-20" diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..abb76527e82 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -992,6 +992,7 @@ class ChatCompletionToolParamFunctionChunk(TypedDict, total=False): description: str parameters: dict strict: bool + eager_input_streaming: ReadOnly[bool] class OpenAIChatCompletionToolParam(TypedDict): @@ -1002,6 +1003,7 @@ class OpenAIChatCompletionToolParam(TypedDict): class ChatCompletionToolParam(OpenAIChatCompletionToolParam, total=False): cache_control: ChatCompletionCachedContent allowed_callers: list[str] + eager_input_streaming: ReadOnly[bool] class Function(TypedDict, total=False): diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 269c351f866..1359bf40dcd 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6370,3 +6370,74 @@ def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(loc assert "tools" in result assert "tool_choice" not in result + + +def _eager_chat_tool(**extra): + return { + "type": "function", + "function": { + "name": "write_file", + "description": "Write a file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, + }, + **extra, + } + + +@pytest.mark.parametrize("flag", [True, False]) +def test_eager_input_streaming_passed_through_from_tool_top_level(flag): + mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming=flag)) + + assert mapped_tool == { + "name": "write_file", + "description": "Write a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, + "type": "custom", + "eager_input_streaming": flag, + } + + +def test_eager_input_streaming_passed_through_from_function(): + tool = _eager_chat_tool() + tool["function"]["eager_input_streaming"] = True + + mapped_tool, _ = AnthropicConfig()._map_tool_helper(tool) + + assert mapped_tool["eager_input_streaming"] is True + assert "eager_input_streaming" not in mapped_tool["input_schema"] + + +def test_eager_input_streaming_absent_stays_absent(): + mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool()) + + assert "eager_input_streaming" not in mapped_tool + + +def test_eager_input_streaming_rejects_non_boolean(): + with pytest.raises(litellm.BadRequestError, match="eager_input_streaming must be a boolean"): + AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming="true")) + + +def test_eager_input_streaming_not_set_on_computer_use_tool(): + computer_tool = { + "type": "computer_20250124", + "function": {"name": "computer", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + "eager_input_streaming": True, + } + + mapped_tool, _ = AnthropicConfig()._map_tool_helper(computer_tool) + + assert mapped_tool["type"] == "computer_20250124" + assert "eager_input_streaming" not in mapped_tool + + +def test_eager_input_streaming_reaches_anthropic_request_tools(): + result = AnthropicConfig().map_openai_params( + non_default_params={"tools": [_eager_chat_tool(eager_input_streaming=True)], "stream": True}, + optional_params={}, + model="claude-sonnet-5", + drop_params=False, + ) + + assert result["tools"][0]["eager_input_streaming"] is True + assert result["tools"][0]["name"] == "write_file" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e6782b70d3e..2318ceb937a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5017,3 +5017,45 @@ def test_redacted_thinking_blocks_never_carry_cache_control(): replayed: Final = outbound["messages"][1]["content"][0] assert replayed["type"] == "redacted_thinking" assert "cache_control" not in replayed + + +EAGER_INPUT_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +@pytest.mark.parametrize("flag", [True, False]) +def test_translate_anthropic_tools_to_openai_carries_eager_input_streaming_onto_tool(flag): + """The per-tool flag lands on the OpenAI tool object, never inside the JSON schema Bedrock sends as inputSchema.""" + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA, "eager_input_streaming": flag}] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert new_tools[0]["eager_input_streaming"] is flag + assert new_tools[0]["function"]["parameters"] == EAGER_INPUT_SCHEMA + assert "eager_input_streaming" not in new_tools[0]["function"] + + +def test_translate_anthropic_tools_to_openai_omits_unset_eager_input_streaming(): + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA}] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert "eager_input_streaming" not in new_tools[0] + assert "eager_input_streaming" not in new_tools[0]["function"]["parameters"] + + +def test_eager_input_streaming_tool_reaches_bedrock_converse_as_beta(): + """An Anthropic Messages request routed to bedrock/converse/ turns the flag into the fine-grained streaming beta.""" + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA, "eager_input_streaming": True}] + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + data: Final = AmazonConverseConfig()._transform_request_helper( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + system_content_blocks=[], + optional_params={"tools": new_tools}, + messages=[{"role": "user", "content": "write a big file"}], + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == ["fine-grained-tool-streaming-2025-05-14"] + assert data["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"] == EAGER_INPUT_SCHEMA diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index ec0bf6b842a..1395489334e 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -859,3 +859,56 @@ def test_bedrock_chat_invoke_tool_search_beta_follows_model_map( ) assert result.get("anthropic_beta") == expected_betas + + +FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +def _chat_invoke_request_with_tools(tools, headers=None): + config = AmazonAnthropicClaudeConfig() + model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params = config.map_openai_params( + non_default_params={"max_tokens": 64, "stream": True, "tools": tools}, + optional_params={}, + model=model, + drop_params=False, + ) + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "write a big file"}], + optional_params=optional_params, + litellm_params={}, + headers=headers or {}, + ) + + +def _eager_openai_tool(name, **extra): + return {"type": "function", "function": {"name": name, "parameters": EAGER_TOOL_SCHEMA}, **extra} + + +def test_bedrock_chat_invoke_eager_input_streaming_tool_adds_beta_and_strips_key(): + result = _chat_invoke_request_with_tools( + [_eager_openai_tool("write_file", eager_input_streaming=True), _eager_openai_tool("read_file")] + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + assert [tool["name"] for tool in result["tools"]] == ["write_file", "read_file"] + assert all("eager_input_streaming" not in tool for tool in result["tools"]) + assert result["tools"][0]["input_schema"] == EAGER_TOOL_SCHEMA + + +def test_bedrock_chat_invoke_eager_input_streaming_false_strips_key_without_beta(): + result = _chat_invoke_request_with_tools([_eager_openai_tool("write_file", eager_input_streaming=False)]) + + assert "anthropic_beta" not in result + assert "eager_input_streaming" not in result["tools"][0] + + +def test_bedrock_chat_invoke_eager_input_streaming_beta_not_duplicated_with_client_header(): + result = _chat_invoke_request_with_tools( + [_eager_openai_tool("write_file", eager_input_streaming=True)], + headers={"anthropic-beta": FINE_GRAINED_TOOL_STREAMING_BETA}, + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 96c78c1cf75..650427f2ae3 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -7400,3 +7400,109 @@ def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it( ) assert result.choices[0].message.tool_calls is None assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} + + +FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +def _eager_openai_tool(**extra): + return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA}, **extra} + + +def _eager_openai_function_tool(**extra): + return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA, **extra}} + + +def _eager_anthropic_tool(**extra): + return {"name": "write_file", "input_schema": EAGER_TOOL_SCHEMA, **extra} + + +def _converse_request(model, tools, headers=None): + return AmazonConverseConfig()._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={"tools": tools}, + messages=[{"role": "user", "content": "write a big file"}], + headers=headers, + ) + + +@pytest.mark.parametrize( + "tool", + [ + _eager_openai_tool(eager_input_streaming=True), + _eager_openai_function_tool(eager_input_streaming=True), + _eager_anthropic_tool(eager_input_streaming=True), + ], + ids=["openai_top_level", "openai_under_function", "anthropic_shape"], +) +def test_eager_input_streaming_tool_adds_fine_grained_tool_streaming_beta(tool): + data = _converse_request("us.anthropic.claude-sonnet-4-5-20250929-v1:0", [tool]) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + tool_spec = data["toolConfig"]["tools"][0]["toolSpec"] + assert tool_spec["name"] == "write_file" + assert "eager_input_streaming" not in tool_spec + assert "eager_input_streaming" not in tool_spec["inputSchema"]["json"] + + +@pytest.mark.parametrize( + "tool", + [ + _eager_openai_tool(eager_input_streaming=False), + _eager_openai_function_tool(eager_input_streaming=False), + _eager_anthropic_tool(eager_input_streaming=False), + _eager_openai_tool(), + ], + ids=["openai_false", "function_false", "anthropic_false", "absent"], +) +def test_eager_input_streaming_false_or_absent_adds_no_beta(tool): + data = _converse_request("us.anthropic.claude-sonnet-4-5-20250929-v1:0", [tool]) + + assert "anthropic_beta" not in data.get("additionalModelRequestFields", {}) + assert "eager_input_streaming" not in data["toolConfig"]["tools"][0]["toolSpec"] + + +def test_eager_input_streaming_beta_only_on_anthropic_models(): + data = _converse_request("amazon.nova-pro-v1:0", [_eager_openai_tool(eager_input_streaming=True)]) + + assert "anthropic_beta" not in data.get("additionalModelRequestFields", {}) + assert data["toolConfig"]["tools"][0]["toolSpec"]["name"] == "write_file" + + +def test_eager_input_streaming_beta_not_duplicated_with_client_header(): + data = _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming=True)], + headers={"anthropic-beta": f"{FINE_GRAINED_TOOL_STREAMING_BETA},interleaved-thinking-2025-05-14"}, + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [ + FINE_GRAINED_TOOL_STREAMING_BETA, + "interleaved-thinking-2025-05-14", + ] + + +def test_eager_input_streaming_beta_never_written_back_into_client_header_list(): + headers = {"anthropic-beta": ["interleaved-thinking-2025-05-14"]} + + data = _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming=True)], + headers=headers, + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [ + "interleaved-thinking-2025-05-14", + FINE_GRAINED_TOOL_STREAMING_BETA, + ] + assert headers == {"anthropic-beta": ["interleaved-thinking-2025-05-14"]} + + +def test_eager_input_streaming_non_boolean_is_a_bad_request(): + with pytest.raises(litellm.BadRequestError, match="eager_input_streaming must be a boolean"): + _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming="true")], + ) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index e43accdb835..b71e8e521f2 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3244,3 +3244,65 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo ) assert result.get("output_config") == {"format": schema_format} + + +FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" + + +def _invoke_request_with_tools(tools, headers=None): + from litellm.types.router import GenericLiteLLMParams + + return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "write a big file"}], + anthropic_messages_optional_request_params={"max_tokens": 4096, "tools": copy.deepcopy(tools), "stream": True}, + litellm_params=GenericLiteLLMParams(), + headers=headers or {}, + ) + + +def _eager_invoke_tool(name, eager_input_streaming): + return { + "name": name, + "description": f"{name} tool", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + "eager_input_streaming": eager_input_streaming, + } + + +def test_bedrock_invoke_eager_input_streaming_tool_adds_beta_and_strips_key(): + result = _invoke_request_with_tools( + [ + _eager_invoke_tool("write_file", True), + _eager_invoke_tool("read_file", False), + {"name": "list_files", "input_schema": {"type": "object", "properties": {}}}, + ] + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + assert [tool["name"] for tool in result["tools"]] == ["write_file", "read_file", "list_files"] + assert all("eager_input_streaming" not in tool for tool in result["tools"]) + assert result["tools"][0]["description"] == "write_file tool" + assert result["tools"][0]["input_schema"] == {"type": "object", "properties": {"path": {"type": "string"}}} + + +def test_bedrock_invoke_eager_input_streaming_false_strips_key_without_beta(): + result = _invoke_request_with_tools([_eager_invoke_tool("write_file", False)]) + + assert "anthropic_beta" not in result + assert result["tools"] == [ + { + "name": "write_file", + "description": "write_file tool", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ] + + +def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_header(): + result = _invoke_request_with_tools( + [_eager_invoke_tool("write_file", True)], + headers={"anthropic-beta": FINE_GRAINED_TOOL_STREAMING_BETA}, + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1b0081cb9c..cc5169dd84c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25970,6 +25970,8 @@ export interface components { /** Allowed Callers */ allowed_callers?: string[]; cache_control?: components["schemas"]["ChatCompletionCachedContent"]; + /** Eager Input Streaming */ + eager_input_streaming?: boolean; function: components["schemas"]["ChatCompletionToolParamFunctionChunk"]; /** Type */ type: "function" | string; @@ -25978,6 +25980,8 @@ export interface components { ChatCompletionToolParamFunctionChunk: { /** Description */ description?: string; + /** Eager Input Streaming */ + eager_input_streaming?: boolean; /** Name */ name: string; /** Parameters */ From a37d1e8b9b752f2c0672bc394e14095ef7c69b31 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:30:08 -0700 Subject: [PATCH 09/33] test: type the eager_input_streaming test helpers and mark constants Final --- .../test_anthropic_chat_transformation.py | 22 +++++++++---------- ...ations_anthropic_claude3_transformation.py | 17 ++++++++------ .../chat/test_converse_transformation.py | 15 ++++++++----- .../test_anthropic_claude3_transformation.py | 9 +++++--- 4 files changed, 36 insertions(+), 27 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 1359bf40dcd..5e179f950a0 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6372,18 +6372,19 @@ def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(loc assert "tool_choice" not in result -def _eager_chat_tool(**extra): +def _eager_chat_function(**extra: object) -> dict[str, object]: return { - "type": "function", - "function": { - "name": "write_file", - "description": "Write a file", - "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, - }, + "name": "write_file", + "description": "Write a file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, **extra, } +def _eager_chat_tool(**extra: object) -> dict[str, object]: + return {"type": "function", "function": _eager_chat_function(), **extra} + + @pytest.mark.parametrize("flag", [True, False]) def test_eager_input_streaming_passed_through_from_tool_top_level(flag): mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming=flag)) @@ -6398,10 +6399,9 @@ def test_eager_input_streaming_passed_through_from_tool_top_level(flag): def test_eager_input_streaming_passed_through_from_function(): - tool = _eager_chat_tool() - tool["function"]["eager_input_streaming"] = True - - mapped_tool, _ = AnthropicConfig()._map_tool_helper(tool) + mapped_tool, _ = AnthropicConfig()._map_tool_helper( + {"type": "function", "function": _eager_chat_function(eager_input_streaming=True)} + ) assert mapped_tool["eager_input_streaming"] is True assert "eager_input_streaming" not in mapped_tool["input_schema"] diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 1395489334e..84db0733227 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,6 +1,7 @@ import asyncio import json import uuid +from typing import Final from unittest.mock import patch import httpx @@ -861,14 +862,16 @@ def test_bedrock_chat_invoke_tool_search_beta_follows_model_map( assert result.get("anthropic_beta") == expected_betas -FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" -EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} -def _chat_invoke_request_with_tools(tools, headers=None): - config = AmazonAnthropicClaudeConfig() - model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" - optional_params = config.map_openai_params( +def _chat_invoke_request_with_tools( + tools: list[dict[str, object]], headers: dict[str, str] | None = None +) -> dict[str, object]: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params: Final = config.map_openai_params( non_default_params={"max_tokens": 64, "stream": True, "tools": tools}, optional_params={}, model=model, @@ -883,7 +886,7 @@ def _chat_invoke_request_with_tools(tools, headers=None): ) -def _eager_openai_tool(name, **extra): +def _eager_openai_tool(name: str, **extra: object) -> dict[str, object]: return {"type": "function", "function": {"name": name, "parameters": EAGER_TOOL_SCHEMA}, **extra} diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 650427f2ae3..cd33e8d34d8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -4,6 +4,7 @@ import os import httpx import pytest +from typing import Final from unittest.mock import MagicMock, patch import litellm @@ -7402,23 +7403,25 @@ def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it( assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} -FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" -EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} -def _eager_openai_tool(**extra): +def _eager_openai_tool(**extra: object) -> dict[str, object]: return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA}, **extra} -def _eager_openai_function_tool(**extra): +def _eager_openai_function_tool(**extra: object) -> dict[str, object]: return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA, **extra}} -def _eager_anthropic_tool(**extra): +def _eager_anthropic_tool(**extra: object) -> dict[str, object]: return {"name": "write_file", "input_schema": EAGER_TOOL_SCHEMA, **extra} -def _converse_request(model, tools, headers=None): +def _converse_request( + model: str, tools: list[dict[str, object]], headers: dict[str, object] | None = None +) -> dict[str, object]: return AmazonConverseConfig()._transform_request_helper( model=model, system_content_blocks=[], diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index b71e8e521f2..7be005c0efe 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -4,6 +4,7 @@ import json import os from datetime import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import Mock import pytest @@ -3246,10 +3247,12 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo assert result.get("output_config") == {"format": schema_format} -FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" -def _invoke_request_with_tools(tools, headers=None): +def _invoke_request_with_tools( + tools: list[dict[str, object]], headers: dict[str, str] | None = None +) -> dict[str, object]: from litellm.types.router import GenericLiteLLMParams return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( @@ -3261,7 +3264,7 @@ def _invoke_request_with_tools(tools, headers=None): ) -def _eager_invoke_tool(name, eager_input_streaming): +def _eager_invoke_tool(name: str, eager_input_streaming: bool) -> dict[str, object]: return { "name": name, "description": f"{name} tool", From e5fd2bec832be2ddf76851f73859d8bc2fb2b957 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:34:29 -0700 Subject: [PATCH 10/33] fix: forward eager_input_streaming through the Responses API tool bridge --- .../transformation.py | 4 +++ .../test_litellm_completion_responses.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 01fb6cb483d..0d1297f63e2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2020,6 +2020,8 @@ class LiteLLMCompletionResponsesConfig: chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") if tool.get("input_examples"): chat_completion_tool["input_examples"] = tool.get("input_examples") + if tool.get("eager_input_streaming") is not None: + chat_completion_tool["eager_input_streaming"] = tool.get("eager_input_streaming") return ResponsesToolChatForm( chat_tools=(cast(ChatCompletionToolParam, chat_completion_tool),), web_search_options=None ) @@ -2096,6 +2098,8 @@ class LiteLLMCompletionResponsesConfig: responses_tool["allowed_callers"] = tool.get("allowed_callers") if tool.get("input_examples") is not None: responses_tool["input_examples"] = tool.get("input_examples") + if tool.get("eager_input_streaming") is not None: + responses_tool["eager_input_streaming"] = tool.get("eager_input_streaming") result.append(responses_tool) else: # mcp or other: pass through unchanged diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..f1924b87a20 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1900,6 +1900,38 @@ class TestToolTransformation: assert "defer_loading" not in result_tool assert "allowed_callers" not in result_tool assert "input_examples" not in result_tool + assert "eager_input_streaming" not in result_tool + + @pytest.mark.parametrize("eager_input_streaming", [True, False]) + def test_transform_function_tools_forwards_eager_input_streaming(self, eager_input_streaming: bool) -> None: + function_tool: Final = { + "type": "function", + "name": "write_file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}, + "eager_input_streaming": eager_input_streaming, + } + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[function_tool] + ) + + assert result_tools[0]["eager_input_streaming"] is eager_input_streaming + + @pytest.mark.parametrize("eager_input_streaming", [True, False]) + def test_chat_completion_tools_to_responses_tools_keeps_eager_input_streaming( + self, eager_input_streaming: bool + ) -> None: + chat_tool: Final = { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object"}}, + "eager_input_streaming": eager_input_streaming, + } + + result_tools: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( + [chat_tool] + ) + + assert result_tools[0]["eager_input_streaming"] is eager_input_streaming def test_transform_code_execution_tools(self): """Test that code_execution tools are passed through as-is""" From 3d3a46fc1fcd66549958807360ad4b182117a673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:08:11 -0700 Subject: [PATCH 11/33] refactor: build the Converse beta list once and keep the OpenAPI snapshot as CI generates it --- .../adapters/transformation.py | 18 +++++++++------- .../bedrock/chat/converse_transformation.py | 21 ++++++++++--------- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index d1047c8d86c..b213b5e30ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -198,6 +198,15 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider) +def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam: + eager_input_streaming: Final = eager_input_streaming_flag(tool) + if eager_input_streaming is None: + return ChatCompletionToolParam(type="function", function=function_chunk) + return ChatCompletionToolParam( + type="function", function=function_chunk, eager_input_streaming=eager_input_streaming + ) + + class AnthropicAdapter: def __init__(self) -> None: pass @@ -810,14 +819,7 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - eager_input_streaming = eager_input_streaming_flag(tool) - tool_param = ( - ChatCompletionToolParam(type="function", function=function_chunk) - if eager_input_streaming is None - else ChatCompletionToolParam( - type="function", function=function_chunk, eager_input_streaming=eager_input_streaming - ) - ) + tool_param = _chat_tool_param(function_chunk, tool) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 87e28054817..1a32fec45e3 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1518,9 +1518,6 @@ class AmazonConverseConfig(BaseConfig): """Process tools and collect anthropic_beta values.""" bedrock_tools: list[ToolBlock] = [] - # Collect anthropic_beta values from user headers - anthropic_beta_list: Final = list(get_anthropic_beta_from_headers(headers or {})) - # Separate pre-formatted Bedrock tools (e.g. systemTool from web_search_options) # from OpenAI-format tools that need transformation via _bedrock_tools_pt filtered_tools: Final = [] @@ -1540,6 +1537,17 @@ class AmazonConverseConfig(BaseConfig): continue filtered_tools.append(tool) + base_model: Final = BedrockModelInfo.get_base_model(model) + client_beta_list: Final = get_anthropic_beta_from_headers(headers or {}) + eager_beta: Final = ( + (ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER,) + if base_model.startswith("anthropic") + and AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) + and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in client_beta_list + else () + ) + anthropic_beta_list: Final = [*client_beta_list, *eager_beta] + # Only separate tools if computer use tools are actually present if filtered_tools and self.is_computer_use_tool_used(filtered_tools, model): # Separate computer use tools from regular function tools @@ -1617,7 +1625,6 @@ class AmazonConverseConfig(BaseConfig): # Opus 4.5 gates ``output_config.effort`` behind a beta header; # Claude 4.6/4.7 accept it without one. - base_model: Final = BedrockModelInfo.get_base_model(model) if base_model.startswith("anthropic"): output_config: Final = additional_request_params.get("output_config") if ( @@ -1632,12 +1639,6 @@ class AmazonConverseConfig(BaseConfig): if ANTHROPIC_EFFORT_BETA_HEADER not in anthropic_beta_list: anthropic_beta_list.append(ANTHROPIC_EFFORT_BETA_HEADER) - if ( - AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) - and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in anthropic_beta_list - ): - anthropic_beta_list.append(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) - # Bedrock Converse: compact_20260112 edits only (+ beta header). AmazonConverseConfig._filter_context_management_for_bedrock_converse( additional_request_params, anthropic_beta_list diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 64daa695316..7a4c966a628 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19370,7 +19370,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 4a951847bbcb25a8c42eb526b0526604d1b3bb1a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:15:39 -0700 Subject: [PATCH 12/33] fix(responses): merge deployment litellm_params into native websocket response.create frames --- litellm/llms/custom_httpx/llm_http_handler.py | 3 + litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/responses/main.py | 25 ++- litellm/responses/streaming_iterator.py | 31 ++- .../types/responses/streaming_websocket.py | 13 ++ .../test_responses_websocket_all_providers.py | 210 ++++++++++++++++++ 6 files changed, 275 insertions(+), 9 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 98fe0014386..477d10a3cbd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -144,6 +144,7 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult +from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CallTypes, @@ -6588,6 +6589,7 @@ class BaseLLMHTTPHandler: litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, + request_defaults: ResponsesWebSocketRequestDefaults | None = None, **kwargs: Any, ): """ @@ -6742,6 +6744,7 @@ class BaseLLMHTTPHandler: output_guardrail_callbacks=_ws_output_guardrail_callbacks, quota_callbacks=_ws_quota_callbacks, authorized_model=model, + request_defaults=request_defaults, ) await streaming.bidirectional_forward() diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8a8d08c6887..2aa2cf15ac1 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19394,7 +19394,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 6dc34bb93ef..6064fa66d91 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import assert_never import litellm @@ -53,6 +53,7 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * +from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import all_litellm_params from litellm.utils import ( @@ -2261,6 +2262,27 @@ def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: return metadata +_EXTRA_BODY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults: + reasoning_effort: Final = kwargs.get("reasoning_effort") + mapped_reasoning: Final = ( + LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) + if kwargs.get("reasoning") is None and isinstance(reasoning_effort, str) + else None + ) + candidate_params: Final[dict[str, object]] = { + **kwargs, + **({"reasoning": mapped_reasoning} if mapped_reasoning is not None else {}), + } + fill_missing: Final = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(candidate_params) + return ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType(dict(fill_missing)), + overrides=MappingProxyType(_EXTRA_BODY_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), + ) + + @client async def _aresponses_websocket( model: str, @@ -2352,5 +2374,6 @@ async def _aresponses_websocket( user_api_key_dict=kwargs.get("user_api_key_dict"), litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, + request_defaults=_build_responses_websocket_request_defaults(kwargs), **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8d766cf1cd0..1c82df8c664 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -54,6 +54,7 @@ if TYPE_CHECKING: PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, + ResponsesWebSocketRequestDefaults, ) from litellm.types.router import LiteLLM_Params @@ -1717,6 +1718,7 @@ class ResponsesWebSocketStreaming: output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, + request_defaults: ResponsesWebSocketRequestDefaults | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1732,6 +1734,7 @@ class ResponsesWebSocketStreaming: # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model + self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1874,12 +1877,23 @@ class ResponsesWebSocketStreaming: modified = True return modified + def _with_request_defaults(self, msg_obj: dict[str, object]) -> dict[str, object]: + if self.request_defaults is None: + return msg_obj + nested: Final = msg_obj.get("response") + if _is_json_object(nested): + return {**msg_obj, "response": self.request_defaults.merged_into(nested)} + return self.request_defaults.merged_into(msg_obj) + async def _mask_response_create(self, message: str) -> str: """ - Enforce the authorized model and apply Presidio PII masking to a - ``response.create`` message before it is forwarded to the upstream - provider. + Merge deployment defaults, enforce the authorized model, and apply + Presidio PII masking to a ``response.create`` message before it is + forwarded to the upstream provider. + - Fills the deployment's ``litellm_params`` request defaults into the + frame the way the HTTP ``/v1/responses`` path does: client-set keys + win, ``extra_body`` entries override. - Overwrites any ``model`` field with the connection-authorized model to prevent deployment-substitution attacks (always applied). - Walks the ``input`` and ``instructions`` fields, calls ``check_pii`` @@ -1889,23 +1903,26 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = _load_json_object(message) + parsed: Final = _load_json_object(message) except (json.JSONDecodeError, TypeError): return message - if msg_obj.get("type") != "response.create": + if parsed.get("type") != "response.create": return message + msg_obj: Final = self._with_request_defaults(parsed) + defaults_applied: Final = msg_obj != parsed + # Always enforce the authorized model, even when PII masking is off. model_modified: Final = self._enforce_authorized_model(msg_obj) if not self.guardrail_callbacks: - return json.dumps(msg_obj) if model_modified else message + return json.dumps(msg_obj) if model_modified or defaults_applied else message if "metadata" not in self.request_data: self.request_data["metadata"] = {} - modified = model_modified + modified = model_modified or defaults_applied guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks) for cb in guardrail_cbs: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) diff --git a/litellm/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index 2aa71647955..f369cbcebf8 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -1,5 +1,7 @@ from __future__ import annotations +from collections.abc import Mapping +from dataclasses import dataclass from typing import Protocol from litellm.types.guardrails import PresidioPerRequestConfig @@ -39,3 +41,14 @@ class PresidioGuardrailCallback(Protocol): presidio_config: PresidioPerRequestConfig | None, request_data: dict[str, object], ) -> str: ... + + +@dataclass(frozen=True, slots=True) +class ResponsesWebSocketRequestDefaults: + """Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket.""" + + fill_missing: Mapping[str, object] + overrides: Mapping[str, object] + + def merged_into(self, request: Mapping[str, object]) -> dict[str, object]: + return {**self.fill_missing, **request, **self.overrides} diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index fe3c4a0640d..918c4a33d40 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1204,6 +1204,216 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +def _deployment_defaults(): + from types import MappingProxyType + + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + return ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType({"reasoning": {"effort": "high"}, "service_tier": "priority"}), + overrides=MappingProxyType({"provider_default": "configured"}), + ) + + +class TestNativeWebSocketDeploymentDefaults: + """The native relay merges deployment litellm_params into every response.create like HTTP does.""" + + def test_builder_maps_router_kwargs_like_the_http_path(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + { + "model": "gpt-5-pro", + "reasoning_effort": "high", + "service_tier": "priority", + "extra_body": {"provider_default": "configured"}, + "temperature": None, + "timeout": 600, + "max_retries": 2, + "caching": False, + "custom_llm_provider": "openai", + "litellm_metadata": {"user_api_key": "hashed"}, + "user_api_key_dict": MagicMock(), + "litellm_logging_obj": MagicMock(), + "websocket": MagicMock(), + } + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"} + assert dict(defaults.overrides) == {"provider_default": "configured"} + + def test_builder_keeps_explicit_reasoning_over_reasoning_effort(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + {"model": "gpt-5-pro", "reasoning": {"effort": "low"}, "reasoning_effort": "high"} + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "low"}} + assert dict(defaults.overrides) == {} + + @pytest.mark.asyncio + async def test_flat_frame_gets_defaults_client_keys_win_extra_body_overrides(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "service_tier": "default", + "provider_default": "client", + } + ) + ) + ) + + assert forwarded == { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "service_tier": "default", + "provider_default": "configured", + "reasoning": {"effort": "high"}, + } + + @pytest.mark.asyncio + async def test_nested_response_frame_gets_defaults_inside_response(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps({"type": "response.create", "response": {"model": "gpt-5-pro", "input": "hi"}}) + ) + ) + + assert forwarded == { + "type": "response.create", + "response": { + "model": "gpt-5-pro", + "input": "hi", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + }, + } + + @pytest.mark.asyncio + async def test_frames_that_need_nothing_pass_through_untouched(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + cancel_frame = json.dumps({"type": "response.cancel"}) + complete_frame = json.dumps( + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "hi", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + ) + + assert await handler._mask_response_create(cancel_frame) is cancel_frame + assert await handler._mask_response_create(complete_frame) is complete_frame + + @pytest.mark.asyncio + async def test_handler_applies_defaults_to_the_first_frame_sent_upstream(self): + import asyncio + from unittest.mock import AsyncMock, patch + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + class FakeBackend: + def __init__(self): + self.sent = [] + + async def send(self, message): + self.sent.append(message) + + async def recv(self, decode=False): + raise RuntimeError("backend closed") + + async def close(self): + pass + + backend = FakeBackend() + + class FakeConnect: + def __init__(self, url, **kwargs): + pass + + async def __aenter__(self): + return backend + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.model_in_websocket_url.return_value = True + mock_config.get_websocket_url.return_value = "wss://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + mock_logging.dispatch_success_handlers = AsyncMock() + + client_ws = MagicMock() + client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await BaseLLMHTTPHandler().async_responses_websocket( + model="gpt-5-pro", + websocket=client_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + first_message=json.dumps({"type": "response.create", "model": "gpt-5-pro", "input": "Say hello"}), + request_defaults=_deployment_defaults(), + ) + await asyncio.sleep(0) + + assert [json.loads(frame) for frame in backend.sent] == [ + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + ] + + @pytest.mark.asyncio + async def test_aresponses_websocket_builds_defaults_from_deployment_kwargs(self, monkeypatch): + import importlib + from unittest.mock import AsyncMock + + responses_main = importlib.import_module("litellm.responses.main") + + stub = MagicMock() + stub.async_responses_websocket = AsyncMock() + monkeypatch.setattr(responses_main, "base_llm_http_handler", stub) + + await responses_main._aresponses_websocket.__wrapped__( + model="openai/gpt-5-pro", + websocket=MagicMock(), + api_key="sk-test", + litellm_logging_obj=MagicMock(), + reasoning_effort="high", + service_tier="priority", + extra_body={"provider_default": "configured"}, + ) + + request_defaults = stub.async_responses_websocket.call_args.kwargs["request_defaults"] + assert dict(request_defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"} + assert dict(request_defaults.overrides) == {"provider_default": "configured"} + + class TestNativeWebSocketGuardrails: @pytest.mark.asyncio async def test_response_create_injects_authorized_model(self): From cf05466a27ad9409d1edcd9697ddb826ca8f93da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:54:18 -0700 Subject: [PATCH 13/33] fix(gemini): map every documented finishReason and reset per-candidate state A content-less candidate is now kept as a choice whenever it carries a finishReason, with the raw value on the choice's provider_specific_fields. NO_IMAGE, IMAGE_RECITATION, IMAGE_OTHER and ESCALATION map to content_filter; UNEXPECTED_TOOL_CALL and MISSING_THOUGHT_SIGNATURE map to stop. The /v1/responses bridge reports content_filter and refusal as incomplete with incomplete_details, and tool calls and reasoning no longer leak from one candidate into the next. --- litellm/litellm_core_utils/core_helpers.py | 5 ++ .../vertex_and_google_ai_studio_gemini.py | 29 ++++---- .../transformation.py | 33 ++++++--- .../litellm_core_utils/test_core_helpers.py | 6 ++ ...test_vertex_and_google_ai_studio_gemini.py | 68 +++++++++++++++++++ .../test_litellm_completion_responses.py | 2 +- 6 files changed, 117 insertions(+), 26 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 2615c58de2c..d29b1fc74ef 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -225,6 +225,11 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "TOO_MANY_TOOL_CALLS": "stop", "MALFORMED_RESPONSE": "stop", "NO_IMAGE": "content_filter", + "IMAGE_RECITATION": "content_filter", + "IMAGE_OTHER": "content_filter", + "ESCALATION": "content_filter", + "UNEXPECTED_TOOL_CALL": "stop", + "MISSING_THOUGHT_SIGNATURE": "stop", # Zhipu GLM "network_error": "stop", "sensitive": "content_filter", diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 61bac4793a8..a95b845718a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1348,6 +1348,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + "ESCALATION", + "UNEXPECTED_TOOL_CALL", + "MISSING_THOUGHT_SIGNATURE", } ) @@ -2243,7 +2248,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): image_response: list[ImageURLListItem] | None = None chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} chat_completion_logprobs: ChoiceLogprobs | None = None - tools: list[ChatCompletionToolCallChunk] | None = [] + tools: list[ChatCompletionToolCallChunk] | None = None functions: ChatCompletionToolCallFunctionChunk | None = None thinking_blocks: list[ChatCompletionThinkingBlock] | None = None reasoning_content: str | None = None @@ -2358,11 +2363,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_invocation_fields["server_side_tool_invocations"] = server_side_tool_invocations chat_completion_message["provider_specific_fields"] = tool_invocation_fields - if candidate.get("finishReason"): - finish_reason_fields = chat_completion_message.get("provider_specific_fields") or {} - finish_reason_fields["native_finish_reason"] = candidate.get("finishReason") - chat_completion_message["provider_specific_fields"] = finish_reason_fields - if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, @@ -2375,15 +2375,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) model_response.choices.append(choice) elif isinstance(model_response, ModelResponse): + native_finish_reason = candidate.get("finishReason") choice = litellm.Choices( finish_reason=VertexGeminiConfig._check_finish_reason( - chat_completion_message, candidate.get("finishReason") + chat_completion_message, native_finish_reason ), index=candidate.get("index", idx), message=chat_completion_message, logprobs=chat_completion_logprobs, enhancements=None, - provider_specific_fields=chat_completion_message.get("provider_specific_fields"), + provider_specific_fields=( + {"native_finish_reason": native_finish_reason} if native_finish_reason is not None else None + ), ) model_response.choices.append(choice) @@ -3181,12 +3184,10 @@ class ModelResponseIterator: self.has_seen_tool_calls = True break - # _process_candidates skips candidates without a "content" part, so a - # content-less chunk leaves choices empty and the downstream streaming - # handler hits IndexError on choices[0]. This covers the final chunk - # (finishReason, no content) and mid-stream metadata-only chunks - # (grounding/web-search/thought, no content and no finishReason — seen - # with web_search + reasoning) by emitting an empty-delta choice. + # _process_candidates skips candidates with neither "content" nor + # "finishReason", so a metadata-only chunk (grounding/web-search/thought, + # seen with web_search + reasoning) leaves choices empty and the downstream + # streaming handler hits IndexError on choices[0]. Emit an empty-delta choice. if not model_response.choices and _candidates: from litellm.types.utils import Delta, StreamingChoices diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7792bc11405..961fffd3d42 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -61,6 +61,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, ChatCompletionUserMessage, GenericChatCompletionMessage, + IncompleteDetails, InputTokensDetails, OpenAIChatCompletionTextObject, OpenAIMcpServerTool, @@ -2295,6 +2296,21 @@ class LiteLLMCompletionResponsesConfig: # Default to completed for unknown finish reasons return "completed" + @staticmethod + def _incomplete_details_for_finish_reason( + finish_reason: str | None, + existing: IncompleteDetails | None, + ) -> IncompleteDetails | None: + if existing is not None: + return existing + match finish_reason: + case "length": + return IncompleteDetails(reason="max_output_tokens") + case "content_filter" | "refusal": + return IncompleteDetails(reason="content_filter") + case _: + return None + @staticmethod def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) -> str: """Bedrock Mantle returns a non-unique, index-based ``call_id`` (``call_0``, @@ -2411,17 +2427,10 @@ class LiteLLMCompletionResponsesConfig: if choices and len(choices) > 0: finish_reason = choices[0].finish_reason - status: Final[ResponsesAPIStatus] = ( - LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(finish_reason) + incomplete_details: Final = LiteLLMCompletionResponsesConfig._incomplete_details_for_finish_reason( + finish_reason=finish_reason, + existing=getattr(chat_completion_response, "incomplete_details", None), ) - incomplete_details = getattr(chat_completion_response, "incomplete_details", None) - if incomplete_details is None and status == "incomplete": - from openai.types.responses.response import IncompleteDetails - - if finish_reason == "length": - incomplete_details = IncompleteDetails(reason="max_output_tokens") - elif finish_reason in ["content_filter", "refusal"]: - incomplete_details = IncompleteDetails(reason="content_filter") responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse( id=chat_completion_response.id, @@ -2447,7 +2456,9 @@ class LiteLLMCompletionResponsesConfig: max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), previous_response_id=getattr(chat_completion_response, "previous_response_id", None), reasoning=None, - status=status, + status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + finish_reason + ), text={}, truncation=getattr(chat_completion_response, "truncation", None), usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index e937be47441..b2ad13c205e 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -151,6 +151,12 @@ class TestMapFinishReasonGemini: ("IMAGE_PROHIBITED_CONTENT", "content_filter"), ("TOO_MANY_TOOL_CALLS", "stop"), ("MALFORMED_RESPONSE", "stop"), + ("NO_IMAGE", "content_filter"), + ("IMAGE_RECITATION", "content_filter"), + ("IMAGE_OTHER", "content_filter"), + ("ESCALATION", "content_filter"), + ("UNEXPECTED_TOOL_CALL", "stop"), + ("MISSING_THOUGHT_SIGNATURE", "stop"), ], ) def test_gemini_finish_reasons(self, gemini_reason, expected): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 216d6f0db7e..4ee199bd9af 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -968,6 +968,12 @@ def test_finish_reason_unspecified_and_malformed_function_call(): # Test new Gemini finish reasons assert finish_reason_mappings["TOO_MANY_TOOL_CALLS"] == "stop" assert finish_reason_mappings["MALFORMED_RESPONSE"] == "stop" + assert finish_reason_mappings["NO_IMAGE"] == "content_filter" + assert finish_reason_mappings["IMAGE_RECITATION"] == "content_filter" + assert finish_reason_mappings["IMAGE_OTHER"] == "content_filter" + assert finish_reason_mappings["ESCALATION"] == "content_filter" + assert finish_reason_mappings["UNEXPECTED_TOOL_CALL"] == "stop" + assert finish_reason_mappings["MISSING_THOUGHT_SIGNATURE"] == "stop" def test_vertex_ai_usage_metadata_response_token_count(): @@ -6219,3 +6225,65 @@ def test_gemini_candidate_other_finish_reasons_no_content(): ) assert responses_length.status == "incomplete" assert responses_length.incomplete_details.reason == "max_output_tokens" + + +def test_gemini_candidate_with_finish_reason_no_content_streaming_chunk(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + chunk: Final = { + "candidates": [{"finishReason": "NO_IMAGE", "index": 0}], + "usageMetadata": {"promptTokenCount": 19, "candidatesTokenCount": 0, "totalTokenCount": 19}, + } + iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) + + streaming_chunk: Final = iterator.chunk_parser(chunk) + + assert len(streaming_chunk.choices) == 1 + assert streaming_chunk.choices[0].finish_reason == "content_filter" + assert streaming_chunk.choices[0].delta.content is None + assert streaming_chunk.choices[0].delta.tool_calls is None + + +def test_gemini_multi_candidate_messages_do_not_share_state(): + config: Final = VertexGeminiConfig() + completion_response: Final = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "Let me check the weather.", "thought": True}, + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}, + ], + }, + "finishReason": "STOP", + "index": 0, + }, + { + "content": {"role": "model", "parts": [{"text": "It is sunny in Paris."}]}, + "finishReason": "STOP", + "index": 1, + }, + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 20, "totalTokenCount": 30}, + } + + resp: Final = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + assert len(resp.choices) == 2 + assert resp.choices[0].finish_reason == "tool_calls" + assert resp.choices[0].message.tool_calls[0].function.name == "get_weather" + assert resp.choices[0].message.reasoning_content == "Let me check the weather." + assert resp.choices[1].finish_reason == "stop" + assert resp.choices[1].message.content == "It is sunny in Paris." + assert resp.choices[1].message.tool_calls is None + assert getattr(resp.choices[1].message, "reasoning_content", None) is None + assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 849c997838d..52d06acfd64 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4909,7 +4909,7 @@ class TestStreamingSnapshotItemIds: def test_transform_chat_completion_response_incomplete_details(): - from openai.types.responses.response import IncompleteDetails + from litellm.types.llms.openai import IncompleteDetails resp_length = ModelResponse( id="resp-length", From 55249c7128bca0c47641f561c30e96040e734d69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:58:58 -0700 Subject: [PATCH 14/33] fix: set vertex gemma-4-26b-a4b-it-maas context window to 262144 --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- .../gemma/test_vertex_ai_gemma_global_endpoint.py | 12 +++++++++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 377e30a23b7..ef4fce37f23 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -50721,7 +50721,7 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 377e30a23b7..ef4fce37f23 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -50721,7 +50721,7 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py index e9b58622a4b..9d08daa0a81 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py @@ -30,7 +30,7 @@ from litellm.types.llms.vertex_ai import VertexPartnerProvider _GEMMA_MODEL_COST_ENTRY = { "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -180,6 +180,16 @@ class TestCreateVertexURLGemma: # --------------------------------------------------------------------------- +def test_gemma_maas_context_window_matches_google(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + info = litellm.get_model_info("vertex_ai/google/gemma-4-26b-a4b-it-maas") + + assert info["max_input_tokens"] == 262144 + assert info["max_output_tokens"] == 128000 + + # --------------------------------------------------------------------------- # Integration tests: verify payloads reach the global OpenAI endpoint # From 417a88daedf6e5e4d2de62fab5a33f4ec5a6ae3f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:13:07 -0700 Subject: [PATCH 15/33] fix(responses): carry dict-valued reasoning_effort and keep the frame type on websocket defaults A deployment whose reasoning_effort is an object is copied through as reasoning the way the HTTP mapper does it instead of being dropped, and the relay re-asserts the response.create frame type after merging extra_body so a type key inside it can never replace it. The lazy OpenAPI snapshot goes back to main: the earlier regeneration came from a Python 3.14 interpreter dedenting docstrings, which CI on 3.12 rejects --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/responses/main.py | 22 +++++++----- litellm/responses/streaming_iterator.py | 2 +- .../test_responses_websocket_all_providers.py | 36 +++++++++++++++++++ 4 files changed, 51 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2aa2cf15ac1..8a8d08c6887 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19394,7 +19394,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 6064fa66d91..a5912bb42b1 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2262,24 +2262,28 @@ def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: return metadata -_EXTRA_BODY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | dict[str, object] | None: + if kwargs.get("reasoning") is not None: + return None + reasoning_effort: Final = kwargs.get("reasoning_effort") + if isinstance(reasoning_effort, str): + return LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) + return _JSON_OBJECT_ADAPTER.validate_python(reasoning_effort) if isinstance(reasoning_effort, Mapping) else None def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults: - reasoning_effort: Final = kwargs.get("reasoning_effort") - mapped_reasoning: Final = ( - LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) - if kwargs.get("reasoning") is None and isinstance(reasoning_effort, str) - else None - ) + default_reasoning: Final = _deployment_reasoning_default(kwargs) candidate_params: Final[dict[str, object]] = { **kwargs, - **({"reasoning": mapped_reasoning} if mapped_reasoning is not None else {}), + **({"reasoning": default_reasoning} if default_reasoning is not None else {}), } fill_missing: Final = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(candidate_params) return ResponsesWebSocketRequestDefaults( fill_missing=MappingProxyType(dict(fill_missing)), - overrides=MappingProxyType(_EXTRA_BODY_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), + overrides=MappingProxyType(_JSON_OBJECT_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 1c82df8c664..8f65552ec20 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1883,7 +1883,7 @@ class ResponsesWebSocketStreaming: nested: Final = msg_obj.get("response") if _is_json_object(nested): return {**msg_obj, "response": self.request_defaults.merged_into(nested)} - return self.request_defaults.merged_into(msg_obj) + return {**self.request_defaults.merged_into(msg_obj), "type": msg_obj["type"]} async def _mask_response_create(self, message: str) -> str: """ diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 918c4a33d40..b05904acdcb 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1252,6 +1252,42 @@ class TestNativeWebSocketDeploymentDefaults: assert dict(defaults.fill_missing) == {"reasoning": {"effort": "low"}} assert dict(defaults.overrides) == {} + def test_builder_copies_dict_valued_reasoning_effort_like_the_http_path(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + {"model": "gpt-5-pro", "reasoning_effort": {"effort": "xhigh", "summary": "auto"}} + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "xhigh", "summary": "auto"}} + + @pytest.mark.asyncio + async def test_extra_body_type_key_never_replaces_the_frame_type(self): + from types import MappingProxyType + + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + handler = _make_streaming( + authorized_model="gpt-5-pro", + request_defaults=ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType({}), + overrides=MappingProxyType({"type": "session.update", "provider_default": "configured"}), + ), + ) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps({"type": "response.create", "model": "gpt-5-pro", "input": "hi"}) + ) + ) + + assert forwarded == { + "type": "response.create", + "model": "gpt-5-pro", + "input": "hi", + "provider_default": "configured", + } + @pytest.mark.asyncio async def test_flat_frame_gets_defaults_client_keys_win_extra_body_overrides(self): handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) From 1adbfbfbb11b534df9c4c75f5937c51ca2d2ec1b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:28:45 -0700 Subject: [PATCH 16/33] fix: strip eager_input_streaming for non-Claude providers next to input_examples --- litellm/main.py | 46 +++++++++++----------- tests/test_litellm/test_main.py | 70 +++++++++++++++++++++++++-------- 2 files changed, 76 insertions(+), 40 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 859e0df142c..ac8fa507728 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1146,37 +1146,35 @@ def responses_api_bridge_check( return model_info, model -def _should_allow_input_examples(custom_llm_provider: str | None, model: str) -> bool: +_ANTHROPIC_ONLY_TOOL_KEYS: Final = frozenset({"input_examples", "eager_input_streaming"}) + + +def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: if custom_llm_provider == "anthropic": return True - if custom_llm_provider == "azure_ai" or custom_llm_provider == "bedrock" or custom_llm_provider == "vertex_ai": - return "claude" in model.lower() + model_lower: Final = model.lower() + if custom_llm_provider == "bedrock": + return "claude" in model_lower or ("arn:" in model_lower and ":bedrock:" in model_lower) + if custom_llm_provider == "azure_ai" or custom_llm_provider == "vertex_ai": + return "claude" in model_lower return False -def _drop_input_examples_from_tool(tool: dict) -> dict: - tool_copy: Final = tool.copy() - tool_copy.pop("input_examples", None) - function = tool_copy.get("function") - if isinstance(function, dict): - function = function.copy() - function.pop("input_examples", None) - tool_copy["function"] = function - return tool_copy +def _without_anthropic_only_tool_keys(tool: dict) -> dict: + kept: Final = {key: value for key, value in tool.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS} + function: Final = tool.get("function") + if not isinstance(function, dict): + return kept + return { + **kept, + "function": {key: value for key, value in function.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS}, + } -def _drop_input_examples_from_tools( - tools: list[dict] | None, -) -> list[dict] | None: +def _drop_anthropic_only_tool_keys(tools: list[dict] | None) -> list[dict] | None: if tools is None: return None - cleaned_tools: Final[list[dict]] = [] - for tool in tools: - if isinstance(tool, dict): - cleaned_tools.append(_drop_input_examples_from_tool(tool)) - else: - cleaned_tools.append(tool) - return cleaned_tools + return [_without_anthropic_only_tool_keys(tool) if isinstance(tool, dict) else tool for tool in tools] class _ProxyAuthHeadersProvider(Protocol): @@ -5360,8 +5358,8 @@ def completion( api_base=api_base, ) - if not _should_allow_input_examples(custom_llm_provider=custom_llm_provider, model=model): - tools = _drop_input_examples_from_tools(tools=tools) + if not _is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): + tools = _drop_anthropic_only_tool_keys(tools=tools) if provider_specific_header is not None: headers.update( diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index cbbac3d247f..bd115c699d5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -349,28 +349,66 @@ def test_bedrock_latency_optimized_inference(): assert json_data["performanceConfig"]["latency"] == "optimized" -def test_strip_input_examples_for_non_anthropic_providers(): +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected"), + [ + ("anthropic", "claude-sonnet-5", True), + ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), + ("bedrock", "us.amazon.nova-2-lite-v1:0", False), + ("vertex_ai", "claude-sonnet-5", True), + ("vertex_ai", "gemini-3.8-flash", False), + ("azure_ai", "claude-sonnet-4-6", True), + ("azure_ai", "gpt-5.6", False), + ("openai", "gpt-5.6", False), + ("gemini", "gemini-3.8-flash", False), + ], +) +def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): + assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected + + +@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) +def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): tools = [ - { - "type": "function", - "name": "example_tool", - "input_examples": [{"foo": "bar"}], - "function": { - "name": "example_tool", - "input_examples": [{"foo": "bar"}], - }, - } + {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, + "opaque_tool", ] - assert not litellm_main._should_allow_input_examples( - custom_llm_provider="openai", model="gpt-4o-mini" + cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) + + assert cleaned == [ + {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, + "opaque_tool", + ] + assert tools[0][key] is True + assert tools[0]["function"][key] is True + + +def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): + api_base: Final = "http://localhost:12346/v1" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) ) - cleaned = litellm_main._drop_input_examples_from_tools(tools=tools) + litellm.completion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Write the file"}], + tools=[ + { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, + "eager_input_streaming": True, + } + ], + api_base=api_base, + api_key="fake_openai_api_key", + ) - assert isinstance(cleaned, list) - assert "input_examples" not in cleaned[0] - assert "input_examples" not in cleaned[0]["function"] + assert mock_route.called + sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] + assert "eager_input_streaming" not in sent_tool + assert sent_tool["function"]["name"] == "write_file" def test_custom_provider_with_extra_headers(): From 66c01cf35c14d06877f7560af7604a5c8e36154c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:39:02 -0700 Subject: [PATCH 17/33] refactor(responses): map finish reasons to incomplete_details through a lookup table The match statement in _incomplete_details_for_finish_reason tripped CodeQL's mixed explicit and implicit returns alert (code-scanning 12640). A module-level MappingProxyType keyed by finish reason gives the same three mappings with one explicit return path --- .../transformation.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 961fffd3d42..7337beba45c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -112,6 +112,9 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +_INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( + MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) +) @dataclass(frozen=True, slots=True) @@ -2303,13 +2306,10 @@ class LiteLLMCompletionResponsesConfig: ) -> IncompleteDetails | None: if existing is not None: return existing - match finish_reason: - case "length": - return IncompleteDetails(reason="max_output_tokens") - case "content_filter" | "refusal": - return IncompleteDetails(reason="content_filter") - case _: - return None + if finish_reason is None: + return None + reason: Final = _INCOMPLETE_REASON_BY_FINISH_REASON.get(finish_reason) + return IncompleteDetails(reason=reason) if reason is not None else None @staticmethod def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) -> str: From 44034c1d5e25c8f1888dfd3f13e847cc44860fab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:43:07 -0700 Subject: [PATCH 18/33] test: source the gemma context window limits and isolate the cost map cache --- .../gemma/test_vertex_ai_gemma_global_endpoint.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py index 9d08daa0a81..85a2124ab02 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py @@ -180,12 +180,11 @@ class TestCreateVertexURLGemma: # --------------------------------------------------------------------------- -def test_gemma_maas_context_window_matches_google(monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - +def test_gemma_maas_context_window_matches_google(local_model_cost_map): info = litellm.get_model_info("vertex_ai/google/gemma-4-26b-a4b-it-maas") + # 262,144 context length and 128,000 maximum output per Google's model page, checked 2026-09-18: + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/maas/google/gemma-4-26b-a4b-it assert info["max_input_tokens"] == 262144 assert info["max_output_tokens"] == 128000 From b4bfd92a2a4ab11df53d886704b621b6e8cd2339 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 13:36:06 -0700 Subject: [PATCH 19/33] refactor(rust): route-neutral callback contract Every legacy callback call from callbacks-legacy now goes through one typed Python shim, litellm.rust_bridge.legacy_callbacks, the only Python module the crate reaches. Before, the crate called Logging methods, litellm.utils hooks, the logging worker, the executor and several litellm globals directly, and its tests retyped those signatures by hand, so an outdated fake could accept a call the real code rejects. python_contract.json lists each shim function's parameters: a Python test pins it to the real signatures and a Rust test pins it to the Rust enum. The lifecycle contract changes to match the Python @client wrapper: - the driver emits CallEvent::Started before begin, so every host sees one start time - RequestContext carries the route-resolved api_key, so legacy pre_call and post_call receive it, and post_call's additional_args match the Python OCR path - Passthrough and its re-aliasing are gone - async deployment hooks always run, and the "no callbacks" shortcut that skipped the logging payload is removed, as in the Python path The OCR api_key is a SecretValue from the wire request onward, so Debug output upstream of the callback contract cannot leak it. host-python's RouteHost now classifies native failures once through classify, and host ops return HostOpError. The OCR route host keeps main's public errors by sending both through the existing Python map_failure. --- litellm-rust/Cargo.lock | 3 + litellm-rust/crates/auth/src/secret.rs | 4 +- .../crates/callbacks-legacy/AGENTS.md | 14 +- .../crates/callbacks-legacy/Cargo.toml | 4 + .../callbacks-legacy/python_contract.json | 98 +++++ .../crates/callbacks-legacy/src/adapter.rs | 101 +++-- .../crates/callbacks-legacy/src/call.rs | 43 +-- .../crates/callbacks-legacy/src/callbacks.rs | 242 +++--------- .../callbacks-legacy/src/legacy_python.rs | 155 ++++++++ .../crates/callbacks-legacy/src/lib.rs | 5 +- .../crates/callbacks-legacy/src/logger.rs | 91 +---- .../callbacks-legacy/src/preparation.rs | 44 +-- .../crates/callbacks-legacy/tests/deferred.rs | 18 +- .../tests/deployment_hooks.rs | 18 +- .../crates/callbacks-legacy/tests/payload.rs | 113 +++--- .../crates/callbacks-legacy/tests/support.rs | 135 ++++--- .../crates/callbacks-legacy/tests/terminal.rs | 46 ++- litellm-rust/crates/callbacks/Cargo.toml | 1 + litellm-rust/crates/callbacks/src/event.rs | 77 +--- litellm-rust/crates/callbacks/src/run.rs | 36 +- litellm-rust/crates/core/src/ocr/handler.rs | 13 +- litellm-rust/crates/core/src/ocr/mod.rs | 4 +- litellm-rust/crates/core/src/ocr/prepare.rs | 4 +- .../crates/core/src/ocr/provider_config.rs | 46 ++- litellm-rust/crates/core/src/ocr/types.rs | 18 +- litellm-rust/crates/core/src/ocr/wire.rs | 4 +- .../tests/azure_document_intelligence_ocr.rs | 2 +- litellm-rust/crates/core/tests/ocr.rs | 22 +- .../crates/core/tests/ocr/document.rs | 152 ++++++++ .../crates/core/tests/ocr/passthrough.rs | 282 -------------- litellm-rust/crates/core/tests/ocr/support.rs | 10 +- litellm-rust/crates/host-python/AGENTS.md | 5 +- .../crates/host-python/src/adapter.rs | 55 ++- .../crates/host-python/src/argument.rs | 51 +++ litellm-rust/crates/host-python/src/driver.rs | 364 +++++++++++++----- litellm-rust/crates/host-python/src/lib.rs | 8 +- .../document_intelligence/transformation.rs | 21 +- .../llms/src/azure_ai/ocr/transformation.rs | 21 +- .../llms/src/base_llm/ocr/transformation.rs | 17 +- .../llms/src/cohere/ocr/transformation.rs | 6 +- .../llms/src/custom_httpx/llm_http_handler.rs | 32 +- .../llms/src/mistral/ocr/transformation.rs | 6 +- .../llms/src/reducto/ocr/transformation.rs | 8 +- .../llms/src/vertex_ai/ocr/transformation.rs | 5 +- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- .../python-bridge/src/routes/ocr/host.rs | 61 +-- .../python-bridge/src/routes/ocr/project.rs | 10 +- litellm/litellm_core_utils/litellm_logging.py | 43 +-- litellm/rust_bridge/legacy_callbacks.py | 291 ++++++++++---- .../rust_bridge/test_legacy_callbacks.py | 19 +- tests/test_litellm_rust/ocr/test_lifecycle.py | 31 +- 51 files changed, 1564 insertions(+), 1297 deletions(-) create mode 100644 litellm-rust/crates/callbacks-legacy/python_contract.json create mode 100644 litellm-rust/crates/callbacks-legacy/src/legacy_python.rs create mode 100644 litellm-rust/crates/core/tests/ocr/document.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/passthrough.rs create mode 100644 litellm-rust/crates/host-python/src/argument.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c359ca19986..0265e0adbc2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2031,6 +2031,7 @@ dependencies = [ name = "litellm-callbacks" version = "0.1.0" dependencies = [ + "litellm-auth", "rstest", "serde_json", "tokio", @@ -2040,11 +2041,13 @@ dependencies = [ name = "litellm-callbacks-legacy" version = "0.1.0" dependencies = [ + "litellm-auth", "litellm-callbacks", "litellm-host-python", "pyo3", "rstest", "serde_json", + "strum", ] [[package]] diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth/src/secret.rs index 3ecb0a835ee..a07fe3eaad9 100644 --- a/litellm-rust/crates/auth/src/secret.rs +++ b/litellm-rust/crates/auth/src/secret.rs @@ -1,6 +1,8 @@ +use serde::Deserialize; use veil::Redact; -#[derive(Redact, Clone)] +#[derive(Redact, Clone, Deserialize)] +#[serde(transparent)] pub struct SecretValue(#[redact(with = "[REDACTED]")] String); impl SecretValue { diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy/AGENTS.md index e4762d3037a..e184e2fb415 100644 --- a/litellm-rust/crates/callbacks-legacy/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy/AGENTS.md @@ -1,15 +1,17 @@ - Target invariants, not completion claims - Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits) - - The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call + - The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call +- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`) + - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it + - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy -- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it - - A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case - - A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run +- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does + - Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's - Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - - Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None` + - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields` + - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view - Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index 96c9c9ed560..3cf9382f857 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -9,8 +9,12 @@ autotests = false [dependencies] litellm-callbacks.workspace = true litellm-host-python.workspace = true + pyo3.workspace = true +strum.workspace = true +serde_json.workspace = true [dev-dependencies] +litellm-auth.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json new file mode 100644 index 00000000000..840c0abfa45 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -0,0 +1,98 @@ +{ + "setup": [ + "call_type", + "args", + "kwargs", + "start_time", + "asynchronous" + ], + "check_limits": [ + "kwargs" + ], + "finalize": [ + "response", + "logger", + "kwargs", + "start_time", + "end_time" + ], + "update_logging": [ + "logger", + "kwargs", + "model", + "optional_params", + "litellm_params", + "custom_llm_provider" + ], + "pre_call": [ + "logger", + "input", + "api_key", + "additional_args" + ], + "post_call": [ + "logger", + "original_response", + "api_key", + "additional_args" + ], + "defers_async_logging": [ + "logger" + ], + "defer_success": [ + "logger", + "pending" + ], + "sync_success_for_async_call": [ + "logger", + "response", + "start", + "end" + ], + "failure_handler": [ + "logger", + "error", + "start", + "end", + "asynchronous" + ], + "submit_success": [ + "logger", + "response", + "start", + "end" + ], + "async_success_handler": [ + "logger", + "response", + "start", + "end" + ], + "enqueue_logging": [ + "coroutine" + ], + "restore_context": [ + "logger" + ], + "custom_pricing_fields": [], + "is_internal_call": [], + "credential_list": [], + "warn_unknown_credential": [ + "name", + "loaded" + ], + "before_deployment_call": [ + "kwargs", + "call_type" + ], + "after_deployment_success": [ + "kwargs", + "response", + "call_type" + ], + "after_deployment_failure": [ + "kwargs", + "error", + "call_type" + ] +} diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index df346506094..7204207cd62 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -4,7 +4,7 @@ use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest}; use litellm_host_python::{ - AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py, + LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -12,6 +12,7 @@ use pyo3::{ prelude::*, types::PyDict, }; +use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, @@ -43,7 +44,7 @@ pub struct LegacyLogging { response: Option>, error: Option>, body: Option>, - headers: Option>, + context: Option, asynchronous: bool, internal: bool, pending: Option, @@ -76,7 +77,7 @@ impl LegacyLogging { response: None, error: None, body: None, - headers: None, + context: None, asynchronous, internal: false, pending: None, @@ -85,8 +86,8 @@ impl LegacyLogging { /// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never /// runs them. - fn deployment_hooks(&self, py: Python<'_>) -> PyResult { - Ok(self.asynchronous && DeploymentHooks::needed(py)?) + fn runs_deployment_hooks(&self) -> bool { + self.asynchronous } fn logger(&self) -> PyResult<&PythonLogger> { @@ -95,13 +96,13 @@ impl LegacyLogging { }) } - fn prepare(&mut self, py: Python<'_>) -> PyResult { + fn prepare(&mut self, py: Python<'_>) -> PyResult { let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind(); self.call.set_kwargs(prepared); - Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py))) + Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py))) } - fn finalize(&mut self, py: Python<'_>) -> PyResult { + fn finalize(&mut self, py: Python<'_>) -> PyResult { finalize( py, &self.response, @@ -112,7 +113,7 @@ impl LegacyLogging { )?; self.response .as_ref() - .map(|response| AdapterStep::Response(response.clone_ref(py))) + .map(|response| LifecycleStep::Response(response.clone_ref(py))) .ok_or_else(missing_state) } @@ -145,9 +146,7 @@ impl LegacyLogging { .get_item("fallbacks")? .is_none_or(|value| value.is_none()) { - if !logger.callbacks_needed(py, "async_success")? { - logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; - } else if logger.defers_async_logging(py) { + if logger.defers_async_logging(py) { let pending = Py::new( py, PendingLogging { @@ -165,12 +164,12 @@ impl LegacyLogging { /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. - fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { + fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { let (Some(logger), Some(error)) = (&self.logger, &self.error) else { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); }; if self.asynchronous && self.internal { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false) && is_cancellation(py, &failure) @@ -178,27 +177,27 @@ impl LegacyLogging { return Err(failure); } if !self.asynchronous { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } match logger.failure(py, error, &self.start, &self.end, true) { Ok(Some(awaitable)) => { self.pending = Some(Pending::AsyncFailure); - Ok(AdapterStep::Await(awaitable)) + Ok(LifecycleStep::Await(awaitable)) } - Ok(None) => Ok(AdapterStep::Done), + Ok(None) => Ok(LifecycleStep::Done), Err(failure) if is_cancellation(py, &failure) => Err(failure), - Err(_) => Ok(AdapterStep::Done), + Err(_) => Ok(LifecycleStep::Done), } } } -impl CallbackAdapter for LegacyLogging { +impl PythonLifecycle for LegacyLogging { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult { + ) -> PyResult { self.call.set_kwargs(arguments); self.start = datetime(py, started_at)?; self.internal = is_internal_call(py)?; @@ -212,9 +211,9 @@ impl CallbackAdapter for LegacyLogging { )?; self.logger = Some(result.logger()?); self.call.set_kwargs(result.kwargs()?); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPreCall); - return Ok(AdapterStep::Await(DeploymentHooks::before_call( + return Ok(LifecycleStep::Await(DeploymentHooks::before_call( py, self.call.kwargs(), self.surface.call_type, @@ -228,18 +227,16 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult { + ) -> PyResult { let logger = self.logger()?; logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; - if !logger.callbacks_needed(py, "payload")? { - logger.record_api_call_start(py)?; - return Ok(AdapterStep::Wire(wire)); - } let body = to_py(py, &wire.body)? .into_bound(py) .cast_into::()?; - for name in context.passthrough_fields.iter() { - if let Some(value) = self.call.lookup(py, name)? { + for (name, sent) in wire.body.as_object().into_iter().flatten() { + if let Some(value) = self.call.lookup(py, name)? + && from_py::(&value).is_ok_and(|caller| caller == *sent) + { body.set_item(name, value)?; } } @@ -248,12 +245,11 @@ impl CallbackAdapter for LegacyLogging { headers.set_item(name, value)?; } self.body = Some(body.clone().unbind()); - self.headers = Some(headers.clone().unbind()); - let api_key = self.call.lookup(py, "api_key")?; + self.context = Some(context.clone()); self.logger()?.pre_call( py, self.surface.input_description, - api_key.as_ref(), + context.api_key.as_ref().map(|api_key| api_key.expose()), &body, &headers, &wire.url, @@ -262,7 +258,7 @@ impl CallbackAdapter for LegacyLogging { .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { body: from_py(&body)?, headers, ..*wire @@ -274,12 +270,12 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult { + ) -> PyResult { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPostCall); - return Ok(AdapterStep::Await(DeploymentHooks::after_success( + return Ok(LifecycleStep::Await(DeploymentHooks::after_success( py, self.call.kwargs(), &self.response, @@ -294,31 +290,35 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult { + ) -> PyResult { match (event, public) { + (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), (CallEvent::ResponseReceived { raw }, _) => { - let logger = self.logger()?; - if logger.callbacks_needed(py, "payload")? { - logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?; - } - Ok(AdapterStep::Done) + let api_key = self + .context + .as_ref() + .and_then(|context| context.api_key.as_ref()) + .map(|api_key| api_key.expose()); + self.logger()? + .post_call(py, &raw.body, api_key, self.body.as_ref())?; + Ok(LifecycleStep::Done) } (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); self.dispatch_success(py)?; - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); if *origin == FailureOrigin::Call && self.logger.is_some() - && self.deployment_hooks(py)? + && self.runs_deployment_hooks() { let error = self.error.as_ref().ok_or_else(missing_state)?; self.pending = Some(Pending::DeploymentFailure); - return Ok(AdapterStep::Await(DeploymentHooks::after_failure( + return Ok(LifecycleStep::Await(DeploymentHooks::after_failure( py, self.call.kwargs(), error, @@ -331,7 +331,7 @@ impl CallbackAdapter for LegacyLogging { } } - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { self.call @@ -345,7 +345,7 @@ impl CallbackAdapter for LegacyLogging { Pending::DeploymentFailure => self.dispatch_failure(py), Pending::AsyncFailure => match result { Err(failure) if is_cancellation(py, &failure) => Err(failure), - _ => Ok(AdapterStep::Done), + _ => Ok(LifecycleStep::Done), }, } } @@ -357,7 +357,7 @@ impl CallbackAdapter for LegacyLogging { error.write_unraisable(py, None); } self.body = None; - self.headers = None; + self.context = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -369,8 +369,7 @@ impl CallbackAdapter for LegacyLogging { visit.call(&self.end)?; visit.call(&self.response)?; visit.call(&self.error)?; - visit.call(&self.body)?; - visit.call(&self.headers) + visit.call(&self.body) } } diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy/src/call.rs index 59090ee8d60..bd1e2525d3d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy/src/call.rs @@ -4,7 +4,7 @@ //! this crate holds them. use litellm_callbacks::{machine::Machine, route::Route}; -use litellm_host_python::{RouteHost, run_call}; +use litellm_host_python::{RouteHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -63,21 +63,6 @@ impl PublicCall { } } -/// The caller's own object for a public argument, as every legacy reader resolves it: the -/// keyword if given, even an explicit `None`, else the bound request's attribute. A route -/// host projecting from the prepared keyword view uses the same rule, so the callbacks -/// and the provider see one object per argument. -pub fn lookup<'py>( - kwargs: &Bound<'py, PyDict>, - request: &Bound<'py, PyAny>, - name: &str, -) -> PyResult>> { - if let Some(value) = kwargs.get_item(name)? { - return Ok(Some(value)); - } - request.getattr_opt(name) -} - /// Runs one native call under the legacy `Logging` contract: the route host projects from /// the keyword view the contract prepares, and the contract observes the call. pub fn run_legacy_call( @@ -121,32 +106,6 @@ mod tests { (call, locals) } - #[test] - fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { - Python::initialize(); - Python::attach(|py| { - let (call, locals) = capture( - py, - c" -key = object() -document = {'type': 'document_url'} -class Request: - api_key = 'from-request' - api_base = 'from-request' - document = document -request = Request() -kwargs = {'api_key': key, 'api_base': None} -", - ); - let key = locals.get_item("key").unwrap().unwrap(); - let document = locals.get_item("document").unwrap().unwrap(); - assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key)); - assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none()); - assert!(call.lookup(py, "document").unwrap().unwrap().is(&document)); - assert!(call.lookup(py, "model").unwrap().is_none()); - }); - } - #[test] fn capture_copies_the_keyword_dict_without_copying_its_values() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index aa586013e75..9464f1d6612 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -6,11 +6,10 @@ use litellm_callbacks::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use crate::legacy_python::{Logging, Wrapper}; use crate::logger::PythonLogger; pub trait LegacyCallbacks { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult; - /// `Logging.update_from_kwargs`: what the logger is told about the request it is /// about to see, with consumed credentials redacted. fn update_from_kwargs( @@ -21,26 +20,24 @@ pub trait LegacyCallbacks { context: &RequestContext, ) -> PyResult<()>; - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>; - - /// `Logging.pre_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.pre_call`. fn pre_call( &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, ) -> PyResult<()>; - /// `Logging.post_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.post_call`. fn post_call( &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, - headers: Option<&Py>, ) -> PyResult<()>; fn defers_async_logging(&self, py: Python<'_>) -> bool; @@ -82,16 +79,6 @@ pub trait LegacyCallbacks { } impl LegacyCallbacks for PythonLogger { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { - if !self.bridge_owned() { - return Ok(true); - } - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("callbacks_needed")? - .call1((self.object(py), phase))? - .extract() - } - fn update_from_kwargs( &self, py: Python<'_>, @@ -100,18 +87,13 @@ impl LegacyCallbacks for PythonLogger { context: &RequestContext, ) -> PyResult<()> { let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect(); - let update = PyDict::new(py); - update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?; - update.set_item("model", &context.model)?; - update.set_item( - "optional_params", - redact( - py, - &to_py(py, &context.optional_params)? - .into_bound(py) - .cast_into::()?, - &secret_fields, - )?, + let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?; + let optional_params = redact( + py, + &to_py(py, &context.optional_params)? + .into_bound(py) + .cast_into::()?, + &secret_fields, )?; let params = PyDict::new(py); params.set_item( @@ -131,15 +113,17 @@ impl LegacyCallbacks for PythonLogger { params.set_item(name, value)?; } } - update.set_item("litellm_params", params)?; - update.set_item("custom_llm_provider", &context.custom_llm_provider)?; - self.object(py) - .call_method("update_from_kwargs", (), Some(&update))?; - Ok(()) - } - - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> { - self.object(py).call_method0("record_api_call_start_time")?; + Logging::Update.call( + py, + ( + self.object(py), + redacted_kwargs, + &context.model, + optional_params, + params, + &context.custom_llm_provider, + ), + )?; Ok(()) } @@ -147,7 +131,7 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, @@ -156,17 +140,7 @@ impl LegacyCallbacks for PythonLogger { additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; additional.set_item("api_base", url)?; - let kwargs = PyDict::new(py); - kwargs.set_item("input", input)?; - kwargs.set_item("api_key", api_key)?; - kwargs.set_item("additional_args", &additional)?; - if self.callbacks_needed(py, "input")? { - self.object(py).call_method("pre_call", (), Some(&kwargs))?; - } else { - self.object(py) - .call_method("_pre_call", (), Some(&kwargs))?; - self.record_api_call_start(py)?; - } + Logging::PreCall.call(py, (self.object(py), input, api_key, &additional))?; Ok(()) } @@ -174,37 +148,28 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, - headers: Option<&Py>, ) -> PyResult<()> { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; - additional.set_item("headers", headers)?; - if self.callbacks_needed(py, "input")? { - let kwargs = PyDict::new(py); - kwargs.set_item("original_response", original_response)?; - kwargs.set_item("additional_args", &additional)?; - self.object(py) - .call_method("post_call", (), Some(&kwargs))?; - } else { - let response = py - .import("json")? - .call_method1("dumps", (original_response,))?; - self.object(py).call_method1( - "record_post_call", - (response, py.None(), py.None(), additional), - )?; - } + Logging::PostCall.call( + py, + (self.object(py), original_response, api_key, &additional), + )?; Ok(()) } + fn defers_async_logging(&self, py: Python<'_>) -> bool { - self.object(py) - .getattr("_defer_async_logging") - .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + Logging::DefersAsync + .call(py, (self.object(py),)) + .and_then(|value| value.extract()) + .unwrap_or(false) } fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - self.object(py).setattr("_native_pending_logging", pending) + Logging::DeferSuccess.call(py, (self.object(py), pending))?; + Ok(()) } fn sync_success_for_async_call( @@ -214,13 +179,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success_async")? { - return Ok(()); - } - self.object(py).call_method1( - "handle_sync_success_callbacks_for_async_calls", - (response, start, end), - )?; + Logging::SyncSuccessForAsyncCall.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -232,34 +191,11 @@ impl LegacyCallbacks for PythonLogger { end: &Option>, asynchronous: bool, ) -> PyResult>> { - if !self.callbacks_needed( - py, - if asynchronous { - "async_failure" - } else { - "sync_failure" - }, - )? { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("failure_bookkeeping")? - .call1((self.object(py), error, start, end, asynchronous))?; - return Ok(None); - } - let trace = py - .import("traceback")? - .getattr("format_exception")? - .call1((error,))?; - let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?; - let value = self.object(py).call_method1( - if asynchronous { - "async_failure_handler" - } else { - "failure_handler" - }, - (error, trace, start, end), - )?; + let value = + Logging::FailureHandler.call(py, (self.object(py), error, start, end, asynchronous))?; Ok(asynchronous.then(|| value.unbind())) } + fn submit_success( &self, py: Python<'_>, @@ -267,22 +203,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success")? { - return self.success_bookkeeping(py, response, start, end, false); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - py.import("litellm.litellm_core_utils.litellm_logging")? - .getattr("executor")? - .call_method1( - "submit", - ( - context.getattr("run")?, - self.object(py).getattr("success_handler")?, - response, - start, - end, - ), - )?; + Logging::SubmitSuccess.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -293,18 +214,9 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "async_success")? { - return self.success_bookkeeping(py, response, start, end, true); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - let worker = py - .import("litellm.litellm_core_utils.logging_worker")? - .getattr("GLOBAL_LOGGING_WORKER")? - .getattr("ensure_initialized_and_enqueue")?; - let coroutine = self - .object(py) - .call_method1("async_success_handler", (response, start, end))?; - let enqueue = context.call_method1("run", (worker, &coroutine)); + let coroutine = + Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") { @@ -315,14 +227,7 @@ impl LegacyCallbacks for PythonLogger { } fn custom_pricing_fields(py: Python<'_>) -> PyResult> { - py.import("litellm.types.utils")? - .getattr("CustomPricingLiteLLMParams")? - .getattr("model_fields")? - .cast_into::()? - .keys() - .iter() - .map(|name| name.extract::()) - .collect() + Logging::CustomPricingFields.call(py, ())?.extract() } fn redact( @@ -347,58 +252,5 @@ fn redact( /// Proxy-internal calls skip the legacy success fan-out. pub fn is_internal_call(py: Python<'_>) -> PyResult { - py.import("litellm._internal_context")? - .getattr("is_internal_call")? - .call_method0("get")? - .extract() -} - -#[cfg(test)] -mod tests { - use pyo3::types::PyDict; - - use super::*; - - fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger { - let locals = PyDict::new(py); - py.run( - c" -import sys -import types -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): - sys.modules.setdefault(name, types.ModuleType(name)) -legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -class Logger: - needed = {'input': False} -logger = Logger() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - PythonLogger::new( - locals.get_item("logger").unwrap().unwrap().unbind(), - bridge_owned, - ) - } - - #[test] - fn a_caller_owned_logger_is_observed_in_full() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, false); - assert!(logger.callbacks_needed(py, "input").unwrap()); - }); - } - - #[test] - fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, true); - assert!(!logger.callbacks_needed(py, "input").unwrap()); - assert!(logger.callbacks_needed(py, "payload").unwrap()); - }); - } + Wrapper::IsInternalCall.call(py, ())?.extract() } diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs new file mode 100644 index 00000000000..a924775070c --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs @@ -0,0 +1,155 @@ +use pyo3::prelude::*; +use strum::{IntoStaticStr, VariantArray}; + +const MODULE: &str = "litellm.rust_bridge.legacy_callbacks"; + +/// Every litellm Python internal the native call still borrows, grouped by the subsystem it +/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today +/// (span tracking, the standard logging payload, spend, callback fan-out) keeps working. +/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a +/// user's own callback is not borrowing and does not belong here. +/// +/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and +/// `python_contract.json` pins each function's parameters on both sides. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LegacyPython { + Wrapper(Wrapper), + Logging(Logging), + DeploymentHooks(DeploymentHooks), +} + +/// The `@client` wrapper around the call: `function_setup`, limits, credentials, +/// response metadata and the correlation context. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Wrapper { + #[strum(serialize = "setup")] + Setup, + #[strum(serialize = "check_limits")] + CheckLimits, + #[strum(serialize = "credential_list")] + CredentialList, + #[strum(serialize = "warn_unknown_credential")] + WarnUnknownCredential, + #[strum(serialize = "is_internal_call")] + IsInternalCall, + #[strum(serialize = "finalize")] + Finalize, + #[strum(serialize = "restore_context")] + RestoreContext, +} + +/// litellm's `Logging` object and the sync and async callback fan-out behind it. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Logging { + #[strum(serialize = "custom_pricing_fields")] + CustomPricingFields, + #[strum(serialize = "update_logging")] + Update, + #[strum(serialize = "pre_call")] + PreCall, + #[strum(serialize = "post_call")] + PostCall, + #[strum(serialize = "defers_async_logging")] + DefersAsync, + #[strum(serialize = "defer_success")] + DeferSuccess, + #[strum(serialize = "sync_success_for_async_call")] + SyncSuccessForAsyncCall, + #[strum(serialize = "submit_success")] + SubmitSuccess, + #[strum(serialize = "async_success_handler")] + AsyncSuccessHandler, + #[strum(serialize = "enqueue_logging")] + Enqueue, + #[strum(serialize = "failure_handler")] + FailureHandler, +} + +/// The `litellm.utils` fan-outs that run every callback's deployment hook. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum DeploymentHooks { + #[strum(serialize = "before_deployment_call")] + BeforeDeploymentCall, + #[strum(serialize = "after_deployment_success")] + AfterDeploymentSuccess, + #[strum(serialize = "after_deployment_failure")] + AfterDeploymentFailure, +} + +impl LegacyPython { + fn name(self) -> &'static str { + match self { + Self::Wrapper(function) => function.into(), + Self::Logging(function) => function.into(), + Self::DeploymentHooks(function) => function.into(), + } + } + + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + py.import(MODULE)?.getattr(self.name())?.call1(args) + } +} + +impl Wrapper { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Wrapper(self).call(py, args) + } +} + +impl Logging { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Logging(self).call(py, args) + } +} + +impl DeploymentHooks { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::DeploymentHooks(self).call(py, args) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use strum::VariantArray; + + use super::{DeploymentHooks, LegacyPython, Logging, Wrapper}; + use crate::test_support::PYTHON_CONTRACT; + + #[test] + fn every_borrowed_function_is_in_the_python_contract() { + let contract: serde_json::Map = + serde_json::from_str(PYTHON_CONTRACT).unwrap(); + let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); + let called: Vec<&str> = Wrapper::VARIANTS + .iter() + .map(|&function| LegacyPython::Wrapper(function)) + .chain( + Logging::VARIANTS + .iter() + .map(|&function| LegacyPython::Logging(function)), + ) + .chain( + DeploymentHooks::VARIANTS + .iter() + .map(|&function| LegacyPython::DeploymentHooks(function)), + ) + .map(LegacyPython::name) + .collect(); + assert_eq!(called.len(), declared.len(), "a function is borrowed twice"); + assert_eq!(called.into_iter().collect::>(), declared); + } +} diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy/src/lib.rs index 06783ac255d..42ffd545e2b 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy/src/lib.rs @@ -2,7 +2,7 @@ //! sync and async callback registries it fans out to, the deployment hooks, the deferred //! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name //! inheritance, budget and retry-count limits). All of it sits behind one -//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and +//! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and //! core never learn which Python object is on the other end. //! //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`] @@ -13,6 +13,7 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod legacy_python; mod logger; mod preparation; #[cfg(test)] @@ -21,7 +22,7 @@ mod test_support; pub(crate) use adapter::LegacyLogging; pub use adapter::LegacySurface; -pub use call::{PublicCall, lookup, run_legacy_call}; +pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; diff --git a/litellm-rust/crates/callbacks-legacy/src/logger.rs b/litellm-rust/crates/callbacks-legacy/src/logger.rs index a0e525000b8..061941f05b9 100644 --- a/litellm-rust/crates/callbacks-legacy/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy/src/logger.rs @@ -5,34 +5,25 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -/// The `Logging` instance one call fans out through, and who owns it. A logger the caller -/// handed in is observed in full, because the caller reads it after the call; one this -/// crate built through `function_setup` is elided wherever no registry needs it. +use crate::legacy_python::{self, Wrapper}; + +/// The `Logging` instance one call fans out through. pub struct PythonLogger { object: Py, - bridge_owned: bool, } impl PythonLogger { - pub(crate) fn new(object: Py, bridge_owned: bool) -> Self { - Self { - object, - bridge_owned, - } + pub(crate) fn new(object: Py) -> Self { + Self { object } } pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> { self.object.bind(py) } - pub(crate) fn bridge_owned(&self) -> bool { - self.bridge_owned - } - pub fn clone_ref(&self, py: Python<'_>) -> Self { Self { object: self.object.clone_ref(py), - bridge_owned: self.bridge_owned, } } @@ -40,34 +31,17 @@ impl PythonLogger { visit.call(&self.object) } - pub fn success_bookkeeping( - &self, - py: Python<'_>, - response: &Option>, - start: &Py, - end: &Option>, - asynchronous: bool, - ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("success_bookkeeping")? - .call1((self.object(py), response, start, end, asynchronous))?; - Ok(()) - } - pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> { - py.import("litellm.utils")? - .getattr("_restore_correlation_context_if_supported")? - .call1((self.object(py),))?; + Wrapper::RestoreContext.call(py, (self.object(py),))?; Ok(()) } } -/// A bare Python object was not obtained from `setup`, so it is caller-owned. impl FromPyObject<'_, '_> for PythonLogger { type Error = PyErr; fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult { - Ok(Self::new(object.to_owned().unbind(), false)) + Ok(Self::new(object.to_owned().unbind())) } } @@ -75,9 +49,7 @@ pub struct SetupResult<'py>(Bound<'py, PyAny>); impl SetupResult<'_> { pub fn logger(&self) -> PyResult { - let object = self.0.getattr("logger")?.unbind(); - let bridge_owned = self.0.getattr("bridge_owned")?.extract()?; - Ok(PythonLogger::new(object, bridge_owned)) + Ok(PythonLogger::new(self.0.getattr("logger")?.unbind())) } pub fn kwargs(&self) -> PyResult> { @@ -93,9 +65,8 @@ pub fn setup<'py>( start: &Py, asynchronous: bool, ) -> PyResult> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("setup")? - .call1((call_type, args, kwargs, start, asynchronous)) + Wrapper::Setup + .call(py, (call_type, args, kwargs, start, asynchronous)) .map(SetupResult) } @@ -107,30 +78,20 @@ pub fn finalize( start: &Py, end: &Option>, ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("finalize")? - .call1((response, logger.object(py), kwargs, start, end))?; + Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?; Ok(()) } pub struct DeploymentHooks; impl DeploymentHooks { - pub fn needed(py: Python<'_>) -> PyResult { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("deployment_callbacks_needed")? - .call0()? - .extract() - } - pub fn before_call( py: Python<'_>, kwargs: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_pre_call_deployment_hook")? - .call1((kwargs, call_type)) + legacy_python::DeploymentHooks::BeforeDeploymentCall + .call(py, (kwargs, call_type)) .map(Bound::unbind) } @@ -140,9 +101,8 @@ impl DeploymentHooks { response: &Option>, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_success_deployment_hook")? - .call1((kwargs, response, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentSuccess + .call(py, (kwargs, response, call_type)) .map(Bound::unbind) } @@ -152,9 +112,8 @@ impl DeploymentHooks { error: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_failure_deployment_hook")? - .call1((kwargs, error, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentFailure + .call(py, (kwargs, error, call_type)) .map(Bound::unbind) } } @@ -185,10 +144,6 @@ class Setup: reads.append('logger') return logger @property - def bridge_owned(self): - reads.append('bridge_owned') - return True - @property def kwargs(self): reads.append('kwargs') return [] @@ -206,7 +161,6 @@ result = Setup() .object(py) .is(locals.get_item("logger").unwrap().unwrap()) ); - assert!(logger.bridge_owned()); assert!( result .kwargs() @@ -220,17 +174,8 @@ result = Setup() .unwrap() .extract::>() .unwrap(), - ["logger", "bridge_owned", "kwargs"] + ["logger", "kwargs"] ); }); } - - #[test] - fn a_logger_extracted_from_a_bare_object_is_caller_owned() { - Python::initialize(); - Python::attach(|py| { - let logger: PythonLogger = py.None().into_bound(py).extract().unwrap(); - assert!(!logger.bridge_owned()); - }); - } } diff --git a/litellm-rust/crates/callbacks-legacy/src/preparation.rs b/litellm-rust/crates/callbacks-legacy/src/preparation.rs index 981b1702f2e..fa1ff9acd4d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/preparation.rs +++ b/litellm-rust/crates/callbacks-legacy/src/preparation.rs @@ -3,6 +3,8 @@ use pyo3::{ types::{PyDict, PyList}, }; +use crate::legacy_python::Wrapper; + struct CredentialEntry<'py>(Bound<'py, PyAny>); impl<'py> CredentialEntry<'py> { @@ -22,18 +24,19 @@ pub fn prepare<'py>( ) -> PyResult> { let arguments = kwargs.copy()?; arguments.set_item("litellm_logging_obj", logger.object(py))?; - let litellm = py.import("litellm")?; - inherit_credentials(py, &litellm, &arguments)?; - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("check_limits")? - .call1((&arguments,))?; + inherit_credentials(py, &arguments, || { + Ok(Wrapper::CredentialList + .call(py, ())? + .cast_into::()?) + })?; + Wrapper::CheckLimits.call(py, (&arguments,))?; Ok(arguments) } -fn inherit_credentials( - py: Python<'_>, - litellm: &Bound<'_, PyModule>, - arguments: &Bound<'_, PyDict>, +fn inherit_credentials<'py>( + py: Python<'py>, + arguments: &Bound<'py, PyDict>, + credential_list: impl FnOnce() -> PyResult>, ) -> PyResult<()> { let Some(requested) = arguments .get_item("litellm_credential_name")? @@ -45,16 +48,13 @@ fn inherit_credentials( return Ok(()); } let requested: String = requested.extract()?; - let credentials = litellm.getattr("credential_list")?.cast_into::()?; + let credentials = credential_list()?; let names = credentials .iter() .map(|credential| CredentialEntry(credential).name()) .collect::>>()?; let Some(index) = names.iter().position(|name| *name == requested) else { - py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1( - "warning", - ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()), - )?; + Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?; return Ok(()); }; let selected = CredentialEntry(credentials.get_item(index)?); @@ -80,19 +80,19 @@ mod tests { } fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> { - let litellm = PyModule::new(py, "credential_host")?; - litellm.setattr( - "credential_list", - locals.get_item("credentials").unwrap().unwrap(), - )?; inherit_credentials( py, - &litellm, &locals .get_item("arguments") .unwrap() .unwrap() .cast_into::()?, + || { + Ok(locals + .get_item("credentials")? + .unwrap() + .cast_into::()?) + }, ) } @@ -304,11 +304,11 @@ arguments = {'litellm_credential_name': 'ocr-test'} fn falsy_credential_names_return_before_loading_credentials() { Python::initialize(); Python::attach(|py| { - let litellm = PyModule::new(py, "credential_host").unwrap(); for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] { let arguments = PyDict::new(py); arguments.set_item("litellm_credential_name", name).unwrap(); - inherit_credentials(py, &litellm, &arguments).unwrap(); + inherit_credentials(py, &arguments, || panic!("credentials must not be loaded")) + .unwrap(); } }); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs index 3daea8840d8..289ea1b2e7f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs @@ -16,7 +16,7 @@ fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { py, PendingLogging { pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind(), true), + logger: PythonLogger::new(local(&locals, "logger").unbind()), response: Some(local(&locals, "response").unbind()), start: py.None(), end: Some(py.None()), @@ -79,22 +79,6 @@ assert logger.calls == [], logger.calls }); } -#[test] -fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c"logger.needed = {'async_success': False}"); - run( - py, - &locals, - c" -pending.release(True) -assert logger.calls == [('success_bookkeeping', True)], logger.calls -", - ); - }); -} - #[rstest] #[case::ordinary_error(c"RuntimeError('queue full')", false)] #[case::cancellation(c"asyncio.CancelledError()", true)] diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 3ceda4441a7..ea3510de17e 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -24,7 +24,7 @@ fn begin<'py>( py: Python<'py>, locals: &Bound<'py, PyDict>, asynchronous: bool, -) -> (LegacyLogging, AdapterStep) { +) -> (LegacyLogging, LifecycleStep) { let mut logging = legacy_call(py, locals, asynchronous); let kwargs = local(locals, "kwargs") .cast_into::() @@ -34,15 +34,15 @@ fn begin<'py>( (logging, step) } -fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> { - let AdapterStep::Arguments(arguments) = step else { +fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { panic!("expected the prepared arguments"); }; arguments.into_bound(py) } -fn awaits_deployment_hook(step: &AdapterStep) -> bool { - matches!(step, AdapterStep::Await(_)) +fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) } #[rstest] @@ -121,7 +121,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs} let step = logging .resume(py, Ok(local(&locals, "replacement").unbind())) .unwrap(); - let AdapterStep::Response(returned) = step else { + let LifecycleStep::Response(returned) = step else { panic!("expected the finalized response"); }; assert!(returned.bind(py).is(local(&locals, "replacement"))); @@ -195,7 +195,7 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle }; assert!(matches!( logging.resume(py, hook_result).unwrap(), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); run( py, @@ -237,7 +237,7 @@ kwargs = {'logger': logger} .unwrap() .unbind(); let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), + LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), step => Ok(step), }); let error = result.err().unwrap(); diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 480bedf8548..68c0a2b1e15 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,7 +1,8 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{AdapterStep, CallbackAdapter}; +use litellm_auth::SecretValue; +use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleStep, PythonLifecycle}; use pyo3::prelude::*; use rstest::rstest; use serde_json::{Value, json}; @@ -23,20 +24,12 @@ class PayloadLogger(StubLogger): def pre_call(self, input, api_key, additional_args): self.record('pre_call', None) self.pre = additional_args + self.pre_api_key = api_key on_pre_call(additional_args) - def _pre_call(self, input, api_key, additional_args): - self.record('_pre_call', None) - - def record_api_call_start_time(self): - self.record('record_api_call_start_time', None) - - def post_call(self, original_response, additional_args): + def post_call(self, original_response, api_key, additional_args): self.record('post_call', None) - self.post = (original_response, additional_args) - - def record_post_call(self, response, *rest): - self.record('record_post_call', response) + self.post = (original_response, api_key, additional_args) request = Request() kwargs = {} @@ -52,16 +45,16 @@ fn document(source: &str) -> Value { json!({"type": "document_url", "document_url": source}) } -fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest { - before_send_with_secrets(script, caller, body, &[]) +fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) } -/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the -/// Python objects `script` binds, then delivers the provider's raw response the way the +/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with +/// the Python objects `script` binds, then delivers the provider's raw response the way the /// driver does and runs the script's `check()`. fn before_send_with_secrets( script: &CStr, - caller: Value, + optional_params: Value, body: Value, secret_fields: &[&str], ) -> WireRequest { @@ -70,15 +63,15 @@ fn before_send_with_secrets( let locals = namespace(py, PAYLOAD_LOGGER); run(py, &locals, script); let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), ..legacy_call(py, &locals, false) }; let context = RequestContext { model: "model".into(), custom_llm_provider: "provider".into(), - optional_params: caller.clone(), - passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body), + optional_params, secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), }; let wire = WireRequest { url: "https://provider.invalid/ocr".into(), @@ -93,10 +86,10 @@ fn before_send_with_secrets( }; assert!(matches!( logging.emit(py, &raw, None).unwrap(), - AdapterStep::Done + LifecycleStep::Done )); run(py, &locals, c"check()"); - let AdapterStep::Wire(wire) = step else { + let LifecycleStep::Wire(wire) = step else { panic!("before_send did not hand back the wire request"); }; *wire @@ -129,11 +122,7 @@ def check(): ")] fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send( - script, - json!({"document": document(DOCUMENT), "pages": [0]}), - body.clone(), - ); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); } @@ -149,7 +138,6 @@ def check(): assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' ", json!({"document": document(DOCUMENT)}), - json!({"document": document(DOCUMENT)}), ); assert_eq!(wire.body["document"], document(EDITED)); } @@ -168,7 +156,6 @@ def check(): assert observed == [False], observed assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} ", - json!({"document": document("https://example.invalid/scan.pdf")}), json!({"document": document(DOCUMENT)}), ); assert_eq!( @@ -177,6 +164,23 @@ def check(): ); } +#[test] +fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); +} + #[rstest] #[case::body( c" @@ -192,7 +196,7 @@ def on_pre_call(args): )] fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({}), body.clone()); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); } @@ -205,7 +209,6 @@ def on_pre_call(args): args['headers']['x-callback'] = 'edited' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -288,7 +291,7 @@ def on_pre_call(args): )] fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({"document": document(DOCUMENT)}), body); + let wire = before_send(script, body); assert_eq!(wire.body, expected); } @@ -302,7 +305,6 @@ def on_pre_call(args): retained['x-retained'] = 'sent' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -314,52 +316,33 @@ def on_pre_call(args): } #[test] -fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() { +fn post_call_receives_the_raw_response_the_route_key_and_the_body_pre_call_saw() { before_send( c" def check(): - original_response, additional_args = logger.post + original_response, api_key, additional_args = logger.post assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == {'complete_input_dict': logger.pre['complete_input_dict']}, additional_args assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] - assert additional_args['headers'] is logger.pre['headers'] ", - json!({}), json!({"document": document(DOCUMENT)}), ); } -#[rstest] -#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])] -#[case::no_input_callback( - c"{'input': False}", - &["_pre_call", "record_api_call_start_time", "record_post_call"] -)] -#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])] -fn payload_callbacks_run_only_for_the_phases_someone_listens_to( - #[case] needed: &CStr, - #[case] expected_calls: &[&str], -) { - let script = std::ffi::CString::new(format!( - " -logger.needed = {needed} +#[test] +fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" def on_pre_call(args): args['complete_input_dict']['include_image_base64'] = True def check(): - assert logger.names() == {expected_calls:?}, logger.calls + assert logger.names() == ['pre_call', 'post_call'], logger.calls ", - needed = needed.to_str().unwrap(), - expected_calls = expected_calls, - )) - .unwrap(); - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(&script, json!({}), body.clone()); - let edited = json!({"document": document(DOCUMENT), "include_image_base64": true}); + json!({"document": document(DOCUMENT)}), + ); assert_eq!( wire.body, - if expected_calls.contains(&"pre_call") { - edited - } else { - body - } + json!({"document": document(DOCUMENT), "include_image_base64": true}) ); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 1663e11963e..444655ea77b 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -5,64 +5,89 @@ use pyo3::types::{PyDict, PyTuple}; use crate::{LegacyLogging, LegacySurface, PublicCall}; -/// Stand-ins for every litellm function the legacy contract calls. Tests share one -/// interpreter and run concurrently, so each stub is installed idempotently and forwards to -/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// The parameters of every `legacy_callbacks` function, as the real module declares them. +/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python +/// signatures, and [`namespace`] binds every fake call against it. +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + +/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests +/// share one interpreter and run concurrently, so each fake is installed idempotently and +/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// Every fake is bound against the contract first, so a call the real module would reject +/// fails here too. const STUBS: &CStr = c" import contextvars +import inspect +import json import sys +import traceback import types -for name in ( - 'litellm', - 'litellm.utils', - 'litellm.types', - 'litellm.types.utils', - 'litellm._internal_context', - 'litellm.litellm_core_utils', - 'litellm.litellm_core_utils.logging_worker', - 'litellm.litellm_core_utils.litellm_logging', - 'litellm.rust_bridge', - 'litellm.rust_bridge.legacy_callbacks', -): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): sys.modules.setdefault(name, types.ModuleType(name)) legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - bridge_owned=True, -) -legacy.deployment_callbacks_needed = lambda: True -legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments) -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record( - 'success_bookkeeping', asynchronous -) -legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record( - 'failure_bookkeeping', asynchronous -) -legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response) +CONTRACT = json.loads(python_contract) -utils = sys.modules['litellm.utils'] -utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook( - 'pre', kwargs, call_type -) -utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[ - 'logger' -].hook('success', response, call_type) -utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[ - 'logger' -].hook('failure', error, call_type) -utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None) -internal = sys.modules['litellm._internal_context'] -if not hasattr(internal, 'is_internal_call'): - internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False) +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) -sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type( - 'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}} -) + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) unraisable = sys.modules.setdefault( @@ -77,20 +102,6 @@ def unraisable_from(owner): return [error for source, error in unraisable.events if source is owner] -class Worker: - def ensure_initialized_and_enqueue(self, coroutine): - return coroutine.enqueue() - - -class Executor: - def submit(self, run, handler, *args): - handler.__self__.record('submit', args) - - -sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker() -sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor() - - class StubCoroutine: def __init__(self, logger): self.logger = logger @@ -106,7 +117,6 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] - self.needed = {} self.hooks = {} self.on_enqueue = lambda coroutine: None @@ -147,6 +157,7 @@ logger = StubLogger() /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); py.run(script, Some(&locals), Some(&locals)).unwrap(); locals diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 9b9d29108f6..3094d7b88d2 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -19,12 +19,16 @@ const TIMING: Timing = Timing { fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), ..legacy_call(py, locals, asynchronous) } } -fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, +) -> LifecycleStep { let response = local(locals, "response").unbind(); logging .emit( @@ -35,7 +39,7 @@ fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLoggi .unwrap() } -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { let failure = PyErr::from_value(local(locals, "failure")); logging .emit( @@ -51,20 +55,14 @@ fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) #[rstest] #[case::sync_listened(false, c"", &["submit"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])] #[case::async_listened( true, c"", &["async_success_handler", "enqueued", "sync_success_for_async_call"] )] -#[case::async_unlistened( - true, - c"logger.needed = {'async_success': False, 'sync_success_async': False}", - &["success_bookkeeping"] -)] #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_only_the_callbacks_that_listen( +fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -76,7 +74,7 @@ fn success_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); let names: Vec = local(&locals, "logger") .call_method0("names") @@ -109,7 +107,10 @@ fn internal_calls_skip_failure_callbacks_only_when_asynchronous( internal: true, ..logged(py, &locals, asynchronous) }; - assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done)); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -157,7 +158,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); assert!( logging @@ -173,14 +174,8 @@ logger = FailingLogger() #[rstest] #[case::sync_listened(false, c"", &["failure_handler"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])] #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -#[case::async_unlistened( - true, - c"logger.needed = {'sync_failure': False, 'async_failure': False}", - &["failure_bookkeeping", "failure_bookkeeping"] -)] -fn failure_reaches_only_the_callbacks_that_listen( +fn failure_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -192,7 +187,10 @@ fn failure_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); let step = fail(py, &locals, &mut logging); let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -227,7 +225,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( fail(py, &locals, &mut logging), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); assert!( logging @@ -265,7 +263,7 @@ fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( }; let expected = result.as_ref().err().map(|error| error.value(py).clone()); match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, AdapterStep::Done)), + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), Err(propagated) => { assert!(!done); assert!(propagated.value(py).is(expected.unwrap())); diff --git a/litellm-rust/crates/callbacks/Cargo.toml b/litellm-rust/crates/callbacks/Cargo.toml index 4b966271478..a68ebc26a8d 100644 --- a/litellm-rust/crates/callbacks/Cargo.toml +++ b/litellm-rust/crates/callbacks/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-auth.workspace = true serde_json.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index e6f88fd9709..3bf12e553b9 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -1,6 +1,6 @@ use std::time::{SystemTime, UNIX_EPOCH}; -use serde_json::{Map, Value}; +use serde_json::Value; /// Seconds since the Unix epoch, on one clock for every host. pub fn epoch_seconds() -> f64 { @@ -33,34 +33,10 @@ pub struct RequestContext { pub custom_llm_provider: String, /// The route's parameters before the provider transformation. pub optional_params: Value, - pub passthrough_fields: Passthrough, /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, -} - -/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to -/// build one is to compare the two, so a route cannot name a key it rewrote. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Passthrough(Vec); - -impl Passthrough { - pub fn unchanged(caller: &Map, body: &Value) -> Self { - Self( - caller - .iter() - .filter(|(name, value)| body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.clone()) - .collect(), - ) - } - - pub fn iter(&self) -> impl Iterator { - self.0.iter().map(String::as_str) - } - - pub fn contains(&self, name: &str) -> bool { - self.0.iter().any(|field| field == name) - } + /// The credential the route resolved for the provider call. + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -78,6 +54,9 @@ pub enum FailureOrigin { #[derive(Clone, Debug, PartialEq)] pub enum CallEvent { + Started { + start_time: f64, + }, ResponseReceived { raw: RawResponse, }, @@ -89,47 +68,3 @@ pub enum CallEvent { origin: FailureOrigin, }, } - -#[cfg(test)] -mod tests { - use rstest::rstest; - use serde_json::json; - - use super::*; - - #[rstest] - #[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])] - #[case::unchanged_nested_object( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}), - &["document"] - )] - #[case::rewritten_value( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}), - &[] - )] - #[case::dropped_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - &[] - )] - #[case::added_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}), - &[] - )] - #[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])] - #[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])] - #[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])] - fn passthrough_is_exactly_the_callers_unchanged_keys( - #[case] caller: Value, - #[case] body: Value, - #[case] expected: &[&str], - ) { - let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body); - assert_eq!(passthrough.iter().collect::>(), expected); - } -} diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 57bf134f345..5accaa2e25e 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -11,6 +11,7 @@ where H: Host, { let start_time = epoch_seconds(); + let _ = host.emit(&CallEvent::Started { start_time }).await; let mut result = None; let outcome = loop { let step = match machine.resume(result.take()).await { @@ -102,6 +103,7 @@ mod tests { async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { self.seen.lock().unwrap().push(match event { + CallEvent::Started { .. } => "started".into(), CallEvent::Succeeded { .. } => "succeeded".into(), CallEvent::Failed { .. } => "failed".into(), other => format!("{other:?}"), @@ -124,7 +126,7 @@ mod tests { assert_eq!(outcome, Ok(())); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "succeeded"] + ["started", "route:project", "route:send", "succeeded"] ); } @@ -133,7 +135,7 @@ mod tests { let host = Recording::default(); let outcome = run(scripted(&[], Err("boom")), &host).await; assert_eq!(outcome, Err("boom")); - assert_eq!(*host.seen.lock().unwrap(), ["failed"]); + assert_eq!(*host.seen.lock().unwrap(), ["started", "failed"]); let host = Recording { fail: Some("send"), @@ -143,7 +145,35 @@ mod tests { assert_eq!(outcome, Err("host failed")); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "failed"] + ["started", "route:project", "route:send", "failed"] ); } + + struct StartTimes(Mutex>); + + impl Host for StartTimes { + async fn route(&self, _: &'static str) -> Result<(), &'static str> { + Ok(()) + } + + async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { + if let CallEvent::Started { start_time } + | CallEvent::Succeeded { + timing: Timing { start_time, .. }, + } = event + { + self.0.lock().unwrap().push(*start_time); + } + Err("observer failed") + } + } + + #[tokio::test] + async fn started_opens_the_call_at_the_terminal_start_time_and_cannot_fail_it() { + let host = StartTimes(Mutex::default()); + assert_eq!(run(scripted(&["project"], Ok(())), &host).await, Ok(())); + let times = host.0.lock().unwrap(); + assert_eq!(times.len(), 2); + assert_eq!(times[0], times[1]); + } } diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 33cb8a8d32a..a6af190eb91 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,5 +1,6 @@ use futures_util::future::BoxFuture; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; +use litellm_auth::SecretValue; +use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -36,6 +37,7 @@ pub(crate) struct OcrCallHooks { custom_llm_provider: &'static str, optional_params: Value, secret_fields: Vec, + api_key: Option, } impl OcrCallHooks { @@ -51,22 +53,19 @@ impl OcrCallHooks { .filter(|name| is_secret_param(name)) .cloned() .collect(), + api_key: request.connection.api_key.clone(), } } } impl CallHooks for OcrCallHooks { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { let context = RequestContext { model: self.model.clone(), custom_llm_provider: self.custom_llm_provider.into(), optional_params: self.optional_params.clone(), - passthrough_fields, secret_fields: self.secret_fields.clone(), + api_key: self.api_key.clone(), }; Box::pin(self.host.before_send(wire, context)) } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index e7f77acc3f8..c977f721a70 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -21,8 +21,8 @@ mod cohere_tests; #[path = "../../tests/deepseek_ocr.rs"] mod deepseek_tests; #[cfg(test)] -#[path = "../../tests/ocr/passthrough.rs"] -mod passthrough_tests; +#[path = "../../tests/ocr/document.rs"] +mod document_tests; #[cfg(test)] #[path = "../../tests/reducto_ocr.rs"] mod reducto_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 24c3f43e2b4..8ac038290b7 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,4 +1,4 @@ -use litellm_auth::{InputSource, Sourced}; +use litellm_auth::{InputSource, SecretValue, Sourced}; use litellm_llms::base_llm::ocr::transformation::{ OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env, }; @@ -22,7 +22,7 @@ pub(crate) fn prepare_request( .config .get_api_key_env_var() .and_then(credential_env) - .map(|value| Sourced::new(value, InputSource::Environment)) + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index d12b8cfee95..14b34ea4564 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -277,12 +277,18 @@ mod tests { #[test] fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Request, @@ -292,7 +298,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("dynamic-key") ); assert_eq!( @@ -318,22 +324,31 @@ mod tests { fn empty_or_missing_dynamic_credentials_preserve_explicit_values( #[case] dynamic_value: Option<&str>, ) { - let dynamic = + let dynamic_key = dynamic_value.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Environment, + ) + }); + let dynamic_base = dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: dynamic.clone(), - dynamic_api_base: dynamic, + dynamic_api_key: dynamic_key, + dynamic_api_base: dynamic_base, }); assert_eq!( connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("explicit-key") ); assert_eq!( @@ -356,11 +371,18 @@ mod tests { ) { let connection = OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params( OcrCredentialInputs { - api_key: explicit_key - .map(|value| Sourced::new(value.into(), InputSource::Deployment)), + api_key: explicit_key.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Deployment, + ) + }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Deployment, @@ -371,7 +393,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), explicit_key.map(|_| "dynamic-key") ); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 75202ed52a5..6316088dec8 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, @@ -56,7 +56,7 @@ pub struct OcrFileContent { /// credentials, and per-field provenance in `input_sources`. #[derive(Clone, Debug, Default)] pub struct OcrConnectionInputs { - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub extra_headers: Map, pub timeout: Option, @@ -237,6 +237,16 @@ mod tests { .unwrap() } + #[test] + fn connection_inputs_debug_hides_the_api_key() { + let inputs = OcrConnectionInputs { + api_key: Some(SecretValue::new("caller-api-key")), + ..OcrConnectionInputs::default() + }; + + assert!(!format!("{inputs:?}").contains("caller-api-key")); + } + #[test] fn from_inputs_applies_connection_overrides_with_field_sources() { let request = LiteLLMOcrRequest::from_inputs( @@ -245,7 +255,7 @@ mod tests { None, Default::default(), OcrConnectionInputs { - api_key: Some(" key ".into()), + api_key: Some(SecretValue::new(" key ")), api_base: Some("".into()), extra_headers: json!({"x-a": "1"}).as_object().unwrap().clone(), timeout: Some(Duration::from_secs(7)), @@ -259,7 +269,7 @@ mod tests { .unwrap(); let api_key = request.credentials.api_key.as_ref().unwrap(); - assert_eq!(api_key.clone().into_value(), "key"); + assert_eq!(api_key.value().expose(), "key"); assert_eq!(api_key.source(), InputSource::Request); assert!(request.credentials.api_base.is_none()); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 29345e38885..b9c60f57e3c 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::InputSource; +use litellm_auth::{InputSource, SecretValue}; use litellm_llms::base_llm::ocr::{ error::Error, transformation::{OcrDocument, decode_request_value}, @@ -44,7 +44,7 @@ pub fn consumed_optional_param_names( pub struct OcrWireRequest { pub model: String, pub document: D, - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, pub extra_headers: Option>, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 01a4e5efb3b..1fc4d6c2b9e 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -69,7 +69,7 @@ async fn rejects_invalid_pages_features_and_format( let result = decode_request(OcrWireRequest { model: "azure_ai/doc-intelligence/prebuilt-read".into(), document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: Some(base), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 1f591d74d5d..ca2e14a7f0d 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -81,7 +81,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { let request = OcrWireRequest { model: "mistral/model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: None, extra_headers: None, @@ -97,7 +97,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { decode_request(OcrWireRequest { model: "model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: Some("unknown".into()), extra_headers: None, @@ -194,6 +194,7 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { + CallEvent::Started { .. } => "started", CallEvent::ResponseReceived { .. } => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", @@ -235,7 +236,7 @@ async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { } #[tokio::test] -async fn before_send_context_names_passthrough_fields_and_secrets() { +async fn before_send_context_names_the_route_and_its_secrets() { let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; let observed = Arc::new(Mutex::new(None)); let captured = observed.clone(); @@ -254,8 +255,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { assert_eq!(context.custom_llm_provider, "mistral"); assert_eq!(context.model, "model"); assert_eq!(wire.body["pages"], json!([0])); - assert!(context.passthrough_fields.contains("pages")); - assert!(context.passthrough_fields.contains("document")); assert!(context.secret_fields.is_empty()); assert_eq!(context.optional_params["req_format"], "native"); @@ -279,7 +278,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { perform_ocr_with(host).await.unwrap(); server.await.unwrap(); let context = observed.lock().unwrap().take().unwrap(); - assert!(!context.passthrough_fields.contains("document")); assert_eq!(context.secret_fields, ["client_secret"]); } @@ -296,7 +294,7 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { server.await.unwrap(); assert_eq!( *events.lock().unwrap(), - ["before_send", "response", "success"] + ["started", "before_send", "response", "success"] ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -311,7 +309,10 @@ async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { ); let error = perform_ocr_with(host).await.unwrap_err(); assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); } #[tokio::test] @@ -330,7 +331,10 @@ async fn upstream_failure_emits_one_terminal_failure() { ); assert!(perform_ocr_with(host).await.is_err()); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs new file mode 100644 index 00000000000..5e10ce3ad2a --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/document.rs @@ -0,0 +1,152 @@ +use litellm_callbacks::event::WireRequest; +use litellm_llms::base_llm::ocr::error::Error; +use rstest::rstest; +use serde_json::{Value, json}; + +use super::test_support::{ + MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, + wire_request_with_document, +}; +use crate::ocr::route::LocalOcrHost; + +#[derive(Clone, Copy, Debug)] +enum Route { + Mistral, + AzureAi, + VertexMistral, + AzureCohereParse, + Cohere, +} + +impl Route { + fn model(self) -> &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } +} + +/// What the host does to the wire request in `before_send`. +#[derive(Clone, Copy, Debug)] +enum Host { + Detached, + ReplacesDocument, +} + +const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + +impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } +} + +struct Sent { + result: Result<(), Error>, + provider_body: Option, +} + +async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } +} + +fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) +} + +#[rstest] +#[case::azure_ai(Route::AzureAi)] +#[case::vertex_mistral(Route::VertexMistral)] +#[case::azure_cohere_parse(Route::AzureCohereParse)] +#[tokio::test] +async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); +} + +#[rstest] +#[tokio::test] +async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, +) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/passthrough.rs b/litellm-rust/crates/core/tests/ocr/passthrough.rs deleted file mode 100644 index 0273b48664d..00000000000 --- a/litellm-rust/crates/core/tests/ocr/passthrough.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::{ - collections::BTreeSet, - sync::{Arc, Mutex}, -}; - -use litellm_callbacks::event::{RequestContext, WireRequest}; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use rstest_reuse::{self, apply, template}; -use serde_json::{Map, Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -#[derive(Clone, Copy, Debug)] -enum Source { - Inline, - Remote, - RemoteWithExtraField, -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - /// What `litellm-callbacks-legacy` does before `pre_call`: every passthrough body key - /// is replaced by the caller's own value. - Realiasing, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send( - self, - caller: &Map, - wire: WireRequest, - context: &RequestContext, - ) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::Realiasing => { - let aliased = context - .passthrough_fields - .contains(&name) - .then(|| caller.get(&name).cloned()) - .flatten() - .unwrap_or(value); - (name, aliased) - } - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - caller: Map, - result: Result<(), Error>, - before_send: Option<(WireRequest, RequestContext)>, - provider_body: Option, -} - -fn caller_document(route: Route, source: Source, document_base: &str) -> Value { - let document_type = route.document_type(); - let remote = format!("{document_base}/scan.png"); - match source { - Source::Inline => { - json!({"type": document_type, document_type: "data:image/png;base64,YWJj"}) - } - Source::Remote => json!({"type": document_type, document_type: remote}), - Source::RemoteWithExtraField => { - json!({"type": document_type, document_type: remote, "document_name": "scan.png"}) - } - } -} - -async fn send(route: Route, source: Source, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document = caller_document(route, source, document_base); - let caller: Map = route - .options() - .as_object() - .unwrap() - .clone() - .into_iter() - .chain([("document".to_string(), document.clone())]) - .collect(); - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host_caller = caller.clone(); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(host.before_send(&host_caller, wire, context)) - }); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - let before_send = observed.lock().unwrap().take(); - Sent { - caller, - result, - before_send, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[template] -#[rstest] -fn every_route_and_source( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, - #[values(Source::Inline, Source::Remote, Source::RemoteWithExtraField)] source: Source, -) { -} - -#[template] -#[rstest] -fn every_route( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { -} - -#[template] -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -fn inlining_routes(#[case] route: Route) {} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn passthrough_fields_are_exactly_the_caller_values_sent_unchanged( - route: Route, - source: Source, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, source, Host::Detached, &document_base).await; - sent.result.unwrap(); - let (wire, context) = sent.before_send.unwrap(); - let passthrough: BTreeSet<&str> = context.passthrough_fields.iter().collect(); - let unchanged: BTreeSet<&str> = sent - .caller - .iter() - .filter(|(name, value)| wire.body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.as_str()) - .collect(); - assert_eq!( - passthrough, - unchanged, - "body: {:#}\ncaller: {:#}", - wire.body, - Value::Object(sent.caller.clone()) - ); -} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn realiasing_leaves_the_provider_request_unchanged(route: Route, source: Source) { - let (document_base, _documents) = document_server().await; - let detached = send(route, source, Host::Detached, &document_base).await; - let realiased = send(route, source, Host::Realiasing, &document_base).await; - detached.result.unwrap(); - realiased.result.unwrap(); - assert_eq!(realiased.provider_body, detached.provider_body); -} - -#[apply(inlining_routes)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document( - route: Route, - #[values(Host::Detached, Host::Realiasing)] host: Host, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Source::Remote, host, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[apply(every_route)] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider(route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send( - route, - Source::Remote, - Host::ReplacesDocument, - &document_base, - ) - .await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index f3adf27cfa6..44313d5f552 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_callbacks::event::WireRequest; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::{CallHooks, OcrClient}, @@ -23,11 +23,7 @@ use crate::ocr::{ pub(crate) struct NoHooks; impl CallHooks for NoHooks { - fn before_send( - &self, - wire: WireRequest, - _passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { Box::pin(async move { Ok(wire) }) } @@ -70,7 +66,7 @@ pub(crate) fn wire_request_with_document( decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some("test-key".into()), + api_key: Some(litellm_auth::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index a3fdd2340b3..903ccb06c77 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,9 +1,10 @@ - Target invariants; implementation and runtime validation may lag these rules -- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `CallbackAdapter`/`RouteHost` traits +- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - - A failure that surfaces inside the call, including a host op the call asked for, is mapped through the route's `map_failure`; a failure in `begin` or `after_success` is raised as is + - A native failure, including one a host op returns as `HostOpError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is + - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index f1bc3142a25..4aa7a2163ca 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -11,7 +11,7 @@ pub fn missing_state() -> PyErr { /// What an adapter step produced: either the value the driver asked for, or a Python /// awaitable the driver hands back to the caller's task before asking again. -pub enum AdapterStep { +pub enum LifecycleStep { Await(Py), Arguments(Py), Wire(Box), @@ -28,55 +28,74 @@ pub enum PublicValue<'a> { /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in /// order: `begin` before the machine starts, `before_send` and `emit` while it runs, /// `after_success` and one terminal `emit` after it completes. Whenever a step returns -/// [`AdapterStep::Await`], the driver awaits it in the caller's task and continues the +/// [`LifecycleStep::Await`], the driver awaits it in the caller's task and continues the /// same step through `resume`. /// /// A step that fails with an ordinary exception fails the call with that exception, /// except on a terminal event, where the adapter is expected to report and swallow its /// own errors. An exception that is not a `PyException`, such as a cancellation, ends /// the call without further dispatch. -pub trait CallbackAdapter: Send + Sync { +pub trait PythonLifecycle: Send + Sync { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult; + ) -> PyResult; fn before_send( &mut self, py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult; + ) -> PyResult; fn after_success( &mut self, py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult; + ) -> PyResult; fn emit( &mut self, py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult; + ) -> PyResult; - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } -/// The Python side of one route: answers the route's own operations, builds the public -/// response and maps failures to public exceptions. -pub trait RouteHost: Send + Sync { - type Route: Route; +/// Why a route operation the host answered did not produce a result: the route's own code +/// rejected it, which the route classifies like any other native failure, or Python code +/// raised, which reaches the caller as it was raised. +#[derive(Debug)] +pub enum HostOpError { + Native(E), + Python(PyErr), +} - /// `arguments` is the keyword view the callback adapter's `begin` produced, not the +impl From for HostOpError { + fn from(error: PyErr) -> Self { + Self::Python(error) + } +} + +/// The Python side of one route: answers the route's own operations, builds the public +/// response and classifies native failures into public exceptions. +pub trait RouteHost: Send + Sync { + type Route: Route; + + /// The public exception a native failure maps to, kept as a value until the driver + /// raises it. + type Failure: Into; + + /// `arguments` is the keyword view the lifecycle's `begin` produced, not the /// caller's own dict. A route host that projects from it inherits whatever that /// adapter rewrote. fn invoke( @@ -84,7 +103,7 @@ pub trait RouteHost: Send + Sync { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: ::Op, - ) -> PyResult<::OpResult>; + ) -> Result<::OpResult, HostOpError<::Error>>; fn complete( &mut self, @@ -92,12 +111,14 @@ pub trait RouteHost: Send + Sync { response: ::Response, ) -> PyResult>; - fn native_error(error: ::Error) -> PyErr; + fn classify( + &self, + py: Python<'_>, + error: ::Error, + ) -> PyResult; fn host_error(error: &PyErr) -> ::Error; - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult; - fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; diff --git a/litellm-rust/crates/host-python/src/argument.rs b/litellm-rust/crates/host-python/src/argument.rs new file mode 100644 index 00000000000..34e07cdfbd5 --- /dev/null +++ b/litellm-rust/crates/host-python/src/argument.rs @@ -0,0 +1,51 @@ +use pyo3::{prelude::*, types::PyDict}; + +/// The caller's own object for a public argument: the keyword if given, even an explicit +/// `None`, else the bound request's attribute. Every reader of a public Python call uses +/// this rule, so the callbacks and the provider see one object per argument. +pub fn lookup<'py>( + kwargs: &Bound<'py, PyDict>, + request: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + if let Some(value) = kwargs.get_item(name)? { + return Ok(Some(value)); + } + request.getattr_opt(name) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { + crate::initialize_python(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +key = object() +document = {'type': 'document_url'} +class Request: + api_key = 'from-request' + api_base = 'from-request' + document = document +request = Request() +kwargs = {'api_key': key, 'api_base': None} +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let item = |name: &str| locals.get_item(name).unwrap().unwrap(); + let kwargs = item("kwargs").cast_into::().unwrap(); + let request = item("request"); + let find = |name: &str| lookup(&kwargs, &request, name).unwrap(); + assert!(find("api_key").unwrap().is(item("key"))); + assert!(find("api_base").unwrap().is_none()); + assert!(find("document").unwrap().is(item("document"))); + assert!(find("model").is_none()); + }); + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 8bda13b44d0..7895d803cb5 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -12,7 +12,9 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use tokio::sync::Mutex; -use crate::adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +use crate::adapter::{ + HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, +}; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -43,6 +45,7 @@ enum Stage { #[derive(Clone, Copy)] enum Expect { + Started, Arguments, Wire, Emitted, @@ -66,7 +69,7 @@ where M: Machine> + 'static, { route: H, - adapter: Box, + adapter: Box, machine: Option>>>, arguments: Option>, started_at: f64, @@ -84,7 +87,7 @@ pub fn run_call( py: Python<'_>, machine: M, route: H, - adapter: Box, + adapter: Box, arguments: Py, asynchronous: bool, ) -> PyResult> @@ -146,9 +149,11 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let arguments = self.arguments.take().ok_or_else(missing_state)?; - match self.adapter.begin(py, arguments, self.started_at) { - Ok(step) => self.on_adapter(py, step, Expect::Arguments), + let started = CallEvent::Started { + start_time: self.started_at, + }; + match self.adapter.emit(py, &started, None) { + Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } } @@ -170,27 +175,28 @@ where fn on_adapter( &mut self, py: Python<'_>, - step: AdapterStep, + step: LifecycleStep, expect: Expect, ) -> PyResult { match (expect, step) { - (_, AdapterStep::Await(awaitable)) => { + (_, LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(expect)); Ok(ExecutionStep::Await(awaitable)) } - (Expect::Arguments, AdapterStep::Arguments(arguments)) => { + (Expect::Started, LifecycleStep::Done) => self.begin(py), + (Expect::Arguments, LifecycleStep::Arguments(arguments)) => { self.arguments = Some(arguments); self.stage = Stage::Call; self.resume_machine(py, None) } - (Expect::Wire, AdapterStep::Wire(wire)) => { + (Expect::Wire, LifecycleStep::Wire(wire)) => { self.resume_machine(py, Some(Ok(HostResult::BeforeSend(wire)))) } - (Expect::Emitted, AdapterStep::Done) => { + (Expect::Emitted, LifecycleStep::Done) => { self.resume_machine(py, Some(Ok(HostResult::Emitted))) } - (Expect::Response, AdapterStep::Response(response)) => self.succeeded(py, response), - (Expect::Terminal, AdapterStep::Done) => match &self.stage { + (Expect::Response, LifecycleStep::Response(response)) => self.succeeded(py, response), + (Expect::Terminal, LifecycleStep::Done) => match &self.stage { Stage::Succeeded(response) => Ok(ExecutionStep::Return(response.clone_ref(py))), Stage::Failed(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())), _ => Err(missing_state()), @@ -199,6 +205,14 @@ where } } + fn begin(&mut self, py: Python<'_>) -> PyResult { + let arguments = self.arguments.take().ok_or_else(missing_state)?; + match self.adapter.begin(py, arguments, self.started_at) { + Ok(step) => self.on_adapter(py, step, Expect::Arguments), + Err(error) => self.adapter_failed(py, error), + } + } + fn adapter_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { match self.stage { Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), @@ -248,14 +262,20 @@ where let answer = match op { HostOp::Route(op) => { let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; - self.route - .invoke(py, arguments.bind(py), op) - .map(HostResult::Route) + match self.route.invoke(py, arguments.bind(py), op) { + Ok(result) => Ok(HostResult::Route(result)), + Err(HostOpError::Native(error)) => { + return self + .resume_core(py, Some(Err(HostFailure::Error(error)))) + .map(Next::Continue); + } + Err(HostOpError::Python(error)) => Err(error), + } } HostOp::BeforeSend { wire, context } => { match self.adapter.before_send(py, wire, &context) { - Ok(AdapterStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Wire)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -264,8 +284,8 @@ where } } HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { - Ok(AdapterStep::Done) => Ok(HostResult::Emitted), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -360,11 +380,29 @@ where self.ended_at.get_or_insert_with(epoch_seconds); let error = match self.interrupted.take() { Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), - None => H::native_error(error), + None => self.classified(py, error), }; self.failure(py, error, FailureOrigin::Call) } + /// The route's public exception for a native failure. When classification itself + /// fails, that failure is raised with the native error's text as its `__context__`. + fn classified(&self, py: Python<'_>, error: ErrorOf) -> PyErr { + let native = error.to_string(); + let classifier_error = match self.route.classify(py, error) { + Ok(failure) => return failure.into(), + Err(classifier_error) => classifier_error, + }; + let attached = classifier_error.value(py).setattr( + "__context__", + PyRuntimeError::new_err(native).into_value(py), + ); + match attached { + Ok(()) => classifier_error, + Err(error) => error, + } + } + fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { let event = CallEvent::Succeeded { timing: self.timing(), @@ -386,18 +424,14 @@ where if is_cancellation(py, &error) { return Err(error); } - let public = match origin { - FailureOrigin::Call => self.route.map_failure(py, &error).unwrap_or(error), - FailureOrigin::Host => error, - }; let event = CallEvent::Failed { timing: self.timing(), origin, }; let step = self .adapter - .emit(py, &event, Some(PublicValue::Error(&public)))?; - self.stage = Stage::Failed(public.into_value(py)); + .emit(py, &event, Some(PublicValue::Error(&error)))?; + self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -489,6 +523,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri #[derive(Clone, Debug, PartialEq, Eq)] struct Error(String); + impl std::fmt::Display for Error { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } + } + struct Synthetic; impl Route for Synthetic { @@ -518,8 +558,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri model: "model".into(), custom_llm_provider: "provider".into(), optional_params: serde_json::json!({}), - passthrough_fields: Default::default(), secret_fields: Vec::new(), + api_key: None, } } @@ -566,25 +606,46 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } + #[derive(Clone, Copy)] + enum OpScript { + Answer, + RaisePython, + RejectNatively, + } + struct SyntheticHost { log: Log, - fail_op: bool, + op: OpScript, + classifier_fails: bool, + } + + /// The fake route's public exception, kept as a value so a test sees what `classify` + /// produced before the driver raises it. + #[derive(Debug, PartialEq, Eq)] + struct Classified(String); + + impl From for PyErr { + fn from(classified: Classified) -> Self { + PyValueError::new_err(format!("classified: {}", classified.0)) + } } impl RouteHost for SyntheticHost { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, _: Python<'_>, arguments: &Bound<'_, PyDict>, op: &'static str, - ) -> PyResult { + ) -> Result> { self.log.push(format!("route:{op}")); - if self.fail_op { - return Err(PyValueError::new_err("op failed")); + match self.op { + OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), + OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), + OpScript::RejectNatively => Err(HostOpError::Native(Error("op rejected".into()))), } - Ok(format!("{op}:{}", arguments.len())) } fn complete(&mut self, py: Python<'_>, response: String) -> PyResult> { @@ -594,22 +655,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unbind()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.log.push(format!("classify:{error}")); + if self.classifier_fails { + return Err(pyo3::exceptions::PyTypeError::new_err("classifier failed")); + } + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - self.log.push("map_failure"); - Ok(PyValueError::new_err(format!( - "mapped: {}", - error.value(py) - ))) - } - fn close(&mut self, _: Python<'_>) { self.log.push("route.close"); } @@ -632,13 +689,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri script: AdapterScript, } - impl CallbackAdapter for SyntheticAdapter { - fn begin(&mut self, _: Python<'_>, arguments: Py, _: f64) -> PyResult { + impl PythonLifecycle for SyntheticAdapter { + fn begin( + &mut self, + _: Python<'_>, + arguments: Py, + _: f64, + ) -> PyResult { self.log.push("begin"); if matches!(self.script, AdapterScript::FailBegin) { return Err(PyValueError::new_err("begin failed")); } - Ok(AdapterStep::Arguments(arguments)) + Ok(LifecycleStep::Arguments(arguments)) } fn before_send( @@ -646,9 +708,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri _: Python<'_>, wire: Box, _: &RequestContext, - ) -> PyResult { + ) -> PyResult { self.log.push("before_send"); - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { url: "rewritten".into(), ..*wire }))) @@ -659,17 +721,17 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, response: Py, _: Timing, - ) -> PyResult { + ) -> PyResult { self.log.push("after_success"); match self.script { - AdapterScript::ReplaceResponse => Ok(AdapterStep::Response( + AdapterScript::ReplaceResponse => Ok(LifecycleStep::Response( "replaced".into_pyobject(py)?.into_any().unbind(), )), AdapterScript::FailAfterSuccess => { Err(PyValueError::new_err("after_success failed")) } AdapterScript::Plain | AdapterScript::FailBegin => { - Ok(AdapterStep::Response(response)) + Ok(LifecycleStep::Response(response)) } } } @@ -679,8 +741,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult { + ) -> PyResult { self.log.push(match (event, public) { + (CallEvent::Started { .. }, None) => "started".into(), (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { format!("succeeded:{}", value.bind(py)) @@ -690,10 +753,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } _ => "unexpected".into(), }); - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } - fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -709,15 +772,31 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn run_scripted( py: Python<'_>, machine: ScriptedMachine, - fail_op: bool, + op: OpScript, script: AdapterScript, asynchronous: bool, ) -> (PyResult>, Vec) { - let log = Log::default(); - let route = SyntheticHost { - log: Log(log.0.clone()), - fail_op, - }; + run_hosted( + py, + machine, + SyntheticHost { + log: Log::default(), + op, + classifier_fails: false, + }, + script, + asynchronous, + ) + } + + fn run_hosted( + py: Python<'_>, + machine: ScriptedMachine, + route: SyntheticHost, + script: AdapterScript, + asynchronous: bool, + ) -> (PyResult>, Vec) { + let log = Log(route.log.0.clone()); let adapter = SyntheticAdapter { log: Log(log.0.clone()), script, @@ -777,7 +856,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::Plain, asynchronous, ); @@ -785,6 +864,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "route:project", "before_send", @@ -800,28 +880,75 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + fn failing_machine() -> ScriptedMachine { + ScriptedMachine { + ops: vec![HostOp::Route("project")], + outcome: Some(Err(Error("provider exploded".into()))), + answers: Vec::new(), + } + } + #[test] - fn machine_failures_are_mapped_and_dispatched_once_as_call_failures() { + fn a_native_failure_is_classified_once_and_reported_classified() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let machine = ScriptedMachine { - ops: vec![HostOp::Route("project")], - outcome: Some(Err(Error("provider exploded".into()))), - answers: Vec::new(), - }; - let (result, log) = run_scripted(py, machine, false, AdapterScript::Plain, false); - let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: provider exploded"); + install_lifecycle_module(py); + for asynchronous in [false, true] { + let (result, log) = run_scripted( + py, + failing_machine(), + OpScript::Answer, + AdapterScript::Plain, + asynchronous, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classified: provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classified: provider exploded", + "adapter.close", + "route.close", + ] + ); + } + }); + } + + #[test] + fn a_native_rejection_from_a_host_operation_is_classified_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RejectNatively, + AdapterScript::Plain, + false, + ); + assert_eq!( + result.unwrap_err().value(py).to_string(), + "classified: op rejected" + ); assert_eq!( log, [ + "started", "begin", "route:project", - "map_failure", - "failed:Call:mapped: provider exploded", + "classify:op rejected", + "failed:Call:classified: op rejected", "adapter.close", "route.close", ] @@ -830,18 +957,72 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } #[test] - fn host_operation_failures_interrupt_the_call_and_keep_the_python_exception() { + fn a_python_exception_from_a_host_operation_is_reported_as_raised() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let (result, log) = - run_scripted(py, success_machine(), true, AdapterScript::Plain, false); + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RaisePython, + AdapterScript::Plain, + false, + ); let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: op failed"); - assert!(!log.contains(&"before_send".to_string())); - assert!(log.contains(&"failed:Call:mapped: op failed".to_string())); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "op failed"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "failed:Call:op failed", + "adapter.close", + "route.close", + ] + ); + }); + } + + #[test] + fn a_failing_classifier_surfaces_with_the_native_error_as_context() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_hosted( + py, + failing_machine(), + SyntheticHost { + log: Log::default(), + op: OpScript::Answer, + classifier_fails: true, + }, + AdapterScript::Plain, + false, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classifier failed"); + let context = error.value(py).getattr("__context__").unwrap(); + assert!(context.is_instance_of::()); + assert_eq!(context.str().unwrap().to_string(), "provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classifier failed", + "adapter.close", + "route.close", + ] + ); }); } @@ -855,7 +1036,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailBegin, false, ); @@ -864,6 +1045,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "failed:Host:begin failed", "adapter.close", @@ -885,7 +1067,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::ReplaceResponse, asynchronous, ); @@ -908,7 +1090,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailAfterSuccess, asynchronous, ); @@ -938,12 +1120,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri struct Cancelling(Log); impl RouteHost for Cancelling { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, py: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, - ) -> PyResult { + ) -> Result> { self.0.push("route"); Err(PyErr::from_value( py.import("asyncio") @@ -952,21 +1135,19 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unwrap() .call0() .unwrap(), - )) + ) + .into()) } fn complete(&mut self, _: Python<'_>, _: String) -> PyResult> { Err(missing_state()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.0.push("classify"); + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, _: Python<'_>, _: &PyErr) -> PyResult { - self.0.push("map_failure"); - Err(missing_state()) - } fn close(&mut self, _: Python<'_>) {} fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { Ok(()) @@ -988,7 +1169,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) .unwrap_err(); assert!(!error.is_instance_of::(py)); - assert_eq!(log.entries(), ["begin", "route", "adapter.close"]); + assert_eq!( + log.entries(), + ["started", "begin", "route", "adapter.close"] + ); }); } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index bb0b5b1c3b1..8889b1513db 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,9 +1,10 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and //! asyncio glue, and the driver that runs a native [`Machine`](litellm_callbacks::machine::Machine) -//! against a Python route host and a callback adapter. Everything here is Python-specific by +//! against a Python route host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. mod adapter; +mod argument; mod callable; mod driver; mod execution; @@ -11,7 +12,10 @@ mod gil; mod handle; mod marshal; -pub use adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +pub use adapter::{ + HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, +}; +pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value}; diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 87945bf8785..2e398d0287e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -150,7 +150,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { api_key: inputs.api_key.and_then(|key| { inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(Some(key)) }), api_base: inputs.api_base.and_then(|base| { @@ -592,12 +592,17 @@ impl AzureDocumentIntelligenceOcrConfig { )?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::super::common_utils::validate_destination(connection, key.source())?; return Ok( @@ -796,7 +801,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2012f740173..7ef051e8986 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -142,12 +142,17 @@ impl AzureAiOcrConfig { super::common_utils::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::common_utils::validate_destination(connection, key.source())?; return Ok(bearer_headers(connection, key.value())); @@ -196,7 +201,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -288,7 +293,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e6fe5d9556d..8321dcfb4ce 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, future::Future, time::Duration}; -use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, @@ -90,21 +90,22 @@ pub enum OcrResponseFormat { #[derive(Clone, Default)] pub struct OcrCredentialInputs { - pub api_key: Option>, - pub dynamic_api_key: Option>, + pub api_key: Option>, + pub dynamic_api_key: Option>, pub api_base: Option>, pub dynamic_api_base: Option>, } impl OcrCredentialInputs { pub fn new( - api_key: Option, + api_key: Option, api_key_source: InputSource, api_base: Option, api_base_source: InputSource, ) -> Self { Self { - api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)), + api_key: nonblank(api_key.as_ref().map(|key| key.expose().to_string())) + .map(|value| Sourced::new(SecretValue::new(value), api_key_source)), dynamic_api_key: None, api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)), dynamic_api_base: None, @@ -159,7 +160,7 @@ fn nonblank(value: Option) -> Option { #[derive(Clone)] pub struct OcrConnection { - pub api_key: Option, + pub api_key: Option, pub api_key_source: InputSource, pub api_base: Option, pub api_base_source: InputSource, @@ -209,7 +210,7 @@ impl Default for OcrConnection { #[derive(Clone, Default)] pub struct ResolvedOcrCredentials { - pub api_key: Option>, + pub api_key: Option>, pub api_base: Option>, } @@ -428,7 +429,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { ResolvedOcrCredentials { api_key: inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(inputs.api_key), api_base: inputs .dynamic_api_base diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index f353c22d8c4..2528c967f41 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -179,8 +179,8 @@ impl CohereParseConfig { } let key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -718,7 +718,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index e93ddee3c50..7635bdd3d04 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -3,9 +3,9 @@ use std::{sync::OnceLock, time::Duration}; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_callbacks::event::WireRequest; use serde::{Serialize, de::DeserializeOwned}; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::{ base_llm::ocr::{ @@ -26,11 +26,7 @@ use crate::{ /// The route's view of one call, handed to provider code that has to reach the /// caller's hooks mid-flight (guardrails on the outgoing body, raw response events). pub trait CallHooks: Send + Sync { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result>; + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result>; fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), E>>; } @@ -231,9 +227,8 @@ pub async fn transform_request_body( config.get_supported_ocr_params(&request.model), )?; config.validate_request_body(&composed)?; - let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed); let changed = hooks - .before_send(wire_request(url, headers, composed), passthrough_fields) + .before_send(wire_request(url, headers, composed)) .await?; if !changed.body.is_object() { return Err(Error::RequestField { @@ -252,21 +247,6 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq } } -fn caller_inputs(request: &PreparedOcrRequest) -> Result, Error> { - let document = request - .caller_document - .then(|| serde_json::to_value(&request.document)) - .transpose() - .map_err(|_| Error::RequestField { - path: "document".into(), - })?; - let params: Map = request.optional_params.clone().into(); - Ok(params - .into_iter() - .chain(document.map(|document| ("document".to_string(), document))) - .collect()) -} - pub fn build_http_request( client: &OcrClient, request: &PreparedOcrRequest, @@ -294,9 +274,7 @@ pub async fn guardrail_document( let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField { path: "document".into(), })?; - let changed = hooks - .before_send(wire_request(url, headers, body), Passthrough::default()) - .await?; + let changed = hooks.before_send(wire_request(url, headers, body)).await?; let document = decode_request_value(changed.body, "guardrail.document")?; Ok((document, changed.headers)) } diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index c2038d0552d..9028f09c5ab 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -135,8 +135,8 @@ impl MistralOcrConfig { } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -212,7 +212,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(str::to_string), + api_key: api_key.map(litellm_auth::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index ca2bae9c3bb..ec876fafb8f 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -442,8 +442,8 @@ fn resolve_headers( } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -629,7 +629,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some("passed-key".into()), + api_key: Some(litellm_auth::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -639,7 +639,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index ea0bcf3d08c..c2cb23d0010 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -134,7 +134,10 @@ impl VertexAiOcrConfig { .vertex_auth() .validate_environment( connection.extra_headers.clone(), - connection.api_key.as_deref(), + connection + .api_key + .as_ref() + .map(litellm_auth::SecretValue::expose), config, &credential_env, ) diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 9932594e2f5..5dccfb4aca8 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,7 +1,7 @@ - Target invariants, not completion claims; these supersede older conflicting bridge guidance - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, `passthrough_fields` re-aliasing) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy + - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy` owns `Logging` dispatch policy - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 9dc891a91d6..77212e3d38e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,9 +1,9 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{RouteHost, missing_state, to_py}; +use litellm_host_python::{HostOpError, RouteHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ - exceptions::PyBaseException, + exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -57,12 +57,8 @@ impl OcrRouteHost { .ok_or_else(missing_state)? .acquire(py) } -} -impl RouteHost for OcrRouteHost { - type Route = Ocr; - - fn invoke( + fn answer( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, @@ -88,6 +84,40 @@ impl RouteHost for OcrRouteHost { } } + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let provider = match &self.data { + OcrHostData::Projected(handles) => handles.provider, + _ => "", + }; + let mapped = py + .import("litellm.rust_bridge.ocr.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), provider))) + .and_then(|mapped| mapped.extract::>().map_err(PyErr::from)); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for OcrRouteHost { + type Route = Ocr; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: OcrOp, + ) -> Result> { + self.answer(py, arguments, op) + .map_err(|error| HostOpError::Python(self.map_failure(py, error))) + } + fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { py.import("litellm.rust_bridge.ocr.route_host")? .getattr("response")? @@ -95,27 +125,14 @@ impl RouteHost for OcrRouteHost { .map(Bound::unbind) } - fn native_error(error: Error) -> PyErr { - ocr_error_to_pyerr(error) + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } fn host_error(error: &PyErr) -> Error { Error::InvalidRequest(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - let provider = match &self.data { - OcrHostData::Projected(handles) => handles.provider, - _ => "", - }; - let mapped: Py = py - .import("litellm.rust_bridge.ocr.route_host")? - .getattr("map_failure")? - .call1((error.value(py), self.request.bind(py), provider))? - .extract()?; - Ok(PyErr::from_value(mapped.into_bound(py).into_any())) - } - fn close(&mut self, _: Python<'_>) { self.data = OcrHostData::Released; } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 7ffa129f85c..5dd2aa804b8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,3 +1,4 @@ +use litellm_auth::SecretValue; use litellm_core::ocr::{ types::{LiteLLMOcrRequest, OcrDocumentInput}, wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, @@ -31,7 +32,7 @@ struct OcrArguments<'a, 'py> { impl<'py> OcrArguments<'_, 'py> { fn lookup(&self, name: &str) -> PyResult> { - litellm_callbacks_legacy::lookup(self.kwargs, self.request, name)? + litellm_host_python::lookup(self.kwargs, self.request, name)? .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) } @@ -47,8 +48,11 @@ impl<'py> OcrArguments<'_, 'py> { self.lookup("document") } - fn api_key(&self) -> PyResult> { - self.lookup("api_key")?.extract() + fn api_key(&self) -> PyResult> { + Ok(self + .lookup("api_key")? + .extract::>()? + .map(SecretValue::new)) } fn api_base(&self) -> PyResult> { diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4a9a65b1485..7248c2f3590 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1265,11 +1265,6 @@ class Logging(LiteLLMLoggingBaseClass): additional_args.get("api_base", "") ) - def record_api_call_start_time(self) -> None: - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] - def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API try: @@ -1334,7 +1329,15 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.record_api_call_start_time() + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + # Set-once first provider-handoff instant. api_call_start_time + # is overwritten on every retry, so it can't measure one-time + # preprocessing; pinning the first attempt excludes retry loops + # + backoff. Logging object only — must NOT go into + # litellm_params["metadata"] (caller request metadata, typed + # Dict[str, str], echoed downstream; a datetime breaks it). + if self.model_call_details.get("first_api_call_start_time") is None: + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1468,21 +1471,16 @@ class Logging(LiteLLMLoggingBaseClass): """ return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) - def record_post_call( - self, original_response: object, input: object, api_key: object, additional_args: dict[str, object] - ) -> None: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" - def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: - self.record_post_call(original_response, input, api_key, additional_args) + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" attr: Literal["warning", "debug"] if self.litellm_request_debug: @@ -2177,7 +2175,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, - build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2202,9 +2199,6 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) - if not build_logging_payload: - return - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2266,7 +2260,6 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, - build_logging_payload: bool = True, ): try: if start_time is None: @@ -2304,7 +2297,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, - build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3328,9 +3320,7 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn( - self, exception, traceback_exception, start_time=None, end_time=None, build_logging_payload: bool = True - ): + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time if end_time is None: @@ -3365,9 +3355,6 @@ class Logging(LiteLLMLoggingBaseClass): metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) - if not build_logging_payload: - return start_time, end_time - ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index e05d9368fa8..406dc55cfea 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -6,24 +6,22 @@ registries it fans out to. It expires with that contract. from __future__ import annotations +import contextvars import datetime -import os +import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Coroutine, Mapping from dataclasses import dataclass from typing import ( TYPE_CHECKING, Final, - Literal, Protocol, - TypeAlias, cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations ) -from typing_extensions import assert_never - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CredentialItem class MetadataUpdater(Protocol): @@ -42,7 +40,6 @@ class MetadataUpdater(Protocol): class CallSetup: logger: Logging kwargs: dict[str, object] - bridge_owned: bool def setup( @@ -61,9 +58,9 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments, bridge_owned=False) + return CallSetup(supplied, arguments) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared, bridge_owned=True) + return CallSetup(logger, prepared) def check_limits(kwargs: Mapping[str, object]) -> None: @@ -93,87 +90,219 @@ def finalize( update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) -def deployment_callbacks_needed() -> bool: - import litellm - from litellm.integrations.custom_logger import CustomLogger +class LoggingSurface(Protocol): + def update_from_kwargs( + self, + kwargs: dict[str, object], + litellm_params: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + model: str | None = None, + user: str | None = None, + **additional_params: object, + ) -> None: ... - return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) + def pre_call( + self, input: object, api_key: object, model: object = None, additional_args: dict[str, object] = ... + ) -> object: ... + + def post_call( + self, + original_response: object, + input: object = None, + api_key: object = None, + additional_args: dict[str, object] = ..., + ) -> object: ... + + def handle_sync_success_callbacks_for_async_calls( + self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: object = None + ) -> None: ... + + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: ... + + def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> Coroutine[object, object, None]: ... + + def success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> None: ... + + def async_success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> Coroutine[object, object, None]: ... -Phase: TypeAlias = Literal[ - "input", "sync_success", "sync_success_async", "async_success", "sync_failure", "async_failure", "payload" -] +if TYPE_CHECKING: + _LOGGING_CONFORMS: type[LoggingSurface] = Logging -def callbacks_needed(logger: Logging, phase: Phase) -> bool: - import litellm - from litellm._logging import ( - _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging +class LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... + + +class DeploymentHook(Protocol): + def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... + + +class DeploymentSuccessHook(Protocol): + def __call__(self, request_data: dict[str, object], response: object, call_type: object) -> Awaitable[object]: ... + + +class DeploymentFailureHook(Protocol): + def __call__(self, request_data: Mapping[str, object], exception: Exception, call_type: str) -> Awaitable[None]: ... + + +def update_logging( + logger: LoggingSurface, + kwargs: dict[str, object], + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + custom_llm_provider: str, +) -> None: + logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, ) - if ( - _is_debugging_on() - or getattr(logger, "litellm_request_debug", False) - or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") - ): - return True - input_needed: Final = bool( - litellm.input_callback - or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_input_callbacks - or callable(getattr(logger, "logger_fn", None)) - or logger.log_raw_request_response - or litellm.log_raw_request_response + +def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: + logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) + + +def post_call( + logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] +) -> None: + logger.post_call(original_response=original_response, api_key=api_key, additional_args=additional_args) + + +def defers_async_logging(logger: LoggingSurface) -> bool: + return bool(getattr(logger, "_defer_async_logging", False)) + + +def defer_success(logger: LoggingSurface, pending: object) -> None: + setattr(logger, "_native_pending_logging", pending) + + +def sync_success_for_async_call( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> None: + logger.handle_sync_success_callbacks_for_async_calls(result=response, start_time=start, end_time=end) + + +def failure_handler( + logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> Coroutine[object, object, None] | None: + trace: Final = "".join(traceback.format_exception(error)) + if asynchronous: + return logger.async_failure_handler(error, trace, start, end) + logger.failure_handler(error, trace, start, end) + return None + + +def submit_success(logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime) -> None: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, logger.success_handler, response, start, end) + + +def async_success_handler( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> Coroutine[object, object, None]: + return logger.async_success_handler(response, start, end) + + +def enqueue_logging(coroutine: Coroutine[object, object, None]) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + worker: Final = cast( # cast-ok: bounded adapter for the untyped logging worker + LoggingWorker, GLOBAL_LOGGING_WORKER ) - match phase: - case "input": - return input_needed - case "sync_success": - return bool(litellm.success_callback or logger.dynamic_success_callbacks) - case "sync_success_async": - return bool( - (litellm.success_callback or logger.dynamic_success_callbacks) - and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks - ) - case "async_success": - return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "sync_failure": - return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) - case "async_failure": - return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "payload": - return bool( - input_needed - or litellm.success_callback - or litellm.failure_callback - or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_success_callbacks - or logger.dynamic_async_success_callbacks - or logger.dynamic_failure_callbacks - or logger.dynamic_async_failure_callbacks - ) - case _: - assert_never(phase) + contextvars.copy_context().run(worker.ensure_initialized_and_enqueue, coroutine) -def success_bookkeeping( - logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_success" if asynchronous else "sync_success" - if logger.should_run_logging(phase): - logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload - result=response, start_time=start, end_time=end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def restore_context(logger: LoggingSurface) -> None: + from litellm.utils import ( + _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # the @client wrapper restores the same correlation context + ) + + _restore_correlation_context_if_supported(logger) -def failure_bookkeeping( - logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_failure" if asynchronous else "sync_failure" - if logger.should_run_logging(phase): - logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload - error, "", start, end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def custom_pricing_fields() -> tuple[str, ...]: + from litellm.types.utils import CustomPricingLiteLLMParams + + return tuple(CustomPricingLiteLLMParams.model_fields) + + +def is_internal_call() -> bool: + from litellm._internal_context import is_internal_call as internal + + return internal.get() + + +def credential_list() -> list[CredentialItem]: + import litellm + + return litellm.credential_list + + +def warn_unknown_credential(name: str, loaded: int) -> None: + from litellm._logging import verbose_logger + + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + name, + loaded, + ) + + +def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitable[object]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentHook, utils.async_pre_call_deployment_hook + ) + return hook(kwargs, call_type) + + +def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]: + from litellm import utils + from litellm.types.utils import CallTypes + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentSuccessHook, utils.async_post_call_success_deployment_hook + ) + return hook(kwargs, response, CallTypes(call_type)) + + +def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_type: str) -> Awaitable[None]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentFailureHook, utils.async_post_call_failure_deployment_hook + ) + return hook(kwargs, error, call_type) diff --git a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py index a4474c85230..a0906c7c5be 100644 --- a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py +++ b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py @@ -1,12 +1,16 @@ import datetime +import inspect from collections.abc import Mapping +from pathlib import Path from types import MappingProxyType from typing import Final import pytest +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.rust_bridge import legacy_callbacks as legacy from litellm.rust_bridge.legacy_callbacks import check_limits, setup _OCR_KWARGS: Final = MappingProxyType( @@ -56,13 +60,12 @@ def _supplied_logger() -> Logging: ) -def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: +def test_setup_reuses_a_supplied_logger() -> None: supplied: Final = _supplied_logger() result: Final = setup( "aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True ) assert result.logger is supplied - assert result.bridge_owned is False @pytest.mark.parametrize( @@ -73,7 +76,15 @@ def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: ], ids=["ocr", "embedding"], ) -def test_setup_owns_every_logger_it_builds(call_type: str, kwargs: Mapping[str, object]) -> None: +def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Mapping[str, object]) -> None: result: Final = setup(call_type, (), kwargs, datetime.datetime.now(), asynchronous=True) - assert result.bridge_owned is True assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] + + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy/python_contract.json" + + +def test_the_rust_contract_matches_the_shim_signatures() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 085ea4a14c0..5fca927bea3 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -4,7 +4,7 @@ import gc import json import threading import weakref -from collections.abc import Coroutine +from collections.abc import Awaitable, Callable, Coroutine from contextvars import ContextVar from typing import Final @@ -400,7 +400,7 @@ async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: Rec @pytest.mark.asyncio @pytest.mark.parametrize("failure", [False, True]) -async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( +async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper( ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: bool, @@ -414,8 +414,12 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( submissions = 0 enqueues = 0 - def deployment(self, *args: object, **kwargs: object) -> None: - self.deployments += 1 + def counting(self, hook: Callable[..., Awaitable[object]]) -> Callable[..., Awaitable[object]]: + async def counted(*args: object, **kwargs: object) -> object: + self.deployments += 1 + return await hook(*args, **kwargs) + + return counted def submit(self, *args: object, **kwargs: object) -> None: self.submissions += 1 @@ -430,7 +434,7 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( "async_post_call_success_deployment_hook", "async_post_call_failure_deployment_hook", ): - monkeypatch.setattr(utils, name, probe.deployment) + monkeypatch.setattr(utils, name, probe.counting(getattr(utils, name))) monkeypatch.setattr(litellm_logging, "executor", probe) monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) if failure: @@ -447,17 +451,16 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( assert response._hidden_params["response_cost"] is not None assert response._hidden_params["_response_ms"] > 0 assert trace_id_var.get() == "callback-free-parent" - assert probe.deployments == probe.submissions == probe.enqueues == 0 + assert probe.deployments == 2 + assert probe.submissions == probe.enqueues == 0 assert len(created_loggers) == 1 logger: Final = created_loggers[0] - assert not hasattr(logger, "_native_pending_logging") - assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] - assert "standard_logging_object" not in logger.model_call_details - assert ( - "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None - ) - assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) - assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) + if failure: + assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] + assert logger.model_call_details["response_cost"] == 0 + else: + assert getattr(logger, "_native_pending_logging", None) is not None + assert "end_time" not in logger.model_call_details @pytest.mark.asyncio From a737e3430a979b78c4d84ceac6d7605aca729d79 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:41:50 -0700 Subject: [PATCH 20/33] test(rust): property and parametrized tests for legacy callback contracts The payload boundary of callbacks-legacy gets a model-based proptest: for any JSON body, caller keywords and callback edit, keywords the route sends unchanged reach pre_call as the caller's own objects, and the wire is the body pre_call received as the callback left it. A parametrized test pins that a keyword the bridge never reads keeps its identity through setup, the deployment hook, check_limits and prepare. Behaviour owned by the real Logging object is pinned end to end in the OCR tests: a hypothesis version of the body property over HTTP, sync hooks seeing no running event loop, retained payloads staying intact after the call, success callbacks sharing one standard logging payload, and state stashed before a blocking deployment hook raises reaching both failure callback families --- litellm-rust/Cargo.toml | 1 + .../crates/callbacks-legacy/Cargo.toml | 1 + .../tests/deployment_hooks.rs | 37 ++++ .../crates/callbacks-legacy/tests/payload.rs | 175 +++++++++++++++++- tests/test_litellm_rust/ocr/test_callbacks.py | 163 +++++++++++++++- 5 files changed, 370 insertions(+), 7 deletions(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index ffdbf64bb49..32f925b8d8b 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,6 +26,7 @@ litellm-token-counter = { path = "crates/token-counter" } litellm-host-python = { path = "crates/host-python" } bytes = "1" +proptest = "1.7.0" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index 3cf9382f857..efacde051ed 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -16,5 +16,6 @@ serde_json.workspace = true [dev-dependencies] litellm-auth.workspace = true +proptest.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index ea3510de17e..7bad09c7890 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -97,6 +97,43 @@ assert checked is prepared }); } +#[rstest] +#[case::synchronous(false)] +#[case::asynchronous(true)] +fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( + #[case] asynchronous: bool, +) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +opaque = object() +hooked = [] +logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} +kwargs = {'logger': logger, 'vendor_extension': opaque} +", + ); + let (mut logging, step) = begin(py, &locals, asynchronous); + let step = match step { + LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), + step => step, + }; + locals.set_item("prepared", arguments(py, step)).unwrap(); + locals.set_item("asynchronous", asynchronous).unwrap(); + run( + py, + &locals, + c" +assert prepared['vendor_extension'] is opaque +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked['vendor_extension'] is opaque +assert hooked == ([opaque] if asynchronous else []), hooked +", + ); + }); +} + #[test] fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 68c0a2b1e15..5c49383bf37 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -2,10 +2,11 @@ use std::ffi::CStr; use litellm_auth::SecretValue; use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleStep, PythonLifecycle}; +use litellm_host_python::{LifecycleStep, PythonLifecycle, to_py}; +use proptest::prelude::*; use pyo3::prelude::*; use rstest::rstest; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; use super::LegacyLogging; use crate::PythonLogger; @@ -57,10 +58,24 @@ fn before_send_with_secrets( optional_params: Value, body: Value, secret_fields: &[&str], +) -> WireRequest { + before_send_bound(&[], script, optional_params, body, secret_fields) +} + +/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. +fn before_send_bound( + bindings: &[(&str, &Value)], + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], ) -> WireRequest { Python::initialize(); Python::attach(|py| { let locals = namespace(py, PAYLOAD_LOGGER); + for &(name, value) in bindings { + locals.set_item(name, to_py(py, value).unwrap()).unwrap(); + } run(py, &locals, script); let mut logging = LegacyLogging { logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), @@ -346,3 +361,159 @@ def check(): json!({"document": document(DOCUMENT), "include_image_base64": true}) ); } + +/// What one pre-call callback does to the payload it is handed. +#[derive(Clone, Debug)] +enum Edit { + Nothing, + Set(String, Value), + Remove(String), + Rebind(Value), + RebindThenSetRetained(String, Value), +} + +impl Edit { + fn script(&self) -> Value { + match self { + Self::Nothing => json!({"kind": "nothing"}), + Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), + Self::Remove(key) => json!({"kind": "remove", "key": key}), + Self::Rebind(value) => json!({"kind": "rebind", "value": value}), + Self::RebindThenSetRetained(key, value) => { + json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) + } + } + } + + /// The legacy contract: the provider is sent the body object `pre_call` received, as + /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and + /// leaves that object alone. + fn sent(&self, body: &Map) -> Value { + let mut sent = body.clone(); + match self { + Self::Nothing | Self::Rebind(_) => {} + Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { + sent.insert(key.clone(), value.clone()); + } + Self::Remove(key) => { + sent.remove(key); + } + } + Value::Object(sent) + } +} + +/// How the caller's keyword for a body key relates to what the route sends under it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Caller { + PassedUnchanged, + RewrittenByTheRoute, + NotPassed, +} + +const MODEL: &CStr = c" +aliased = {} +def on_pre_call(args): + body = args['complete_input_dict'] + aliased.update({name: body[name] is kwargs[name] for name in unchanged}) + kind = edit['kind'] + if kind == 'set': + body[edit['key']] = edit['value'] + elif kind == 'remove': + body.pop(edit['key'], None) + elif kind == 'rebind': + args['complete_input_dict'] = edit['value'] + elif kind == 'rebind_then_set_retained': + args['complete_input_dict'] = {} + body[edit['key']] = edit['value'] +def check(): + assert aliased == {name: True for name in unchanged}, aliased + assert logger.names() == ['pre_call', 'post_call'], logger.calls +"; + +fn json_value() -> impl Strategy { + let leaf = prop_oneof![ + Just(Value::Null), + any::().prop_map(Value::from), + any::().prop_map(Value::from), + any::() + .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) + .prop_map(Value::from), + ".{0,8}".prop_map(Value::from), + ]; + leaf.prop_recursive(3, 24, 4, |inner| { + prop_oneof![ + prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), + prop::collection::btree_map(key(), inner, 0..4) + .prop_map(|fields| Value::Object(fields.into_iter().collect())), + ] + }) +} + +fn key() -> impl Strategy { + "[a-z]{1,6}" +} + +fn caller() -> impl Strategy { + prop_oneof![ + Just(Caller::PassedUnchanged), + Just(Caller::RewrittenByTheRoute), + Just(Caller::NotPassed), + ] +} + +fn edit() -> impl Strategy { + prop_oneof![ + Just(Edit::Nothing), + (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), + key().prop_map(Edit::Remove), + json_value().prop_map(Edit::Rebind), + (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), + ] +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(128))] + + /// For any body, any caller keywords and any callback edit: every keyword the route + /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is + /// sent exactly what the model says, so a callback that edits nothing changes nothing. + #[test] + fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( + fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), + edit in edit(), + ) { + let body: Map = fields + .iter() + .map(|(name, (value, _))| (name.clone(), value.clone())) + .collect(); + let kwargs: Map = fields + .iter() + .filter_map(|(name, (value, caller))| match caller { + Caller::PassedUnchanged => Some((name.clone(), value.clone())), + Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), + Caller::NotPassed => None, + }) + .collect(); + let unchanged: Value = fields + .iter() + .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) + .map(|(name, _)| Value::from(name.clone())) + .collect(); + + let wire = before_send_bound( + &[ + ("kwargs", &Value::Object(kwargs)), + ("unchanged", &unchanged), + ("edit", &edit.script()), + ], + MODEL, + json!({}), + Value::Object(body.clone()), + &[], + ); + + prop_assert_eq!(wire.body, edit.sent(&body)); + prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } +} diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 27cdcc4d997..7cc19b8c090 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -1,24 +1,28 @@ import asyncio import copy +import gc import queue import threading from typing import Final import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, + call_native, call_native_aocr, call_native_ocr, request_body, request_headers, ) -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -123,9 +127,7 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ "callbacks": [Retain(), Edit()], } response: Final = ( - await call_native_aocr(ocr_server, **arguments) - if asynchronous - else call_native_ocr(ocr_server, **arguments) + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) ) assert aliases == [True] @@ -291,6 +293,157 @@ def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registere assert "log_failure_event" not in recorder.names +JSON_SCALARS: Final = ( + st.none() + | st.booleans() + | st.integers(min_value=-(2**63), max_value=2**63 - 1) + | st.floats(allow_nan=False, allow_infinity=False) + | st.text(max_size=8) +) +JSON_VALUES: Final = st.recursive( + JSON_SCALARS, + lambda children: st.lists(children, max_size=3) | st.dictionaries(st.text(max_size=6), children, max_size=3), + max_leaves=8, +) + + +LATEST_EDITS: Final[list[dict[str, object]]] = [] + + +class ApplyLatestEdits(CustomLogger): + """Registrations can outlive one hypothesis example, so every instance applies the current example's edits.""" + + def __init__(self, latest: list[dict[str, object]]) -> None: + super().__init__() + self.latest = latest + + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs).update(copy.deepcopy(self.latest[-1])) + + +@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +@given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3)) +def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it( + ocr_server: RecordingServer, edits: dict[str, object] +) -> None: + LATEST_EDITS.append(edits) + + call_native_ocr_with_callbacks(ocr_server, [ApplyLatestEdits(LATEST_EDITS)]) + + assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} + + +@pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"]) +def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None: + recorder: Final = RecordingLogger() + + call_native_ocr_with_callbacks(ocr_server, [recorder]) + + [event] = recorder.wait_for(hook) + assert event.loop is None + assert (event.thread is threading.current_thread()) == (hook == "log_pre_api_call") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + retained: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + retained.append((kwargs, request_body(kwargs), request_headers(kwargs))) + + await call_native(ocr_server, asynchronous, callbacks=[Retain()]) + await drain_logging() + gc.collect() + + [(details, body, headers)] = retained + assert body == ocr_server.requests[0].body + assert headers + assert all(ocr_server.requests[0].headers[name] == value for name, value in headers.items()) + assert details["additional_args"]["complete_input_dict"] is body + assert details["additional_args"]["headers"] is headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("family", ["sync", "async"]) +async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None: + queued: Final = [] + finished: Final = threading.Event() + + def queue_payload(kwargs: dict[str, object]) -> None: + queued.append(kwargs["standard_logging_object"]) + + def strip_payload(kwargs: dict[str, object]) -> None: + payload: Final = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["stripped-by-a-later-callback"] = True + finished.set() + + class QueuePayload(CustomLogger): + if family == "sync": + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + queue_payload(kwargs) + + else: + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + queue_payload(kwargs) + + class StripPayload(CustomLogger): + if family == "sync": + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + strip_payload(kwargs) + + else: + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + strip_payload(kwargs) + + await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()]) + await drain_logging() + + assert await asyncio.to_thread(finished.wait, 10) + assert [payload["stripped-by-a-later-callback"] for payload in queued] == [True] + + +@pytest.mark.asyncio +async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks( + ocr_server: RecordingServer, +) -> None: + token: Final = object() + observed: Final = [] + + class Blocked(Exception): + pass + + class Block(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token + raise Blocked("blocked after the provider answered") + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("success", None, None)) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs.get("blocked-by"), kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs.get("blocked-by"), kwargs["exception"])) + + litellm.callbacks.append(Block()) + + with pytest.raises(Blocked) as raised: + await call_native_aocr(ocr_server) + await drain_logging() + + assert observed == [("sync", token, raised.value), ("async", token, raised.value)] + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context( From 30f7b8442bd8b7258ad4bc5a41a9b85546d7a83c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:42:59 -0700 Subject: [PATCH 21/33] chore(rust): lock proptest --- litellm-rust/Cargo.lock | 62 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0265e0adbc2..bf58a81a6c5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2044,6 +2044,7 @@ dependencies = [ "litellm-auth", "litellm-callbacks", "litellm-host-python", + "proptest", "pyo3", "rstest", "serde_json", @@ -2601,6 +2602,25 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand 0.9.5", + "rand_chacha 0.9.0", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -2682,6 +2702,12 @@ dependencies = [ "serde", ] +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quinn" version = "0.11.11" @@ -2845,6 +2871,15 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core 0.9.5", +] + [[package]] name = "rayon" version = "1.12.0" @@ -3244,6 +3279,18 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + [[package]] name = "ryu" version = "1.0.23" @@ -4099,6 +4146,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "unicase" version = "2.9.0" @@ -4212,6 +4265,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + [[package]] name = "walkdir" version = "2.5.0" From ab27ee7efcbe982991095dbfbf5bd01662f8e69c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:44:19 -0700 Subject: [PATCH 22/33] test(rust): fix request count and header case in new OCR callback tests --- tests/test_litellm_rust/ocr/test_callbacks.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 7cc19b8c090..45b99d19d90 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -326,6 +326,7 @@ class ApplyLatestEdits(CustomLogger): def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it( ocr_server: RecordingServer, edits: dict[str, object] ) -> None: + ocr_server.expected_requests = None LATEST_EDITS.append(edits) call_native_ocr_with_callbacks(ocr_server, [ApplyLatestEdits(LATEST_EDITS)]) @@ -362,7 +363,7 @@ async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( [(details, body, headers)] = retained assert body == ocr_server.requests[0].body assert headers - assert all(ocr_server.requests[0].headers[name] == value for name, value in headers.items()) + assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items()) assert details["additional_args"]["complete_input_dict"] is body assert details["additional_args"]["headers"] is headers From dd2e6c17bc19b839ff0a5e8500dfe40a7612f430 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:56:13 +0000 Subject: [PATCH 23/33] fix(rust): preserve OCR callback headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/callbacks-legacy/src/adapter.rs | 12 ++++++++++-- .../crates/callbacks-legacy/src/callbacks.rs | 3 +++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 7204207cd62..53aa6ce9d2e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -44,6 +44,7 @@ pub struct LegacyLogging { response: Option>, error: Option>, body: Option>, + headers: Option>, context: Option, asynchronous: bool, internal: bool, @@ -77,6 +78,7 @@ impl LegacyLogging { response: None, error: None, body: None, + headers: None, context: None, asynchronous, internal: false, @@ -245,6 +247,7 @@ impl PythonLifecycle for LegacyLogging { headers.set_item(name, value)?; } self.body = Some(body.clone().unbind()); + self.headers = Some(headers.clone().unbind()); self.context = Some(context.clone()); self.logger()?.pre_call( py, @@ -299,8 +302,13 @@ impl PythonLifecycle for LegacyLogging { .as_ref() .and_then(|context| context.api_key.as_ref()) .map(|api_key| api_key.expose()); - self.logger()? - .post_call(py, &raw.body, api_key, self.body.as_ref())?; + self.logger()?.post_call( + py, + &raw.body, + api_key, + self.body.as_ref(), + self.headers.as_ref(), + )?; Ok(LifecycleStep::Done) } (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index 9464f1d6612..9fcfe98368e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -38,6 +38,7 @@ pub trait LegacyCallbacks { original_response: &str, api_key: Option<&str>, body: Option<&Py>, + headers: Option<&Py>, ) -> PyResult<()>; fn defers_async_logging(&self, py: Python<'_>) -> bool; @@ -150,9 +151,11 @@ impl LegacyCallbacks for PythonLogger { original_response: &str, api_key: Option<&str>, body: Option<&Py>, + headers: Option<&Py>, ) -> PyResult<()> { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; + additional.set_item("headers", headers)?; Logging::PostCall.call( py, (self.object(py), original_response, api_key, &additional), From 4627ec4ea8f4ab6b53388f079bab54075d05a0ea Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:58:14 +0000 Subject: [PATCH 24/33] test(rust): cover post-call header identity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/callbacks-legacy/tests/payload.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 5c49383bf37..43128bc38ea 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -331,15 +331,19 @@ def on_pre_call(args): } #[test] -fn post_call_receives_the_raw_response_the_route_key_and_the_body_pre_call_saw() { +fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { before_send( c" def check(): original_response, api_key, additional_args = logger.post assert original_response == 'raw response', original_response assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) - assert additional_args == {'complete_input_dict': logger.pre['complete_input_dict']}, additional_args + assert additional_args == { + 'complete_input_dict': logger.pre['complete_input_dict'], + 'headers': logger.pre['headers'], + }, additional_args assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] + assert additional_args['headers'] is logger.pre['headers'] ", json!({"document": document(DOCUMENT)}), ); From 1a52bae77950ddab49f1d53a0a511fdaf1d7b570 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:32:28 -0700 Subject: [PATCH 25/33] add streaming message --- .../callbacks-legacy/python_contract.json | 17 ++ .../crates/callbacks-legacy/src/adapter.rs | 90 +++++++- .../callbacks-legacy/src/legacy_python.rs | 30 ++- .../crates/callbacks-legacy/tests/support.rs | 5 + litellm-rust/crates/callbacks/src/event.rs | 4 + litellm-rust/crates/callbacks/src/host.rs | 21 ++ litellm-rust/crates/callbacks/src/route.rs | 5 + litellm-rust/crates/callbacks/src/run.rs | 4 + litellm-rust/crates/core/src/machine/mod.rs | 17 +- .../crates/core/src/messages/handler.rs | 128 +++++------- litellm-rust/crates/core/src/messages/mod.rs | 36 +++- .../crates/core/src/messages/prepare.rs | 11 +- .../crates/core/src/messages/route.rs | 197 ++++++++++++++++++ litellm-rust/crates/core/src/ocr/route.rs | 2 + litellm-rust/crates/core/tests/ocr.rs | 2 + .../crates/host-python/src/adapter.rs | 8 + litellm-rust/crates/host-python/src/driver.rs | 77 ++++++- litellm-rust/crates/host-python/src/handle.rs | 13 ++ .../crates/python-bridge/src/marshal.rs | 16 -- .../python-bridge/src/routes/messages.rs | 88 -------- .../python-bridge/src/routes/messages/host.rs | 186 +++++++++++++++++ .../python-bridge/src/routes/messages/mod.rs | 64 ++++++ .../crates/python-bridge/src/routes/mod.rs | 31 --- .../python-bridge/src/routes/ocr/host.rs | 4 + litellm/rust_bridge/_native.pyi | 28 +-- litellm/rust_bridge/catalog.py | 1 + litellm/rust_bridge/failures.py | 43 ++++ litellm/rust_bridge/legacy_callbacks.py | 80 +++++++ litellm/rust_bridge/lifecycle.py | 137 ++++++++++-- litellm/rust_bridge/messages/entrypoints.py | 10 +- litellm/rust_bridge/messages/route_host.py | 2 +- litellm/rust_bridge/ocr/route_host.py | 40 +--- .../rust_bridge/native_route_wheel_test.py | 31 +-- .../test_litellm/rust_bridge/test_bindings.py | 4 +- .../test_litellm/rust_bridge/test_catalog.py | 4 + .../test_litellm/rust_bridge/test_runtime.py | 1 - tests/test_litellm_rust/messages/__init__.py | 0 .../messages/test_callbacks.py | 175 ++++++++++++++++ .../support/recording_server.py | 16 +- tests/test_litellm_rust/support/requests.py | 35 ++++ 40 files changed, 1319 insertions(+), 344 deletions(-) create mode 100644 litellm-rust/crates/core/src/messages/route.rs delete mode 100644 litellm-rust/crates/python-bridge/src/routes/messages.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/host.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/mod.rs create mode 100644 tests/test_litellm_rust/messages/__init__.py create mode 100644 tests/test_litellm_rust/messages/test_callbacks.py diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json index 840c0abfa45..a09bdc711a3 100644 --- a/litellm-rust/crates/callbacks-legacy/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -94,5 +94,22 @@ "kwargs", "error", "call_type" + ], + "stream_opened": [ + "logger" + ], + "stream_success": [ + "logger", + "request_body", + "chunks", + "start", + "end", + "first_chunk" + ], + "stream_failure": [ + "logger", + "request_body", + "chunks", + "error" ] } diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 53aa6ce9d2e..9d28db92add 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -2,7 +2,9 @@ //! raises is answered with the same `Logging` calls, in the same order, as the Python //! `@client` path makes them. -use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest}; +use litellm_callbacks::event::{ + CallEvent, FailureOrigin, RequestContext, Timing, WireRequest, epoch_seconds, +}; use litellm_host_python::{ LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, }; @@ -10,14 +12,16 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::PyDict, + types::{PyDict, PyList}, }; use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, - finalize, is_internal_call, prepare, setup, + finalize, is_internal_call, + legacy_python::Streaming, + prepare, setup, }; /// What the legacy contract needs to know about the route it is logging. @@ -28,6 +32,12 @@ pub struct LegacySurface { pub input_description: &'static str, } +/// What the Messages stream iterator keeps for its end-of-stream billing. +struct DeliveredStream { + chunks: Py, + first_chunk: Option>, +} + enum Pending { DeploymentPreCall, DeploymentPostCall, @@ -46,6 +56,7 @@ pub struct LegacyLogging { body: Option>, headers: Option>, context: Option, + stream: Option, asynchronous: bool, internal: bool, pending: Option, @@ -80,6 +91,7 @@ impl LegacyLogging { body: None, headers: None, context: None, + stream: None, asynchronous, internal: false, pending: None, @@ -163,6 +175,49 @@ impl LegacyLogging { logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } + fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { + let logger = self.logger()?; + let billed = Streaming::Success.call( + py, + ( + logger.object(py), + &self.body, + &stream.chunks, + &self.start, + &self.end, + &stream.first_chunk, + ), + ); + match billed { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(logger.object(py))); + Ok(()) + } + result => result.map(|_| ()), + } + } + + /// A failure after the stream reached the caller bills the delivered chunks as + /// partial usage. The sync path has no loop to schedule that on, so it falls back to + /// the plain failure handler. + fn stream_failure(&mut self, py: Python<'_>) -> PyResult { + let (Some(logger), Some(error), Some(stream)) = (&self.logger, &self.error, &self.stream) + else { + return Ok(LifecycleStep::Done); + }; + if !self.asynchronous { + return self.dispatch_failure(py); + } + match Streaming::Failure.call(py, (logger.object(py), &self.body, &stream.chunks, error)) { + Ok(awaitable) => { + self.pending = Some(Pending::AsyncFailure); + Ok(LifecycleStep::Await(awaitable.unbind())) + } + Err(failure) if is_cancellation(py, &failure) => Err(failure), + Err(_) => Ok(LifecycleStep::Done), + } + } + /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. @@ -296,6 +351,22 @@ impl PythonLifecycle for LegacyLogging { ) -> PyResult { match (event, public) { (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), + (CallEvent::Opened, _) => { + Streaming::Opened.call(py, (self.logger()?.object(py),))?; + self.stream = Some(DeliveredStream { + chunks: PyList::empty(py).unbind(), + first_chunk: None, + }); + Ok(LifecycleStep::Done) + } + (CallEvent::Delivered, Some(PublicValue::Chunk(chunk))) => { + let stream = self.stream.as_mut().ok_or_else(missing_state)?; + if stream.first_chunk.is_none() { + stream.first_chunk = Some(datetime(py, epoch_seconds())?); + } + stream.chunks.bind(py).append(chunk)?; + Ok(LifecycleStep::Done) + } (CallEvent::ResponseReceived { raw }, _) => { let api_key = self .context @@ -314,12 +385,18 @@ impl PythonLifecycle for LegacyLogging { (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); - self.dispatch_success(py)?; + match &self.stream { + Some(stream) => self.stream_success(py, stream)?, + None => self.dispatch_success(py)?, + } Ok(LifecycleStep::Done) } (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); + if self.stream.is_some() { + return self.stream_failure(py); + } if *origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() @@ -366,6 +443,7 @@ impl PythonLifecycle for LegacyLogging { } self.body = None; self.context = None; + self.stream = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -377,6 +455,10 @@ impl PythonLifecycle for LegacyLogging { visit.call(&self.end)?; visit.call(&self.response)?; visit.call(&self.error)?; + if let Some(stream) = &self.stream { + visit.call(&stream.chunks)?; + visit.call(&stream.first_chunk)?; + } visit.call(&self.body) } } diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs index a924775070c..7f5c77c1735 100644 --- a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs +++ b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs @@ -16,6 +16,7 @@ pub(crate) enum LegacyPython { Wrapper(Wrapper), Logging(Logging), DeploymentHooks(DeploymentHooks), + Streaming(Streaming), } /// The `@client` wrapper around the call: `function_setup`, limits, credentials, @@ -76,12 +77,25 @@ pub(crate) enum DeploymentHooks { AfterDeploymentFailure, } +/// The Messages stream iterator's logging: the stream flag, the end-of-stream billing +/// from the delivered chunks, and the partial-usage failure path. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Streaming { + #[strum(serialize = "stream_opened")] + Opened, + #[strum(serialize = "stream_success")] + Success, + #[strum(serialize = "stream_failure")] + Failure, +} + impl LegacyPython { fn name(self) -> &'static str { match self { Self::Wrapper(function) => function.into(), Self::Logging(function) => function.into(), Self::DeploymentHooks(function) => function.into(), + Self::Streaming(function) => function.into(), } } @@ -111,6 +125,15 @@ impl Logging { } } +impl Streaming { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Streaming(self).call(py, args) + } +} + impl DeploymentHooks { pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> where @@ -126,7 +149,7 @@ mod tests { use strum::VariantArray; - use super::{DeploymentHooks, LegacyPython, Logging, Wrapper}; + use super::{DeploymentHooks, LegacyPython, Logging, Streaming, Wrapper}; use crate::test_support::PYTHON_CONTRACT; #[test] @@ -147,6 +170,11 @@ mod tests { .iter() .map(|&function| LegacyPython::DeploymentHooks(function)), ) + .chain( + Streaming::VARIANTS + .iter() + .map(|&function| LegacyPython::Streaming(function)), + ) .map(LegacyPython::name) .collect(); assert_eq!(called.len(), declared.len(), "a function is borrowed twice"); diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 444655ea77b..42ca184eb16 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -84,6 +84,11 @@ FAKES = { 'success', response, call_type ), 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), } assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) for name, fake in FAKES.items(): diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index 3bf12e553b9..d19e973b812 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -60,6 +60,10 @@ pub enum CallEvent { ResponseReceived { raw: RawResponse, }, + /// The call streams and its stream was handed to the caller. + Opened, + /// One chunk of an open stream reached the caller. + Delivered, Succeeded { timing: Timing, }, diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/callbacks/src/host.rs index 2392718a18d..eef3e1da8d5 100644 --- a/litellm-rust/crates/callbacks/src/host.rs +++ b/litellm-rust/crates/callbacks/src/host.rs @@ -11,12 +11,25 @@ pub enum HostOp { context: Box, }, Emit(CallEvent), + /// The response streams: the host hands the caller a stream and answers once the + /// caller asks for the first chunk or goes away. + Open(R::StreamHead), + /// The next chunk of an open stream, answered once the caller asks for the one after. + Deliver(R::Chunk), } pub enum HostResult { Route(R::OpResult), BeforeSend(Box), Emitted, + Demand(Demand), +} + +/// Whether the caller of a streamed call still reads it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Demand { + More, + Detached, } /// A host answer that is either available now or arrives once the host's own @@ -42,4 +55,12 @@ pub trait Host: Send + Sync { fn emit(&self, _event: &CallEvent) -> impl Future> + Send { async { Ok(()) } } + + fn open(&self, _head: R::StreamHead) -> impl Future> + Send { + async { Ok(Demand::More) } + } + + fn deliver(&self, _chunk: R::Chunk) -> impl Future> + Send { + async { Ok(Demand::More) } + } } diff --git a/litellm-rust/crates/callbacks/src/route.rs b/litellm-rust/crates/callbacks/src/route.rs index 97738c8da8b..8ab2b125760 100644 --- a/litellm-rust/crates/callbacks/src/route.rs +++ b/litellm-rust/crates/callbacks/src/route.rs @@ -6,4 +6,9 @@ pub trait Route: Send + Sync + 'static { type Error: Clone + Send + Sync + 'static; type Op: Send + 'static; type OpResult: Send + 'static; + /// One piece of a streamed response, handed to the caller as it arrives. A route + /// that never streams uses `Infallible`. + type Chunk: Send + 'static; + /// What the route knows once a streamed response starts, before its first chunk. + type StreamHead: Send + 'static; } diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 5accaa2e25e..705c5504c01 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -26,6 +26,8 @@ where .await .map(|wire| HostResult::BeforeSend(Box::new(wire))), HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted), + HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), + HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), }; match answer { Ok(answer) => result = Some(answer), @@ -61,6 +63,8 @@ mod tests { type Error = &'static str; type Op = &'static str; type OpResult = (); + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } struct Scripted { diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index 279a2d65c97..929f0a423c4 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -9,7 +9,7 @@ use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; use litellm_callbacks::{ event::{CallEvent, RequestContext, WireRequest}, - host::{HostOp, HostResult}, + host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, }; @@ -88,6 +88,21 @@ where _ => Err(MachineFault::Mismatch.into()), } } + + pub async fn open(&self, head: R::StreamHead) -> Result { + self.demand(HostOp::Open(head)).await + } + + pub async fn deliver(&self, chunk: R::Chunk) -> Result { + self.demand(HostOp::Deliver(chunk)).await + } + + async fn demand(&self, op: HostOp) -> Result { + match self.invoke(op).await? { + HostResult::Demand(demand) => Ok(demand), + _ => Err(MachineFault::Mismatch.into()), + } + } } enum Execution { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index b95402b1a7a..22e2c398ff7 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,88 +1,54 @@ -use litellm_llms::custom_httpx::http_handler::http_request; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use std::time::Duration; -use super::{ - Error, client::http_client, common_utils::truncate_error_body, - prepare::prepare_provider_request, +use litellm_llms::{ + base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, + custom_httpx::{http_handler::http_request, transport::Error as TransportError}, }; -use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, messages::types::MessagesRequest}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use serde_json::Value; -pub(super) async fn execute_messages_provider_call( - request: MessagesRequest<'_>, +use super::{Error, client::http_client, common_utils::truncate_error_body}; + +pub(super) fn network(error: reqwest::Error) -> Error { + Error::Transport(TransportError::Network(error.to_string())) +} + +pub(super) async fn send( + url: &str, + headers: &[(String, String)], + body: &Value, + timeout: Option, +) -> Result { + let builder = headers.iter().fold( + http_client().post(url).json(body), + |builder, (key, value)| builder.header(key, value), + ); + let builder = match timeout { + Some(duration) => builder.timeout(duration), + None => builder, + }; + http_request(builder).await.map_err(network) +} + +pub(super) async fn provider_error(response: reqwest::Response) -> Error { + let status = response.status().as_u16(); + match response.text().await { + Ok(text) => Error::Transport(TransportError::Http { + status, + body: truncate_error_body(&text), + }), + Err(error) => network(error), + } +} + +pub(super) fn decode_response( + config: &dyn BaseAnthropicMessagesConfig, + model: &str, + text: &str, ) -> Result { - let request = prepare_provider_request(request)?; - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - - let status = response.status(); - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - - if !status.is_success() { - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); - } - - let response = serde_json::from_str(&text) + let response = serde_json::from_str(text) .map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?; - request - .config - .transform_anthropic_messages_response(&request.model, response) + config + .transform_anthropic_messages_response(model, response) .map_err(Error::from) } - -pub(super) async fn execute_messages_provider_stream( - request: MessagesRequest<'_>, -) -> Result { - let request = prepare_provider_request(request)?; - if request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(Error::Unsupported("streaming messages for this provider")); - } - - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - let status = response.status(); - if !status.is_success() { - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); - } - Ok(response) -} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index c3d7bea48ff..e36c6668efe 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -1,11 +1,8 @@ //! The Anthropic Messages call, the Rust equivalent of Python's //! `litellm.messages()`. //! -//! [`messages`] is the top-level entrypoint: give it a model, a body, and -//! credentials, and it resolves the provider, transforms the request, calls the -//! provider, and returns a typed non-streaming response. [`messages_stream`] -//! is the streaming variant; it hands the raw upstream response back so a host -//! can splice the event stream to its own caller. +//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs +//! it in process for a caller that already holds the request and wants the message. mod error; pub mod types; @@ -14,17 +11,34 @@ mod client; mod common_utils; mod handler; mod prepare; -use handler::{execute_messages_provider_call, execute_messages_provider_stream}; +pub mod route; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; +use serde_json::Value; use crate::messages::types::MessagesRequest; pub async fn messages(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_call(request).await -} - -pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_stream(request).await + let Value::Object(body) = request.body else { + return Err(Error::InvalidRequest( + "messages body must be an object".into(), + )); + }; + let call = MessagesCall { + model: request.model.into(), + body, + api_key: request.api_key.map(Into::into), + api_base: request.api_base.map(Into::into), + custom_llm_provider: request.custom_llm_provider.map(Into::into), + extra_headers: request.extra_headers, + timeout: request.timeout, + }; + match litellm_callbacks::run::run(messages_machine(), &LocalMessagesHost::new(call)).await? { + MessagesOutput::Message(message) => Ok(message), + MessagesOutput::Streamed => Err(Error::Unsupported( + "streamed responses need a streaming host", + )), + } } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 8b676803871..850f9108869 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -2,6 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_l use litellm_llms::base_llm::anthropic_messages::transformation::{ BaseAnthropicMessagesConfig, MessagesAuthStrategy, }; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde_json::{Map, Value}; use super::{ @@ -37,10 +38,14 @@ pub(super) fn prepare_provider_request( let headers = validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?; - let typed_request = serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + let typed_request: AnthropicMessagesRequest = + serde_json::from_value(request.body).map_err(|err| { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = config.transform_anthropic_messages_request(AnthropicMessagesRequest { + model: model.clone(), + ..typed_request })?; - let transformed = config.transform_anthropic_messages_request(typed_request)?; let body = serde_json::to_value(transformed).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs new file mode 100644 index 00000000000..a680b5ce1ec --- /dev/null +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -0,0 +1,197 @@ +use std::{sync::Mutex, time::Duration}; + +use bytes::Bytes; +use litellm_auth::SecretValue; +use litellm_callbacks::{ + event::{CallEvent, RawResponse, RequestContext, WireRequest}, + host::{Demand, Host}, + route::Route, +}; +use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use serde_json::{Map, Value}; + +use super::{ + Error, + common_utils::messages_provider_config, + handler::{decode_response, network, provider_error, send}, + prepare::prepare_provider_request, + types::MessagesRequest, +}; +use crate::{ + constants::ANTHROPIC_MESSAGES_PROVIDER, + machine::{HostChannel, MachineFault, RouteMachine}, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MessagesOp { + ProjectRequest, +} + +pub enum MessagesOpResult { + Request(Box), +} + +/// The caller's request as the host projects it. +pub struct MessagesCall { + pub model: String, + pub body: Map, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub timeout: Option, +} + +impl MessagesCall { + fn streams(&self) -> bool { + self.body.get("stream").and_then(Value::as_bool) == Some(true) + } +} + +pub enum MessagesOutput { + Message(AnthropicMessagesResponse), + /// Every chunk already reached the host through `Deliver`. + Streamed, +} + +pub struct Messages; + +impl Route for Messages { + type Response = MessagesOutput; + type Error = Error; + type Op = MessagesOp; + type OpResult = MessagesOpResult; + type Chunk = Bytes; + type StreamHead = (); +} + +impl From for Error { + fn from(fault: MachineFault) -> Self { + Self::InvalidRequest(match fault { + MachineFault::Abandoned => "messages host driver was abandoned".into(), + MachineFault::Protocol(message) => format!("messages {message}"), + MachineFault::Mismatch => "invalid messages host operation result".into(), + }) + } +} + +pub type MessagesHost = HostChannel; +pub type MessagesMachine = RouteMachine; + +/// Whether this route serves the request, decided before any callback runs so a host +/// can still run its own path. +pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool { + let provider = get_custom_llm_provider(model, custom_llm_provider) + .map(|resolved| resolved.custom_llm_provider) + .or(custom_llm_provider); + match provider { + Some(ANTHROPIC_MESSAGES_PROVIDER) => true, + Some(provider) => !stream && messages_provider_config(provider).is_some(), + None => false, + } +} + +/// The in-process host for a request already in hand. It answers projection once and +/// observes nothing. +pub struct LocalMessagesHost { + call: Mutex>, +} + +impl LocalMessagesHost { + pub fn new(call: MessagesCall) -> Self { + Self { + call: Mutex::new(Some(call)), + } + } +} + +impl Host for LocalMessagesHost { + async fn route(&self, op: MessagesOp) -> Result { + match op { + MessagesOp::ProjectRequest => self + .call + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .map(|call| MessagesOpResult::Request(Box::new(call))) + .ok_or_else(|| { + Error::InvalidRequest("messages request was already projected".into()) + }), + } + } +} + +pub fn messages_machine() -> MessagesMachine { + RouteMachine::new(|host| Box::pin(execute(host))) +} + +async fn execute(host: MessagesHost) -> Result { + let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?; + let stream = call.streams(); + let request = prepare_provider_request(MessagesRequest { + model: &call.model, + body: Value::Object(call.body.clone()), + api_key: call.api_key.as_deref(), + api_base: call.api_base.as_deref(), + custom_llm_provider: call.custom_llm_provider.as_deref(), + extra_headers: call.extra_headers.clone(), + timeout: call.timeout, + })?; + if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER { + return Err(Error::Unsupported("streaming messages for this provider")); + } + let context = RequestContext { + model: request.model.clone(), + custom_llm_provider: request.provider.clone(), + optional_params: Value::Object( + call.body + .iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ), + secret_fields: Vec::new(), + api_key: call.api_key.clone().map(SecretValue::new), + }; + let wire = host + .before_send( + WireRequest { + url: request.url, + headers: request.upstream_headers, + body: request.body, + }, + context, + ) + .await?; + let response = send(&wire.url, &wire.headers, &wire.body, request.timeout).await?; + if !response.status().is_success() { + return Err(provider_error(response).await); + } + if stream { + return relay(&host, response).await; + } + let text = response.text().await.map_err(network)?; + host.emit(CallEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; + decode_response(request.config, &request.model, &text).map(MessagesOutput::Message) +} + +/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading +/// ends the upstream read, and the call completes with what it delivered. +async fn relay( + host: &MessagesHost, + mut response: reqwest::Response, +) -> Result { + if host.open(()).await? == Demand::Detached { + return Ok(MessagesOutput::Streamed); + } + while let Some(chunk) = response.chunk().await.map_err(network)? { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(MessagesOutput::Streamed) +} diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index ac4237651da..e6ee45c64a8 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -39,6 +39,8 @@ impl Route for Ocr { type Error = Error; type Op = OcrOp; type OpResult = OcrOpResult; + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } impl TokenRoute for Ocr { diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index ca2e14a7f0d..779d2037bb3 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -198,6 +198,8 @@ fn event_name(event: &CallEvent) -> &'static str { CallEvent::ResponseReceived { .. } => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", + CallEvent::Opened => "opened", + CallEvent::Delivered => "delivered", } } diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 4aa7a2163ca..795977438fd 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -23,6 +23,7 @@ pub enum LifecycleStep { pub enum PublicValue<'a> { Response(&'a Py), Error(&'a PyErr), + Chunk(&'a Py), } /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in @@ -111,6 +112,13 @@ pub trait RouteHost: Send + Sync { response: ::Response, ) -> PyResult>; + /// One streamed chunk as the caller receives it. + fn chunk( + &mut self, + py: Python<'_>, + chunk: ::Chunk, + ) -> PyResult>; + fn classify( &self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 7895d803cb5..c59ba1925dc 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -3,7 +3,7 @@ use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; -use litellm_callbacks::host::{HostOp, HostResult, HostStep}; +use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; use litellm_callbacks::route::Route; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; @@ -38,6 +38,7 @@ struct MachineState { enum Stage { Begin, Call, + Streaming, AfterSuccess, Succeeded(Py), Failed(Py), @@ -56,6 +57,8 @@ enum Expect { enum Pending { Native, Adapter(Expect), + /// The stream handed to the caller waits for its next read or its close. + Consumer, } enum Next { @@ -121,7 +124,14 @@ where } match driver.resume(None)? { ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Await(_) => Err(PyRuntimeError::new_err("sync call suspended")), + ExecutionStep::Open => py + .import("litellm.rust_bridge.lifecycle")? + .getattr("SyncStream")? + .call1((Py::new(py, Execution::suspended(driver))?,)) + .map(Bound::unbind), + ExecutionStep::Await(_) | ExecutionStep::Yield(_) => { + Err(PyRuntimeError::new_err("sync call suspended")) + } } } @@ -162,6 +172,14 @@ where self.run_steps(py, HostStep::Ready(result)) } (Some(Pending::Native), Some(Err(error))) => self.interrupt(py, error), + (Some(Pending::Consumer), Some(read)) => { + let demand = if read.is_ok() { + Demand::More + } else { + Demand::Detached + }; + self.resume_machine(py, Some(Ok(HostResult::Demand(demand)))) + } (Some(Pending::Adapter(expect)), Some(result)) => { match self.adapter.resume(py, result) { Ok(step) => self.on_adapter(py, step, expect), @@ -216,7 +234,7 @@ where fn adapter_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { match self.stage { Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), - Stage::Call => self.interrupt(py, error), + Stage::Call | Stage::Streaming => self.interrupt(py, error), Stage::Succeeded(_) | Stage::Failed(_) => Err(error), } } @@ -283,6 +301,8 @@ where Err(error) => Err(error), } } + HostOp::Open(_) => return self.opened(py).map(Next::Return), + HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), Ok(LifecycleStep::Await(awaitable)) => { @@ -299,6 +319,40 @@ where } } + fn opened(&mut self, py: Python<'_>) -> PyResult { + self.stage = Stage::Streaming; + match self.adapter.emit(py, &CallEvent::Opened, None) { + Ok(LifecycleStep::Done) => { + self.pending = Some(Pending::Consumer); + Ok(ExecutionStep::Open) + } + Ok(_) => Err(missing_state()), + Err(error) => self.interrupt(py, error), + } + } + + fn delivered( + &mut self, + py: Python<'_>, + chunk: as Route>::Chunk, + ) -> PyResult { + let chunk = match self.route.chunk(py, chunk) { + Ok(chunk) => chunk, + Err(error) => return self.interrupt(py, error), + }; + let observed = + self.adapter + .emit(py, &CallEvent::Delivered, Some(PublicValue::Chunk(&chunk))); + match observed { + Ok(LifecycleStep::Done) => { + self.pending = Some(Pending::Consumer); + Ok(ExecutionStep::Yield(chunk)) + } + Ok(_) => Err(missing_state()), + Err(error) => self.interrupt(py, error), + } + } + fn interrupt(&mut self, py: Python<'_>, error: PyErr) -> PyResult { let cancelled = is_cancellation(py, &error); let native = H::host_error(&error); @@ -369,6 +423,9 @@ where Ok(public) => public, Err(error) => return self.failure(py, error, FailureOrigin::Call), }; + if let Stage::Streaming = self.stage { + return self.succeeded(py, public); + } self.stage = Stage::AfterSuccess; match self.adapter.after_success(py, public, self.timing()) { Ok(step) => self.on_adapter(py, step, Expect::Response), @@ -536,6 +593,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri type Error = Error; type Op = &'static str; type OpResult = String; + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } /// Yields the scripted ops in order, then completes or fails as scripted. @@ -574,6 +633,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri HostResult::Route(value) => value, HostResult::BeforeSend(wire) => wire.url, HostResult::Emitted => "emitted".into(), + HostResult::Demand(demand) => format!("{demand:?}"), }); } if !self.ops.is_empty() { @@ -648,6 +708,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { + match chunk {} + } + fn complete(&mut self, py: Python<'_>, response: String) -> PyResult> { self.log.push("complete"); Ok(pyo3::types::PyString::new(py, &response) @@ -1138,6 +1202,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) .into()) } + fn chunk( + &mut self, + _: Python<'_>, + chunk: std::convert::Infallible, + ) -> PyResult> { + match chunk {} + } fn complete(&mut self, _: Python<'_>, _: String) -> PyResult> { Err(missing_state()) } diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index d8cd6c92130..10abbadbda5 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -8,6 +8,10 @@ use pyo3::prelude::*; pub enum ExecutionStep { Return(Py), Await(Py), + /// The call streams: the caller gets a stream over this execution, which stays + /// suspended until the stream asks for a chunk. + Open, + Yield(Py), } pub trait ExecutionBody: Send + Sync { @@ -34,6 +38,13 @@ impl Execution { } } + /// An execution already started elsewhere and now waiting for its next input. + pub fn suspended(body: impl ExecutionBody + 'static) -> Self { + Self { + state: ExecutionState::Suspended(Box::new(body)), + } + } + fn advance( slf: &Bound<'_, Self>, py: Python<'_>, @@ -64,6 +75,8 @@ impl Execution { let step = body.resume(result)?; let (tag, value, suspended) = match step { ExecutionStep::Await(value) => ("Await", value, true), + ExecutionStep::Open => ("Open", py.None(), true), + ExecutionStep::Yield(value) => ("Yield", value, true), ExecutionStep::Return(value) => ("Complete", value, false), }; let step = py diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index ea4077b102f..2aba51cc4ff 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -18,10 +18,6 @@ pub(crate) struct RouteOptions { pub(crate) timeout: Option, } -pub(crate) fn body_argument(value: &Bound<'_, PyAny>) -> PyResult> { - required_object("body", from_py_argument(value)?) -} - pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult> { match from_py_argument(value)? { Value::Array(values) => Ok(values), @@ -192,18 +188,6 @@ mod tests { json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) ); - let body = py - .eval( - c"{'model': 'claude', 'metadata': {'user': '1'}}", - None, - None, - ) - .unwrap(); - assert_eq!( - Value::Object(body_argument(&body).unwrap()), - json!({"model": "claude", "metadata": {"user": "1"}}) - ); - let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap(); assert_eq!( optional_params_argument(¶ms).unwrap(), diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs deleted file mode 100644 index daec931c92e..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ /dev/null @@ -1,88 +0,0 @@ -use litellm_core::messages::{Error, messages as run_messages, types::MessagesRequest}; -use litellm_host_python::{run_async, run_sync}; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use pyo3::prelude::*; -use serde_json::{Map, Value}; - -use crate::{ - errors::messages_error_to_pyerr, - marshal::{RouteOptions, body_argument, extra_headers_argument, optional_timeout}, -}; - -async fn execute( - body: Map, - options: RouteOptions, -) -> Result { - let RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout, - } = options; - run_messages(MessagesRequest { - model: &model, - body: Value::Object(body), - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[expect( - clippy::too_many_arguments, - reason = "one parameter per Python keyword" -)] -pub(crate) fn messages( - py: Python<'_>, - model: String, - #[pyo3(from_py_with = body_argument)] body: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let options = RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout: optional_timeout(timeout_seconds), - }; - run_sync(py, execute(body, options), messages_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[expect( - clippy::too_many_arguments, - reason = "one parameter per Python keyword" -)] -pub(crate) fn amessages<'py>( - py: Python<'py>, - model: String, - #[pyo3(from_py_with = body_argument)] body: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let options = RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout: optional_timeout(timeout_seconds), - }; - run_async(py, execute(body, options), messages_error_to_pyerr) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs new file mode 100644 index 00000000000..0c87aea1ca2 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -0,0 +1,186 @@ +use bytes::Bytes; +use litellm_core::messages::{ + Error, + route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, +}; +use litellm_host_python::{HostOpError, RouteHost, from_py, lookup, to_py}; +use litellm_llms::custom_httpx::transport::Error as TransportError; +use pyo3::{ + exceptions::{PyException, PyValueError}, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyBytes, PyDict}, +}; +use serde_json::{Map, Value}; + +use crate::{ + errors::{RustUpstreamError, messages_error_to_pyerr}, + marshal::{optional_timeout, python_timeout_seconds}, +}; + +/// The Anthropic Messages body fields a caller may pass besides `model` and `messages`, +/// as `AnthropicMessagesRequestOptionalParams` declares them. +const BODY_FIELDS: [&str; 20] = [ + "max_tokens", + "metadata", + "stop_sequences", + "stream", + "system", + "temperature", + "thinking", + "tool_choice", + "tools", + "top_k", + "inference_geo", + "top_p", + "mcp_servers", + "context_management", + "container", + "output_format", + "speed", + "output_config", + "cache_control", + "reasoning_effort", +]; + +/// The Python side of the Messages route: projects the prepared arguments and builds the +/// public response, chunks and exceptions. +pub(super) struct MessagesRouteHost { + request: Py, +} + +impl MessagesRouteHost { + pub(super) fn new(request: Py) -> Self { + Self { request } + } + + fn project(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + let request = self.request.bind(py); + let argument = |name: &str| -> PyResult>> { + Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none())) + }; + let string = |name: &str| -> PyResult> { + argument(name)?.map(|value| value.extract()).transpose() + }; + let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; + let messages = + argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?; + let fields = BODY_FIELDS + .iter() + .filter_map(|name| match argument(name) { + Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect::>>()?; + let body = [ + ("model".to_string(), Value::String(model.clone())), + ("messages".to_string(), from_py(&messages)?), + ] + .into_iter() + .chain(fields) + .collect::>(); + let timeout = argument("timeout")? + .map(|value| python_timeout_seconds(py, value.unbind())) + .transpose()? + .flatten(); + Ok(MessagesCall { + model, + body, + api_key: string("api_key")?, + api_base: string("api_base")?, + custom_llm_provider: string("custom_llm_provider")?, + extra_headers: argument("extra_headers")? + .map(|value| from_py(&value)) + .transpose()?, + timeout: optional_timeout(timeout), + }) + } + + fn provider(&self, py: Python<'_>) -> String { + self.request + .bind(py) + .getattr("custom_llm_provider") + .and_then(|value| value.extract::>()) + .ok() + .flatten() + .unwrap_or_else(|| "anthropic".into()) + } + + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let mapped = py + .import("litellm.rust_bridge.messages.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py)))) + .and_then(|mapped| { + mapped + .extract::>() + .map_err(PyErr::from) + }); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for MessagesRouteHost { + type Route = Messages; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: MessagesOp, + ) -> Result> { + match op { + MessagesOp::ProjectRequest => self + .project(py, arguments) + .map(|call| MessagesOpResult::Request(Box::new(call))) + .map_err(|error| HostOpError::Python(self.map_failure(py, error))), + } + } + + fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult> { + match response { + MessagesOutput::Message(message) => py + .import("litellm.rust_bridge.messages.route_host")? + .getattr("response")? + .call1((to_py(py, &message)?,)) + .map(Bound::unbind), + MessagesOutput::Streamed => Ok(py.None()), + } + } + + fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { + Ok(PyBytes::new(py, &chunk).into_any().unbind()) + } + + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + let native = match error { + Error::Transport(TransportError::Http { status, body }) => { + let error = RustUpstreamError::new_err((status, body)); + error + .value(py) + .setattr("headers", Vec::<(String, String)>::new())?; + error + } + other => messages_error_to_pyerr(other), + }; + Ok(self.map_failure(py, native)) + } + + fn host_error(error: &PyErr) -> Error { + Error::InvalidRequest(error.to_string()) + } + + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.request) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs new file mode 100644 index 00000000000..b4259aca5ba --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -0,0 +1,64 @@ +mod host; + +use host::MessagesRouteHost; +use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; +use litellm_core::messages::route::{messages_machine, supports}; +use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, +}; + +use crate::errors::RustBridgeDeclined; + +const SURFACE: LegacySurface = LegacySurface { + call_type: "anthropic_messages", + input_description: "Messages", +}; + +fn run_messages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult> { + let model: String = request.getattr("model")?.extract()?; + let provider: Option = request.getattr("custom_llm_provider")?.extract()?; + let stream = request + .getattr("stream")? + .extract::>()? + .unwrap_or(false); + if !supports(&model, provider.as_deref(), stream) { + return Err(RustBridgeDeclined::new_err( + "the Rust Messages route does not serve this provider", + )); + } + run_legacy_call( + py, + SURFACE, + PublicCall::capture(&request, &args, &kwargs)?, + messages_machine(), + MessagesRouteHost::new(request.unbind()), + asynchronous, + ) +} + +#[pyfunction] +pub(crate) fn messages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + run_messages(py, request, args, kwargs, false) +} + +#[pyfunction] +pub(crate) fn amessages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + run_messages(py, request, args, kwargs, true) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f59e32a28e2..2d6b849a6b1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -22,11 +22,6 @@ mod tests { "atranscription", "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", ), - ( - "messages", - "amessages", - "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)", - ), ( "chat_completions", "achat_completions", @@ -113,25 +108,6 @@ value = Broken() ); assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string()); - let invalid_body = PyList::empty(py); - let sync_messages_error = module - .getattr("messages") - .and_then(|function| function.call1(("model", &invalid_body))) - .expect_err("sync Messages should reject a non-dict body"); - let async_messages_error = module - .getattr("amessages") - .and_then(|function| function.call1(("model", &invalid_body))) - .expect_err("async Messages should reject a non-dict body"); - - assert_eq!( - sync_messages_error.to_string(), - "ValueError: body must be a dict" - ); - assert_eq!( - async_messages_error.to_string(), - sync_messages_error.to_string() - ); - let invalid_headers = PyList::empty(py); let kwargs = PyDict::new(py); kwargs @@ -193,13 +169,6 @@ value = Broken() headers_kwargs .set_item("extra_headers", &invalid) .expect("kwargs should accept extra_headers"); - let invalid_body = PyList::empty(py); - let error = module - .getattr("messages") - .and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs))) - .expect_err("body should be validated before headers"); - assert_eq!(error.to_string(), "ValueError: body must be a dict"); - let invalid_payload = PyModule::new(py, "invalid_payload").expect("invalid payload should be created"); let error = module diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 77212e3d38e..19211671fd7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -125,6 +125,10 @@ impl RouteHost for OcrRouteHost { .map(Bound::unbind) } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { + match chunk {} + } + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 488e278cca7..9f959c056de 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,9 +1,11 @@ from asyncio import Future -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... @@ -39,23 +41,15 @@ def atranscription( timeout_seconds: float | None = None, ) -> Future[dict[str, object]]: ... def messages( - model: str, - body: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - custom_llm_provider: str | None = None, - extra_headers: Mapping[str, object] | None = None, - timeout_seconds: float | None = None, -) -> dict[str, object]: ... + request: LiteLLMMessagesRequest, + args: tuple[object, ...], + kwargs: dict[str, object], +) -> AnthropicMessagesResponse | Iterator[bytes]: ... def amessages( - model: str, - body: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - custom_llm_provider: str | None = None, - extra_headers: Mapping[str, object] | None = None, - timeout_seconds: float | None = None, -) -> Future[dict[str, object]]: ... + request: LiteLLMMessagesRequest, + args: tuple[object, ...], + kwargs: dict[str, object], +) -> Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]]: ... def chat_completions_decline( model: str, messages: Sequence[object], diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9efbbfa2e9e..d843a874fe3 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -59,6 +59,7 @@ Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( Rule(Route.OCR, Rollout.RUST_OPT_OUT), + Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/rust_bridge/failures.py b/litellm/rust_bridge/failures.py index b714341fe43..80805b7ff69 100644 --- a/litellm/rust_bridge/failures.py +++ b/litellm/rust_bridge/failures.py @@ -5,8 +5,37 @@ from __future__ import annotations from collections.abc import Mapping from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper +import httpx +import openai +from pydantic import TypeAdapter, ValidationError + import litellm +_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) +_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) + + +class UpstreamFailure(Exception): + def __init__(self, response: httpx.Response, cause: Exception) -> None: + super().__init__(str(cause)) + self.message: Final = str(cause) + self.response: Final = response + self.status_code: Final = response.status_code + self.__cause__ = cause + + +def _upstream_failure(error: Exception, api_base: str | None) -> Exception: + try: + status, body = _UPSTREAM_ARGS.validate_python(error.args) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + except ValidationError: + return error + http_request: Final = httpx.Request("POST", api_base or "https://docs.litellm.ai/docs") + return UpstreamFailure( + httpx.Response(status, content=body.encode(), headers=headers, request=http_request), + error, + ) + class ExceptionMapper(Protocol): def __call__( @@ -35,3 +64,17 @@ def map_failure(error: Exception, model: str, request_provider: str, kwargs: Map except Exception as public_error: public_error.__context__ = error return public_error + + +def map_native_failure( + error: Exception, model: str, request_provider: str, kwargs: Mapping[str, object], api_base: str | None = None +) -> Exception: + """`map_failure`, reading a native `(status, body)` provider failure as the HTTP response it was.""" + original: Final = _upstream_failure(error, api_base) + public_error: Final = map_failure(original, model, request_provider, kwargs) + if isinstance(original, UpstreamFailure) and public_error.__context__ is original: + public_error.__context__ = error + if isinstance(public_error, openai.APIStatusError): + public_error.response = original.response + public_error.status_code = original.status_code + return public_error diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index 406dc55cfea..fd0fac1799a 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -6,6 +6,7 @@ registries it fans out to. It expires with that contract. from __future__ import annotations +import asyncio import contextvars import datetime import traceback @@ -160,6 +161,21 @@ class LoggingWorker(Protocol): def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... +class StreamingLogBuilder(Protocol): + def __call__( + self, + *, + litellm_logging_obj: Logging, + passthrough_success_handler_obj: object, + url_route: str, + request_body: dict[str, object], + endpoint_type: object, + start_time: datetime.datetime, + raw_bytes: list[bytes], + end_time: datetime.datetime, + ) -> Coroutine[object, object, None]: ... + + class DeploymentHook(Protocol): def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... @@ -306,3 +322,67 @@ def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_t DeploymentFailureHook, utils.async_post_call_failure_deployment_hook ) return hook(kwargs, error, call_type) + + +def stream_opened(logger: Logging) -> None: + logger.stream = True + logger.model_call_details["stream"] = True + + +def stream_success( + logger: Logging, + request_body: dict[str, object], + chunks: list[bytes], + start: datetime.datetime, + end: datetime.datetime, + first_chunk: datetime.datetime | None, +) -> None: + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + if first_chunk is not None: + logger.completion_start_time = first_chunk + logger.model_call_details["completion_start_time"] = first_chunk + build: Final = cast( # cast-ok: bounded adapter for the untyped pass-through logging builder + StreamingLogBuilder, + PassThroughStreamingHandler._route_streaming_logging_to_handler, # pyright: ignore[reportPrivateUsage] # the Messages stream iterator bills through the same builder + ) + coroutine: Final = build( + litellm_logging_obj=logger, + passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + url_route="/v1/messages", + request_body=request_body, + endpoint_type=EndpointType.ANTHROPIC, + start_time=start, + raw_bytes=chunks, + end_time=end, + ) + if getattr(logger, "_on_deferred_stream_complete", None) is not None: + logger._deferred_stream_complete_args = (coroutine,) # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release reads this slot + return + try: + asyncio.get_running_loop() + except RuntimeError: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, asyncio.run, coroutine) + return + enqueue_logging(coroutine) + + +def stream_failure( + logger: Logging, request_body: dict[str, object], chunks: list[bytes], error: Exception +) -> Coroutine[object, object, None]: + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + return PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=logger, + endpoint_type=EndpointType.ANTHROPIC, + request_body=request_body, + raw_bytes=chunks, + exception=error, + ) diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index d903021b6f3..4096d386964 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,8 +1,8 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import AsyncIterator, Awaitable, Iterator from dataclasses import dataclass -from typing import Protocol +from typing import Final, Protocol @dataclass(frozen=True, slots=True) @@ -15,28 +15,133 @@ class Complete: value: object +@dataclass(frozen=True, slots=True) +class Open: + value: None + + +@dataclass(frozen=True, slots=True) +class Yield: + value: object + + +Settled = Complete | Open | Yield +Step = Await | Settled + + class Execution(Protocol): - def start(self) -> Await | Complete: ... + def start(self) -> Step: ... - def resume_value(self, value: object) -> Await | Complete: ... + def resume_value(self, value: object) -> Step: ... - def resume_error(self, error: BaseException) -> Await | Complete: ... + def resume_error(self, error: BaseException) -> Step: ... def close(self) -> None: ... +class StreamClosed(Exception): + """Tells a streaming execution that its caller stopped reading.""" + + +async def _settle(execution: Execution, step: Step) -> Settled: + while isinstance(step, Await): + try: + value = await step.awaitable # rebind-ok: each selected await produces the next protocol input + except GeneratorExit: + raise + except BaseException as error: + step = execution.resume_error(error) # rebind-ok: advance the execution protocol + else: + step = execution.resume_value(value) # rebind-ok: advance the execution protocol + return step + + +def _settled(step: Step) -> Settled: + if isinstance(step, Await): + raise RuntimeError("sync call suspended") + return step + + async def drive(execution: Execution) -> object: + handed_off = False # rebind-ok: set once the execution belongs to the returned stream try: - step = execution.start() # rebind-ok: the execution protocol advances after each selected await - while isinstance(step, Await): - try: - value = await step.awaitable # rebind-ok: each selected await produces the next protocol input - except GeneratorExit: - raise - except BaseException as error: - step = execution.resume_error(error) # rebind-ok: advance the execution protocol - else: - step = execution.resume_value(value) # rebind-ok: advance the execution protocol + step: Final = await _settle(execution, execution.start()) + if isinstance(step, Open): + handed_off = True + return Stream(execution) return step.value finally: - execution.close() + if not handed_off: + execution.close() + + +class Stream(AsyncIterator[object]): + """A streamed native call: each read resumes the execution until its next chunk.""" + + def __init__(self, execution: Execution) -> None: + self._execution: Final = execution + self._done = False + + def __aiter__(self) -> Stream: + return self + + async def __anext__(self) -> object: + if self._done: + raise StopAsyncIteration + try: + step: Final = await _settle(self._execution, self._execution.resume_value(None)) + except BaseException: + self._finish() + raise + if isinstance(step, Yield): + return step.value + self._finish() + raise StopAsyncIteration + + async def aclose(self) -> None: + if self._done: + return + try: + await _settle(self._execution, self._execution.resume_error(StreamClosed())) + finally: + self._finish() + + def _finish(self) -> None: + self._done = True + self._execution.close() + + +class SyncStream(Iterator[object]): + """The sync form of `Stream`; its execution never suspends on an awaitable.""" + + def __init__(self, execution: Execution) -> None: + self._execution: Final = execution + self._done = False + + def __iter__(self) -> SyncStream: + return self + + def __next__(self) -> object: + if self._done: + raise StopIteration + try: + step: Final = _settled(self._execution.resume_value(None)) + except BaseException: + self._finish() + raise + if isinstance(step, Yield): + return step.value + self._finish() + raise StopIteration + + def close(self) -> None: + if self._done: + return + try: + _settled(self._execution.resume_error(StreamClosed())) + finally: + self._finish() + + def _finish(self) -> None: + self._done = True + self._execution.close() diff --git a/litellm/rust_bridge/messages/entrypoints.py b/litellm/rust_bridge/messages/entrypoints.py index 46565bfd46a..d25c906c4c1 100644 --- a/litellm/rust_bridge/messages/entrypoints.py +++ b/litellm/rust_bridge/messages/entrypoints.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping, Sequence from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables @@ -26,7 +26,7 @@ class NativeMessages(Protocol): request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object], - ) -> AnthropicMessagesResponse: ... + ) -> AnthropicMessagesResponse | Iterator[bytes]: ... class NativeAmessages(Protocol): @@ -35,7 +35,7 @@ class NativeAmessages(Protocol): request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object], - ) -> Awaitable[AnthropicMessagesResponse]: ... + ) -> Awaitable[AnthropicMessagesResponse | AsyncIterator[bytes]]: ... def _messages_binding(value: object) -> NativeMessages | None: @@ -50,5 +50,5 @@ def _amessages_binding(value: object) -> NativeAmessages | None: return cast("NativeAmessages", value) # cast-ok: callable validated at the native binding boundary -NATIVE_MESSAGES: Final = NativeBinding("anthropic_messages_handler", validate=_messages_binding) -NATIVE_AMESSAGES: Final = NativeBinding("anthropic_messages", validate=_amessages_binding) +NATIVE_MESSAGES: Final = NativeBinding("messages", validate=_messages_binding) +NATIVE_AMESSAGES: Final = NativeBinding("amessages", validate=_amessages_binding) diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index 1aff6c7f75d..beef0f81eca 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -20,4 +20,4 @@ def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]: def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception: - return failures.map_failure(error, request.model, request_provider, arguments(request)) + return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) diff --git a/litellm/rust_bridge/ocr/route_host.py b/litellm/rust_bridge/ocr/route_host.py index 277fdceb734..bfbd5c11d4e 100644 --- a/litellm/rust_bridge/ocr/route_host.py +++ b/litellm/rust_bridge/ocr/route_host.py @@ -4,40 +4,17 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -import httpx -import openai -from pydantic import TypeAdapter, ValidationError +from pydantic import TypeAdapter import litellm from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse from litellm.rust_bridge import failures +from litellm.rust_bridge.failures import UpstreamFailure from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest +__all__ = ("UpstreamFailure", "arguments", "map_failure", "response") + _RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object]) -_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) -_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) - - -class UpstreamFailure(Exception): - def __init__(self, response: httpx.Response, cause: Exception) -> None: - super().__init__(str(cause)) - self.message: Final = str(cause) - self.response: Final = response - self.status_code: Final = response.status_code - self.__cause__ = cause - - -def _upstream_failure(error: Exception, request: LiteLLMOcrRequest) -> Exception: - try: - status, body = _UPSTREAM_ARGS.validate_python(error.args) - headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) - except ValidationError: - return error - http_request: Final = httpx.Request("POST", request.api_base or "https://docs.litellm.ai/docs") - return UpstreamFailure( - httpx.Response(status, content=body.encode(), headers=headers, request=http_request), - error, - ) def response(value: Mapping[str, object]) -> OCRResponse: @@ -61,11 +38,4 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: model=request.model.removeprefix(f"{request_provider}/"), llm_provider=request_provider, ) - original: Final = _upstream_failure(error, request) - public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request)) - if isinstance(original, UpstreamFailure) and public_error.__context__ is original: - public_error.__context__ = error - if isinstance(public_error, openai.APIStatusError): - public_error.response = original.response - public_error.status_code = original.status_code - return public_error + return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 4fa4c0b95ec..0b442f1f269 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -73,7 +73,7 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"transcription", "messages", "chat_completions"}: + if route not in {"transcription", "chat_completions"}: raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") @@ -89,10 +89,6 @@ def assert_native_request( assert path == "/v1/messages" assert headers.get("x-api-key") == "sk-native" assert body["model"] == "claude-sonnet-4-5" - if route == "messages": - assert body["max_tokens"] == 16 - assert body["messages"][0]["content"] == "hello-from-messages" - return assert body["max_tokens"] == 17 assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] @@ -132,17 +128,6 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: "language": "en", }, } - if route == "messages": - return common | { - "model": "claude-sonnet-4-5", - "body": { - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hello-from-messages"}], - }, - "api_key": "sk-native", - "custom_llm_provider": "anthropic", - } if route == "chat_completions": return common | { "model": "anthropic/claude-sonnet-4-5", @@ -165,8 +150,6 @@ def assert_success(route: str, response: object) -> None: def success_value(route: str, response: dict[object, object]) -> object: if route == "transcription": return response["text"] - if route == "messages": - return response["content"][0]["text"] return response["choices"][0]["message"]["content"] @@ -181,7 +164,7 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None: def exercise_sync(native: object, api_base: str) -> None: - for route in ("transcription", "messages", "chat_completions"): + for route in ("transcription", "chat_completions"): function: Final = getattr(native, route) assert_success(route, function(**route_kwargs(route, api_base, "success"))) try: @@ -193,7 +176,7 @@ def exercise_sync(native: object, api_base: str) -> None: async def exercise_async(native: object, api_base: str) -> None: - for route in ("transcription", "messages", "chat_completions"): + for route in ("transcription", "chat_completions"): function: Final = getattr(native, f"a{route}") assert_success(route, await function(**route_kwargs(route, api_base, "success"))) try: @@ -206,11 +189,11 @@ async def exercise_async(native: object, api_base: str) -> None: async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), + asyncio.gather(*(native.achat_completions(**route_kwargs("chat_completions", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: - assert_success("messages", response) + assert_success("chat_completions", response) def exercise_routes(native_path: Path, api_base: str) -> object: @@ -223,8 +206,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object: def exercise_signal(native: object, api_base: str) -> int: try: - native.messages( - **route_kwargs("messages", api_base, "hang"), + native.chat_completions( + **route_kwargs("chat_completions", api_base, "hang"), ) except KeyboardInterrupt: sys.stdout.write("KeyboardInterrupt\n") diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index 72390b79141..b882a1bb8c2 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -43,8 +43,8 @@ def test_binding_validates_native_attribute( ROUTE_BINDINGS: Final = ( ("completion", chat_completions.NATIVE_COMPLETION), ("acompletion", chat_completions.NATIVE_ACOMPLETION), - ("anthropic_messages_handler", messages.NATIVE_MESSAGES), - ("anthropic_messages", messages.NATIVE_AMESSAGES), + ("messages", messages.NATIVE_MESSAGES), + ("amessages", messages.NATIVE_AMESSAGES), ("responses", responses.NATIVE_RESPONSES), ("aresponses", responses.NATIVE_ARESPONSES), ("ocr", ocr.NATIVE_OCR), diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index 2c737b0160e..e9fdbf859f4 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -40,6 +40,10 @@ def test_shipped_decisions( enabled: Final = environment == "1" if environment is not None else process is not False assert catalog.rollout(context) is Rollout.RUST_OPT_OUT assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) + elif route is Route.MESSAGES: + enabled: Final = environment == "1" if environment is not None else process is True + assert catalog.rollout(context) is Rollout.RUST_OPT_IN + assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) elif route is Route.TRANSCRIPTION and provider == "bedrock": assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index ade0ae549fb..fa6c0b30413 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -157,7 +157,6 @@ def test_context_outside_rule_stays_on_python() -> None: ( Context(Route.CHAT_COMPLETIONS, provider="anthropic"), Context(Route.CHAT_COMPLETIONS, provider="bedrock"), - Context(Route.MESSAGES, provider="anthropic"), Context(Route.RESPONSES, provider="openai"), Context(Route.TRANSCRIPTION, provider="openai"), ), diff --git a/tests/test_litellm_rust/messages/__init__.py b/tests/test_litellm_rust/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py new file mode 100644 index 00000000000..b55bc47d640 --- /dev/null +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -0,0 +1,175 @@ +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import ( + MESSAGES, + MESSAGES_EVENTS, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + request_body, +) + +pytestmark = pytest.mark.requires_rust_extension + +STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: + return { + "model": MESSAGES_MODEL, + "messages": [dict(message) for message in MESSAGES], + "max_tokens": 64, + "api_key": "test-key", + "api_base": server.base_url, + **kwargs, + } + + +def assert_served_natively(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +@pytest.mark.asyncio +async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + response: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success") + ) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + sent: Final = messages_server.requests[0] + assert sent.path == "/v1/messages" + assert sent.body == {"model": "claude-sonnet-5", "messages": list(MESSAGES), "max_tokens": 64, "stream": False} + pre_call: Final = recorder.wait_for("log_pre_api_call") + assert request_body(pre_call[0].kwargs) == sent.body + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].call_type == "anthropic_messages" + assert success[0].kwargs["litellm_call_id"] == "messages-success" + assert success[0].response.choices[0].message.content == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None: + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["temperature"] = 0.25 + + await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()])) + + assert messages_server.requests[0].body["temperature"] == 0.25 + + +@pytest.mark.asyncio +async def test_native_messages_provider_error_reaches_caller_and_failure_callbacks_as_one_public_error( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue( + ResponseSpec(body={"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}, status=400) + ) + observed: Final = [] + + class Observe(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs["exception"])) + + with pytest.raises(litellm.BadRequestError) as raised: + await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()])) + + assert_served_natively(messages_server) + assert [phase for phase, _ in observed] == ["sync", "async"] + assert all(error is raised.value for _, error in observed) + + +def sse_payload() -> bytes: + return b"".join(STREAM.payloads()) + + +@pytest.mark.asyncio +async def test_native_messages_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + first: Final = await anext(stream) + await drain_logging() + assert "async_log_success_event" not in recorder.names + rest: Final = [chunk async for chunk in stream] + + assert first + b"".join(rest) == sse_payload() + assert_served_natively(messages_server) + assert messages_server.requests[0].body["stream"] is True + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].kwargs["stream"] is True + assert success[0].kwargs["completion_start_time"] is not None + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_native_messages_stream_closed_early_logs_success_once_for_what_was_delivered( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + await anext(stream) + await stream.aclose() + + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + with pytest.raises(StopAsyncIteration): + await anext(stream) + + +def test_native_sync_messages_stream_relays_provider_events_and_logs_success_once( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) + assert isinstance(stream, Iterator) + + assert b"".join(stream) == sse_payload() + assert_served_natively(messages_server) + assert len(recorder.wait_for("async_log_success_event")) == 1 + + +def test_native_sync_messages_returns_the_provider_message(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder])) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + assert len(recorder.wait_for("log_success_event")) == 1 diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 228ed2cc454..3eea47751d3 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -25,6 +25,12 @@ class ResponseSpec: status: int = 200 headers: dict[str, str] = field(default_factory=dict) delay: float = 0 + events: tuple[tuple[str, object], ...] = () + + def payloads(self) -> tuple[bytes, ...]: + if not self.events: + return (json.dumps(self.body).encode(),) + return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events) @dataclass @@ -73,15 +79,17 @@ def recording_service() -> Iterator[RecordingServer]: response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) if response.delay: time.sleep(response.delay) - payload: Final = json.dumps(response.body).encode() + payloads: Final = response.payloads() self.send_response(response.status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) + self.send_header("Content-Type", "text/event-stream" if response.events else "application/json") + self.send_header("Content-Length", str(sum(len(payload) for payload in payloads))) for name, value in response.headers.items(): self.send_header(name, value) self.end_headers() try: - self.wfile.write(payload) + for payload in payloads: + self.wfile.write(payload) + self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index b60cf5eac02..c9cf81b83ca 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -12,6 +12,41 @@ OCR_RESPONSE: Final = { "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, } +MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5" +MESSAGES: Final = ({"role": "user", "content": "Hello"},) +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Hello from native Messages"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} +MESSAGES_EVENTS: Final = ( + ("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello from native Messages"}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) + def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: return { From 3a5b7c12ef119474a596cd58f59807cce5854fb5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:35:20 -0700 Subject: [PATCH 26/33] refactor(rust): separate machine events from the Python lifecycle's events --- .../crates/callbacks-legacy/src/adapter.rs | 57 +++++++-------- .../tests/deployment_hooks.rs | 11 ++- .../crates/callbacks-legacy/tests/payload.rs | 8 +-- .../crates/callbacks-legacy/tests/terminal.rs | 14 ++-- litellm-rust/crates/callbacks/src/event.rs | 16 +++-- litellm-rust/crates/callbacks/src/host.rs | 4 +- litellm-rust/crates/callbacks/src/run.rs | 5 +- litellm-rust/crates/core/src/machine/mod.rs | 4 +- .../crates/core/src/messages/route.rs | 4 +- litellm-rust/crates/core/src/ocr/handler.rs | 4 +- .../tests/azure_document_intelligence_ocr.rs | 8 +-- litellm-rust/crates/core/tests/ocr.rs | 9 ++- litellm-rust/crates/core/tests/reducto_ocr.rs | 8 +-- .../crates/host-python/src/adapter.rs | 33 ++++++--- litellm-rust/crates/host-python/src/driver.rs | 69 ++++++++++--------- litellm-rust/crates/host-python/src/lib.rs | 2 +- 16 files changed, 142 insertions(+), 114 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 9d28db92add..a67da2188fb 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -3,10 +3,10 @@ //! `@client` path makes them. use litellm_callbacks::event::{ - CallEvent, FailureOrigin, RequestContext, Timing, WireRequest, epoch_seconds, + FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds, }; use litellm_host_python::{ - LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, + LifecycleEvent, LifecycleStep, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -346,28 +346,11 @@ impl PythonLifecycle for LegacyLogging { fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult { - match (event, public) { - (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), - (CallEvent::Opened, _) => { - Streaming::Opened.call(py, (self.logger()?.object(py),))?; - self.stream = Some(DeliveredStream { - chunks: PyList::empty(py).unbind(), - first_chunk: None, - }); - Ok(LifecycleStep::Done) - } - (CallEvent::Delivered, Some(PublicValue::Chunk(chunk))) => { - let stream = self.stream.as_mut().ok_or_else(missing_state)?; - if stream.first_chunk.is_none() { - stream.first_chunk = Some(datetime(py, epoch_seconds())?); - } - stream.chunks.bind(py).append(chunk)?; - Ok(LifecycleStep::Done) - } - (CallEvent::ResponseReceived { raw }, _) => { + match event { + LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { let api_key = self .context .as_ref() @@ -382,7 +365,7 @@ impl PythonLifecycle for LegacyLogging { )?; Ok(LifecycleStep::Done) } - (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { + LifecycleEvent::Succeeded { timing, response } => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { @@ -391,13 +374,17 @@ impl PythonLifecycle for LegacyLogging { } Ok(LifecycleStep::Done) } - (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Failed { + timing, + origin, + error, + } => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); if self.stream.is_some() { return self.stream_failure(py); } - if *origin == FailureOrigin::Call + if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() { @@ -412,10 +399,26 @@ impl PythonLifecycle for LegacyLogging { } self.dispatch_failure(py) } - _ => Err(missing_state()), } } + fn opened(&mut self, py: Python<'_>) -> PyResult<()> { + Streaming::Opened.call(py, (self.logger()?.object(py),))?; + self.stream = Some(DeliveredStream { + chunks: PyList::empty(py).unbind(), + first_chunk: None, + }); + Ok(()) + } + + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { + let stream = self.stream.as_mut().ok_or_else(missing_state)?; + if stream.first_chunk.is_none() { + stream.first_chunk = Some(datetime(py, epoch_seconds())?); + } + stream.chunks.bind(py).append(chunk) + } + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 7bad09c7890..f4602739548 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; +use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -217,13 +217,12 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle .resume(py, Ok(local(&locals, "kwargs").unbind())) .unwrap(); let failure = PyErr::from_value(local(&locals, "failure")); - let failed = CallEvent::Failed { + let failed = LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Call, + error: &failure, }; - let step = logging - .emit(py, &failed, Some(PublicValue::Error(&failure))) - .unwrap(); + let step = logging.emit(py, failed).unwrap(); assert!(awaits_deployment_hook(&step)); let hook_result = if cancelled { Err(CancelledError::new_err("cancelled")) diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 43128bc38ea..67ad4ab8a2a 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,8 +1,8 @@ use std::ffi::CStr; use litellm_auth::SecretValue; -use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleStep, PythonLifecycle, to_py}; +use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; use pyo3::prelude::*; use rstest::rstest; @@ -94,13 +94,13 @@ fn before_send_bound( body, }; let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = CallEvent::ResponseReceived { + let raw = MachineEvent::ResponseReceived { raw: RawResponse { body: "raw response".into(), }, }; assert!(matches!( - logging.emit(py, &raw, None).unwrap(), + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), LifecycleStep::Done )); run(py, &locals, c"check()"); diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 3094d7b88d2..5688f70f387 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; +use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -33,8 +33,10 @@ fn succeed( logging .emit( py, - &CallEvent::Succeeded { timing: TIMING }, - Some(PublicValue::Response(&response)), + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, ) .unwrap() } @@ -44,11 +46,11 @@ fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) logging .emit( py, - &CallEvent::Failed { + LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Host, + error: &failure, }, - Some(PublicValue::Error(&failure)), ) .unwrap() } diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index d19e973b812..182dab657d3 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -52,18 +52,20 @@ pub enum FailureOrigin { Host, } +/// What a machine reports while it runs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MachineEvent { + ResponseReceived { raw: RawResponse }, +} + +/// What an in-process host observes: the machine's own events between the driver's +/// start and terminal ones. #[derive(Clone, Debug, PartialEq)] pub enum CallEvent { Started { start_time: f64, }, - ResponseReceived { - raw: RawResponse, - }, - /// The call streams and its stream was handed to the caller. - Opened, - /// One chunk of an open stream reached the caller. - Delivered, + Machine(MachineEvent), Succeeded { timing: Timing, }, diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/callbacks/src/host.rs index eef3e1da8d5..aba35185a18 100644 --- a/litellm-rust/crates/callbacks/src/host.rs +++ b/litellm-rust/crates/callbacks/src/host.rs @@ -1,6 +1,6 @@ use std::future::Future; -use crate::event::{CallEvent, RequestContext, WireRequest}; +use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest}; use crate::route::Route; /// One suspension point of a native call, performed by the host. @@ -10,7 +10,7 @@ pub enum HostOp { wire: Box, context: Box, }, - Emit(CallEvent), + Emit(MachineEvent), /// The response streams: the host hands the caller a stream and answers once the /// caller asks for the first chunk or goes away. Open(R::StreamHead), diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 705c5504c01..6a0c08fba68 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -25,7 +25,10 @@ where .before_send(*wire, &context) .await .map(|wire| HostResult::BeforeSend(Box::new(wire))), - HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted), + HostOp::Emit(event) => host + .emit(&CallEvent::Machine(event)) + .await + .map(|()| HostResult::Emitted), HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), }; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index 929f0a423c4..d6db488159e 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -8,7 +8,7 @@ use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; use litellm_callbacks::{ - event::{CallEvent, RequestContext, WireRequest}, + event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, @@ -82,7 +82,7 @@ where } } - pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> { + pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { match self.invoke(HostOp::Emit(event)).await? { HostResult::Emitted => Ok(()), _ => Err(MachineFault::Mismatch.into()), diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index a680b5ce1ec..2fd2f79907a 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -3,7 +3,7 @@ use std::{sync::Mutex, time::Duration}; use bytes::Bytes; use litellm_auth::SecretValue; use litellm_callbacks::{ - event::{CallEvent, RawResponse, RequestContext, WireRequest}, + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, route::Route, }; @@ -172,7 +172,7 @@ async fn execute(host: MessagesHost) -> Result { return relay(&host, response).await; } let text = response.text().await.map_err(network)?; - host.emit(CallEvent::ResponseReceived { + host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: text.clone() }, }) .await?; diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index a6af190eb91..aff1eded2cc 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,6 +1,6 @@ use futures_util::future::BoxFuture; use litellm_auth::SecretValue; -use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; +use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -71,7 +71,7 @@ impl CallHooks for OcrCallHooks { } fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(self.host.emit(CallEvent::ResponseReceived { + Box::pin(self.host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: String::from_utf8_lossy(body).into_owned(), }, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 1fc4d6c2b9e..62544931dfc 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::CallEvent; +use litellm_callbacks::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; @@ -263,7 +263,7 @@ async fn accepted_response_emits_response_received_before_polling() { json!({}), )) .with_observer(move |event| { - let CallEvent::ResponseReceived { raw } = event else { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { return; }; match request_count.lock().unwrap().len() { @@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { mod transformation { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::CallEvent; + use litellm_callbacks::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::transformation::OcrDocument; use serde_json::{Value, json}; @@ -646,7 +646,7 @@ mod transformation { json!({}), )) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed .lock() .unwrap() diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 779d2037bb3..a6001951361 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use litellm_callbacks::{ - event::{CallEvent, WireRequest}, + event::{CallEvent, MachineEvent, WireRequest}, host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; @@ -195,11 +195,9 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { CallEvent::Started { .. } => "started", - CallEvent::ResponseReceived { .. } => "response", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", - CallEvent::Opened => "opened", - CallEvent::Delivered => "delivered", } } @@ -377,6 +375,7 @@ async fn drive_until( intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) } HostOp::Emit(event) => { + let event = CallEvent::Machine(event); ops.push(event_name(&event)); host.emit(&event) .await @@ -420,7 +419,7 @@ async fn invalid_provider_response_emits_response_received_before_normalization_ let observed = responses_received.clone(); let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed.lock().unwrap().push(raw.body.clone()); } }, diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 59891b16e90..8c037889a17 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, WireRequest}; +use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; use rstest::rstest; use serde_json::{Value, json}; @@ -139,7 +139,7 @@ async fn response_received_stays_after_reducto_upload_and_parse() { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } @@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() { } mod transformation { - use litellm_callbacks::event::{CallEvent, WireRequest}; + use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::{ base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, reducto::ocr::transformation::*, @@ -506,7 +506,7 @@ mod transformation { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 795977438fd..f946dbc763a 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, RequestContext, Timing, WireRequest}; +use litellm_callbacks::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; use litellm_callbacks::route::Route; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; @@ -19,11 +19,22 @@ pub enum LifecycleStep { Done, } -/// The host-typed value the driver attaches to a terminal event. -pub enum PublicValue<'a> { - Response(&'a Py), - Error(&'a PyErr), - Chunk(&'a Py), +/// What a lifecycle observes: the driver's start, the machine's own events, and one +/// terminal event carrying the public value the caller receives. +pub enum LifecycleEvent<'a> { + Started { + start_time: f64, + }, + Machine(&'a MachineEvent), + Succeeded { + timing: Timing, + response: &'a Py, + }, + Failed { + timing: Timing, + origin: FailureOrigin, + error: &'a PyErr, + }, } /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in @@ -61,10 +72,16 @@ pub trait PythonLifecycle: Send + Sync { fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult; + /// The call streams and its stream was handed to the caller. The caller is not + /// inside an await here, so this step and `delivered` cannot suspend. + fn opened(&mut self, py: Python<'_>) -> PyResult<()>; + + /// One chunk of an open stream is about to reach the caller. + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()>; + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index c59ba1925dc..25f78d7013e 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; +use litellm_callbacks::event::{FailureOrigin, Timing, epoch_seconds}; use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; use litellm_callbacks::route::Route; @@ -13,7 +13,7 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, + HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -159,10 +159,10 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let started = CallEvent::Started { + let started = LifecycleEvent::Started { start_time: self.started_at, }; - match self.adapter.emit(py, &started, None) { + match self.adapter.emit(py, started) { Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } @@ -303,7 +303,7 @@ where } HostOp::Open(_) => return self.opened(py).map(Next::Return), HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), - HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { + HostOp::Emit(event) => match self.adapter.emit(py, LifecycleEvent::Machine(&event)) { Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); @@ -321,12 +321,11 @@ where fn opened(&mut self, py: Python<'_>) -> PyResult { self.stage = Stage::Streaming; - match self.adapter.emit(py, &CallEvent::Opened, None) { - Ok(LifecycleStep::Done) => { + match self.adapter.opened(py) { + Ok(()) => { self.pending = Some(Pending::Consumer); Ok(ExecutionStep::Open) } - Ok(_) => Err(missing_state()), Err(error) => self.interrupt(py, error), } } @@ -340,15 +339,11 @@ where Ok(chunk) => chunk, Err(error) => return self.interrupt(py, error), }; - let observed = - self.adapter - .emit(py, &CallEvent::Delivered, Some(PublicValue::Chunk(&chunk))); - match observed { - Ok(LifecycleStep::Done) => { + match self.adapter.delivered(py, &chunk) { + Ok(()) => { self.pending = Some(Pending::Consumer); Ok(ExecutionStep::Yield(chunk)) } - Ok(_) => Err(missing_state()), Err(error) => self.interrupt(py, error), } } @@ -461,12 +456,11 @@ where } fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { - let event = CallEvent::Succeeded { + let event = LifecycleEvent::Succeeded { timing: self.timing(), + response: &response, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Response(&response)))?; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Succeeded(response); self.on_adapter(py, step, Expect::Terminal) } @@ -481,13 +475,12 @@ where if is_cancellation(py, &error) { return Err(error); } - let event = CallEvent::Failed { + let event = LifecycleEvent::Failed { timing: self.timing(), origin, + error: &error, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Error(&error)))?; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -541,7 +534,7 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{RequestContext, WireRequest}; + use litellm_callbacks::event::{MachineEvent, RequestContext, WireRequest}; use litellm_callbacks::machine::{Interrupted, Step}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -803,23 +796,33 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult { - self.log.push(match (event, public) { - (CallEvent::Started { .. }, None) => "started".into(), - (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), - (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { - format!("succeeded:{}", value.bind(py)) + self.log.push(match event { + LifecycleEvent::Started { .. } => "started".into(), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + format!("response:{}", raw.body) } - (CallEvent::Failed { origin, .. }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Succeeded { response, .. } => { + format!("succeeded:{}", response.bind(py)) + } + LifecycleEvent::Failed { origin, error, .. } => { format!("failed:{origin:?}:{}", error.value(py)) } - _ => "unexpected".into(), }); Ok(LifecycleStep::Done) } + fn opened(&mut self, _: Python<'_>) -> PyResult<()> { + self.log.push("opened"); + Ok(()) + } + + fn delivered(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + self.log.push("delivered"); + Ok(()) + } + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -899,7 +902,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri wire: Box::new(wire()), context: Box::new(context()), }, - HostOp::Emit(CallEvent::ResponseReceived { + HostOp::Emit(MachineEvent::ResponseReceived { raw: litellm_callbacks::event::RawResponse { body: "raw".into() }, }), ], diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8889b1513db..3738a9b1c3c 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -13,7 +13,7 @@ mod handle; mod marshal; pub use adapter::{ - HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, + HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; pub use argument::lookup; pub use callable::wrap_failure; From c2679757b3ad93bd9f8def74021e78be42a025c7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:36:42 -0700 Subject: [PATCH 27/33] refactor(rust): put stream billing on the legacy surface --- .../callbacks-legacy/python_contract.json | 3 ++ .../crates/callbacks-legacy/src/adapter.rs | 37 +++++++++++++++---- .../crates/callbacks-legacy/src/lib.rs | 2 +- .../crates/callbacks-legacy/tests/support.rs | 1 + .../crates/host-python/src/adapter.rs | 6 +-- litellm-rust/crates/host-python/src/driver.rs | 6 +-- .../python-bridge/src/routes/messages/mod.rs | 6 ++- .../python-bridge/src/routes/ocr/mod.rs | 1 + litellm/rust_bridge/legacy_callbacks.py | 14 +++++-- 9 files changed, 53 insertions(+), 23 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json index a09bdc711a3..8a7f3b98f47 100644 --- a/litellm-rust/crates/callbacks-legacy/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -100,6 +100,8 @@ ], "stream_success": [ "logger", + "url_route", + "endpoint_type", "request_body", "chunks", "start", @@ -108,6 +110,7 @@ ], "stream_failure": [ "logger", + "endpoint_type", "request_body", "chunks", "error" diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index a67da2188fb..4e9167100e0 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -30,6 +30,16 @@ pub struct LegacySurface { pub call_type: &'static str, /// What `Logging.pre_call` is told the input was. pub input_description: &'static str, + /// How a streamed response is billed; `None` for a route that never streams. + pub stream: Option, +} + +/// The pass-through billing a streamed response goes through once its chunks are in. +#[derive(Clone, Copy, Debug)] +pub struct PassThroughStream { + pub url_route: &'static str, + /// A value of Python's `EndpointType`. + pub endpoint_type: &'static str, } /// What the Messages stream iterator keeps for its end-of-stream billing. @@ -177,10 +187,13 @@ impl LegacyLogging { fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { let logger = self.logger()?; + let billing = self.surface.stream.ok_or_else(missing_state)?; let billed = Streaming::Success.call( py, ( logger.object(py), + billing.url_route, + billing.endpoint_type, &self.body, &stream.chunks, &self.start, @@ -201,14 +214,25 @@ impl LegacyLogging { /// partial usage. The sync path has no loop to schedule that on, so it falls back to /// the plain failure handler. fn stream_failure(&mut self, py: Python<'_>) -> PyResult { - let (Some(logger), Some(error), Some(stream)) = (&self.logger, &self.error, &self.stream) + let (Some(logger), Some(error), Some(stream), Some(billing)) = + (&self.logger, &self.error, &self.stream, self.surface.stream) else { return Ok(LifecycleStep::Done); }; if !self.asynchronous { return self.dispatch_failure(py); } - match Streaming::Failure.call(py, (logger.object(py), &self.body, &stream.chunks, error)) { + let scheduled = Streaming::Failure.call( + py, + ( + logger.object(py), + billing.endpoint_type, + &self.body, + &stream.chunks, + error, + ), + ); + match scheduled { Ok(awaitable) => { self.pending = Some(Pending::AsyncFailure); Ok(LifecycleStep::Await(awaitable.unbind())) @@ -343,11 +367,7 @@ impl PythonLifecycle for LegacyLogging { self.finalize(py) } - fn emit( - &mut self, - py: Python<'_>, - event: LifecycleEvent<'_>, - ) -> PyResult { + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult { match event { LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done), LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { @@ -403,6 +423,9 @@ impl PythonLifecycle for LegacyLogging { } fn opened(&mut self, py: Python<'_>) -> PyResult<()> { + if self.surface.stream.is_none() { + return Err(missing_state()); + } Streaming::Opened.call(py, (self.logger()?.object(py),))?; self.stream = Some(DeliveredStream { chunks: PyList::empty(py).unbind(), diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy/src/lib.rs index 42ffd545e2b..eaa1a8b714e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy/src/lib.rs @@ -21,7 +21,7 @@ mod preparation; mod test_support; pub(crate) use adapter::LegacyLogging; -pub use adapter::LegacySurface; +pub use adapter::{LegacySurface, PassThroughStream}; pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 42ca184eb16..d3cc32e301f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -197,6 +197,7 @@ pub(crate) fn legacy_call( LegacySurface { call_type: "test", input_description: "test input", + stream: None, }, call, asynchronous, diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index f946dbc763a..e70e4a8c58f 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -69,11 +69,7 @@ pub trait PythonLifecycle: Send + Sync { timing: Timing, ) -> PyResult; - fn emit( - &mut self, - py: Python<'_>, - event: LifecycleEvent<'_>, - ) -> PyResult; + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult; /// The call streams and its stream was handed to the caller. The caller is not /// inside an await here, so this step and `delivered` cannot suspend. diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 25f78d7013e..d6b31b27314 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -793,11 +793,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } - fn emit( - &mut self, - py: Python<'_>, - event: LifecycleEvent<'_>, - ) -> PyResult { + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult { self.log.push(match event { LifecycleEvent::Started { .. } => "started".into(), LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index b4259aca5ba..8c42315ac59 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,7 +1,7 @@ mod host; use host::MessagesRouteHost; -use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; +use litellm_callbacks_legacy::{LegacySurface, PassThroughStream, PublicCall, run_legacy_call}; use litellm_core::messages::route::{messages_machine, supports}; use pyo3::{ prelude::*, @@ -13,6 +13,10 @@ use crate::errors::RustBridgeDeclined; const SURFACE: LegacySurface = LegacySurface { call_type: "anthropic_messages", input_description: "Messages", + stream: Some(PassThroughStream { + url_route: "/v1/messages", + endpoint_type: "anthropic", + }), }; fn run_messages( diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index b5bb941708d..8afa1e2a906 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -15,6 +15,7 @@ use pyo3::{ const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", input_description: "OCR document processing", + stream: None, }; const ASYNC_SURFACE: LegacySurface = LegacySurface { diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index fd0fac1799a..bac40442ce5 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -331,6 +331,8 @@ def stream_opened(logger: Logging) -> None: def stream_success( logger: Logging, + url_route: str, + endpoint_type: str, request_body: dict[str, object], chunks: list[bytes], start: datetime.datetime, @@ -353,9 +355,9 @@ def stream_success( coroutine: Final = build( litellm_logging_obj=logger, passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, - url_route="/v1/messages", + url_route=url_route, request_body=request_body, - endpoint_type=EndpointType.ANTHROPIC, + endpoint_type=EndpointType(endpoint_type), start_time=start, raw_bytes=chunks, end_time=end, @@ -374,14 +376,18 @@ def stream_success( def stream_failure( - logger: Logging, request_body: dict[str, object], chunks: list[bytes], error: Exception + logger: Logging, + endpoint_type: str, + request_body: dict[str, object], + chunks: list[bytes], + error: Exception, ) -> Coroutine[object, object, None]: from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType return PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=logger, - endpoint_type=EndpointType.ANTHROPIC, + endpoint_type=EndpointType(endpoint_type), request_body=request_body, raw_bytes=chunks, exception=error, From 19ffb584eb6b9c28118f96fb47642d1fca5442b3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:37:05 -0700 Subject: [PATCH 28/33] refactor(rust): rename litellm-callbacks to litellm-host and HostOpError to InvokeError --- litellm-rust/Cargo.lock | 28 +++++++++---------- litellm-rust/Cargo.toml | 2 +- .../crates/callbacks-legacy/AGENTS.md | 2 +- .../crates/callbacks-legacy/Cargo.toml | 2 +- .../crates/callbacks-legacy/src/adapter.rs | 2 +- .../crates/callbacks-legacy/src/call.rs | 2 +- .../crates/callbacks-legacy/src/callbacks.rs | 2 +- .../tests/deployment_hooks.rs | 2 +- .../crates/callbacks-legacy/tests/payload.rs | 2 +- .../crates/callbacks-legacy/tests/terminal.rs | 2 +- litellm-rust/crates/core/Cargo.toml | 2 +- litellm-rust/crates/core/src/machine/auth.rs | 2 +- litellm-rust/crates/core/src/machine/mod.rs | 2 +- litellm-rust/crates/core/src/messages/mod.rs | 2 +- .../crates/core/src/messages/route.rs | 4 +-- litellm-rust/crates/core/src/ocr/client.rs | 2 +- litellm-rust/crates/core/src/ocr/handler.rs | 2 +- litellm-rust/crates/core/src/ocr/route.rs | 4 +-- .../tests/azure_document_intelligence_ocr.rs | 4 +-- litellm-rust/crates/core/tests/ocr.rs | 6 ++-- .../crates/core/tests/ocr/document.rs | 2 +- litellm-rust/crates/core/tests/ocr/support.rs | 4 +-- litellm-rust/crates/core/tests/reducto_ocr.rs | 4 +-- litellm-rust/crates/host-python/AGENTS.md | 4 +-- litellm-rust/crates/host-python/Cargo.toml | 2 +- .../crates/host-python/src/adapter.rs | 10 +++---- litellm-rust/crates/host-python/src/driver.rs | 26 ++++++++--------- litellm-rust/crates/host-python/src/lib.rs | 4 +-- .../crates/{callbacks => host}/Cargo.toml | 2 +- .../crates/{callbacks => host}/src/event.rs | 0 .../crates/{callbacks => host}/src/host.rs | 0 .../crates/{callbacks => host}/src/lib.rs | 0 .../crates/{callbacks => host}/src/machine.rs | 0 .../crates/{callbacks => host}/src/route.rs | 0 .../crates/{callbacks => host}/src/run.rs | 0 litellm-rust/crates/llms/Cargo.toml | 2 +- .../llms/src/custom_httpx/llm_http_handler.rs | 2 +- .../python-bridge/src/routes/messages/host.rs | 6 ++-- .../python-bridge/src/routes/ocr/host.rs | 6 ++-- 39 files changed, 75 insertions(+), 75 deletions(-) rename litellm-rust/crates/{callbacks => host}/Cargo.toml (91%) rename litellm-rust/crates/{callbacks => host}/src/event.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/host.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/lib.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/machine.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/route.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/run.rs (100%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index bf58a81a6c5..5a8a613204f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2027,22 +2027,12 @@ dependencies = [ "tokio", ] -[[package]] -name = "litellm-callbacks" -version = "0.1.0" -dependencies = [ - "litellm-auth", - "rstest", - "serde_json", - "tokio", -] - [[package]] name = "litellm-callbacks-legacy" version = "0.1.0" dependencies = [ "litellm-auth", - "litellm-callbacks", + "litellm-host", "litellm-host-python", "proptest", "pyo3", @@ -2060,8 +2050,8 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-aws", - "litellm-callbacks", "litellm-core-utils", + "litellm-host", "litellm-llms", "litellm-types", "mime_guess", @@ -2114,12 +2104,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-host" +version = "0.1.0" +dependencies = [ + "litellm-auth", + "rstest", + "serde_json", + "tokio", +] + [[package]] name = "litellm-host-python" version = "0.1.0" dependencies = [ "futures-util", - "litellm-callbacks", + "litellm-host", "pyo3", "pyo3-async-runtimes", "pythonize", @@ -2143,9 +2143,9 @@ dependencies = [ "litellm-auth-aws", "litellm-auth-azure", "litellm-auth-gcp", - "litellm-callbacks", "litellm-core-utils", "litellm-framing", + "litellm-host", "litellm-types", "reqwest 0.12.28", "rstest", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 32f925b8d8b..de6eacc62ee 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -10,7 +10,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } -litellm-callbacks = { path = "crates/callbacks" } +litellm-host = { path = "crates/host" } litellm-callbacks-legacy = { path = "crates/callbacks-legacy" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy/AGENTS.md index e184e2fb415..8b2e1c15f6e 100644 --- a/litellm-rust/crates/callbacks-legacy/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy/AGENTS.md @@ -11,7 +11,7 @@ - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view + - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-host`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view - Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index efacde051ed..023c13d912b 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true autotests = false [dependencies] -litellm-callbacks.workspace = true +litellm-host.workspace = true litellm-host-python.workspace = true pyo3.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 4e9167100e0..6c013cd1ea5 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -2,7 +2,7 @@ //! raises is answered with the same `Logging` calls, in the same order, as the Python //! `@client` path makes them. -use litellm_callbacks::event::{ +use litellm_host::event::{ FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds, }; use litellm_host_python::{ diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy/src/call.rs index bd1e2525d3d..b37790f60a8 100644 --- a/litellm-rust/crates/callbacks-legacy/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy/src/call.rs @@ -3,7 +3,7 @@ //! lifetime. No other callback host has that obligation, which is why nothing outside //! this crate holds them. -use litellm_callbacks::{machine::Machine, route::Route}; +use litellm_host::{machine::Machine, route::Route}; use litellm_host_python::{RouteHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index 9fcfe98368e..5f04224e6d7 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -2,7 +2,7 @@ //! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls //! duplication. All of it expires with the legacy callback contract. -use litellm_callbacks::event::{RequestContext, WireRequest}; +use litellm_host::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index f4602739548..52c5e47f83f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,6 +1,6 @@ use std::ffi::CStr; -use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host::event::{FailureOrigin, Timing}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 67ad4ab8a2a..5459b36af27 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_auth::SecretValue; -use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; use pyo3::prelude::*; diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 5688f70f387..f68209233f2 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,6 +1,6 @@ use std::ffi::CStr; -use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host::event::{FailureOrigin, Timing}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index db6cfc4b340..3995a235778 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -9,7 +9,7 @@ autotests = false [dependencies] litellm-types.workspace = true litellm-core-utils.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/core/src/machine/auth.rs b/litellm-rust/crates/core/src/machine/auth.rs index 6a3e4daf6ee..cf91458ae43 100644 --- a/litellm-rust/crates/core/src/machine/auth.rs +++ b/litellm-rust/crates/core/src/machine/auth.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -use litellm_callbacks::route::Route; +use litellm_host::route::Route; use super::{HostChannel, MachineFault}; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index d6db488159e..ffcefc663ab 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -7,7 +7,7 @@ mod auth; use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; -use litellm_callbacks::{ +use litellm_host::{ event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index e36c6668efe..c07a83c7e43 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -33,7 +33,7 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result Ok(message), MessagesOutput::Streamed => Err(Error::Unsupported( "streamed responses need a streaming host", diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 2fd2f79907a..549f848049f 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -2,12 +2,12 @@ use std::{sync::Mutex, time::Duration}; use bytes::Bytes; use litellm_auth::SecretValue; -use litellm_callbacks::{ +use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; +use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, route::Route, }; -use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 03782d91f24..c05622932b1 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -12,7 +12,7 @@ pub async fn perform( client: &OcrClient, request: LiteLLMOcrRequest, ) -> Result { - litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await + litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } pub async fn ocr(request: LiteLLMOcrRequest) -> Result { diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index aff1eded2cc..bbf9cfa0e02 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,6 +1,6 @@ use futures_util::future::BoxFuture; use litellm_auth::SecretValue; -use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index e6ee45c64a8..d3711c34fa7 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use litellm_auth::ResolvedCredential; -use litellm_callbacks::{ +use litellm_host::{ event::{CallEvent, RequestContext, WireRequest}, route::Route, }; @@ -175,7 +175,7 @@ impl LocalOcrHost { } } -impl litellm_callbacks::host::Host for LocalOcrHost { +impl litellm_host::host::Host for LocalOcrHost { async fn route(&self, op: OcrOp) -> Result { match op { OcrOp::ProjectRequest => self diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 62544931dfc..3cbe6fe3159 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, MachineEvent}; +use litellm_host::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; @@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { mod transformation { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{CallEvent, MachineEvent}; + use litellm_host::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::transformation::OcrDocument; use serde_json::{Value, json}; diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index a6001951361..41a650945bc 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,6 +1,6 @@ use std::sync::{Arc, Mutex}; -use litellm_callbacks::{ +use litellm_host::{ event::{CallEvent, MachineEvent, WireRequest}, host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, @@ -820,7 +820,7 @@ impl Host for CallerTokenHost { async fn before_send( &self, wire: WireRequest, - _: &litellm_callbacks::event::RequestContext, + _: &litellm_host::event::RequestContext, ) -> Result { let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); let authorization = wire @@ -855,7 +855,7 @@ async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_ trace: Mutex::new(Vec::new()), }; - litellm_callbacks::run::run(ocr_machine(ocr_client()), &host) + litellm_host::run::run(ocr_machine(ocr_client()), &host) .await .unwrap(); server.await.unwrap(); diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs index 5e10ce3ad2a..855548dc6bf 100644 --- a/litellm-rust/crates/core/tests/ocr/document.rs +++ b/litellm-rust/crates/core/tests/ocr/document.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::WireRequest; +use litellm_host::event::WireRequest; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index 44313d5f552..b368a754656 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; -use litellm_callbacks::event::WireRequest; +use litellm_host::event::WireRequest; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::{CallHooks, OcrClient}, @@ -45,7 +45,7 @@ pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result Result { - litellm_callbacks::run::run(ocr_machine(ocr_client()), &host).await + litellm_host::run::run(ocr_machine(ocr_client()), &host).await } pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 8c037889a17..83e7754122b 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; +use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; use rstest::rstest; use serde_json::{Value, json}; @@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() { } mod transformation { - use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::{ base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, reducto::ocr::transformation::*, diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 903ccb06c77..5aca13eeb18 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,9 +1,9 @@ - Target invariants; implementation and runtime validation may lag these rules - Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits - - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features + - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - - A native failure, including one a host op returns as `HostOpError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is + - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index ae0cebada59..e2c83fe1081 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] futures-util.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true pythonize.workspace = true diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index e70e4a8c58f..3a4cb49be4d 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,5 +1,5 @@ -use litellm_callbacks::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; -use litellm_callbacks::route::Route; +use litellm_host::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; +use litellm_host::route::Route; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -89,12 +89,12 @@ pub trait PythonLifecycle: Send + Sync { /// rejected it, which the route classifies like any other native failure, or Python code /// raised, which reaches the caller as it was raised. #[derive(Debug)] -pub enum HostOpError { +pub enum InvokeError { Native(E), Python(PyErr), } -impl From for HostOpError { +impl From for InvokeError { fn from(error: PyErr) -> Self { Self::Python(error) } @@ -117,7 +117,7 @@ pub trait RouteHost: Send + Sync { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: ::Op, - ) -> Result<::OpResult, HostOpError<::Error>>; + ) -> Result<::OpResult, InvokeError<::Error>>; fn complete( &mut self, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index d6b31b27314..392d36e10f4 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,10 +2,10 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_callbacks::event::{FailureOrigin, Timing, epoch_seconds}; -use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; -use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; -use litellm_callbacks::route::Route; +use litellm_host::event::{FailureOrigin, Timing, epoch_seconds}; +use litellm_host::host::{Demand, HostOp, HostResult, HostStep}; +use litellm_host::machine::{HostFailure, Machine, MachineStep}; +use litellm_host::route::Route; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -13,7 +13,7 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -282,12 +282,12 @@ where let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; match self.route.invoke(py, arguments.bind(py), op) { Ok(result) => Ok(HostResult::Route(result)), - Err(HostOpError::Native(error)) => { + Err(InvokeError::Native(error)) => { return self .resume_core(py, Some(Err(HostFailure::Error(error)))) .map(Next::Continue); } - Err(HostOpError::Python(error)) => Err(error), + Err(InvokeError::Python(error)) => Err(error), } } HostOp::BeforeSend { wire, context } => { @@ -534,8 +534,8 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{MachineEvent, RequestContext, WireRequest}; - use litellm_callbacks::machine::{Interrupted, Step}; + use litellm_host::event::{MachineEvent, RequestContext, WireRequest}; + use litellm_host::machine::{Interrupted, Step}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -692,12 +692,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri _: Python<'_>, arguments: &Bound<'_, PyDict>, op: &'static str, - ) -> Result> { + ) -> Result> { self.log.push(format!("route:{op}")); match self.op { OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), - OpScript::RejectNatively => Err(HostOpError::Native(Error("op rejected".into()))), + OpScript::RejectNatively => Err(InvokeError::Native(Error("op rejected".into()))), } } @@ -899,7 +899,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri context: Box::new(context()), }, HostOp::Emit(MachineEvent::ResponseReceived { - raw: litellm_callbacks::event::RawResponse { body: "raw".into() }, + raw: litellm_host::event::RawResponse { body: "raw".into() }, }), ], outcome: Some(Ok("done".into())), @@ -1189,7 +1189,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, - ) -> Result> { + ) -> Result> { self.0.push("route"); Err(PyErr::from_value( py.import("asyncio") diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 3738a9b1c3c..583a4eb91b6 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,5 +1,5 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and -//! asyncio glue, and the driver that runs a native [`Machine`](litellm_callbacks::machine::Machine) +//! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine) //! against a Python route host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. @@ -13,7 +13,7 @@ mod handle; mod marshal; pub use adapter::{ - HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; pub use argument::lookup; pub use callable::wrap_failure; diff --git a/litellm-rust/crates/callbacks/Cargo.toml b/litellm-rust/crates/host/Cargo.toml similarity index 91% rename from litellm-rust/crates/callbacks/Cargo.toml rename to litellm-rust/crates/host/Cargo.toml index a68ebc26a8d..ebabe5ccbdc 100644 --- a/litellm-rust/crates/callbacks/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "litellm-callbacks" +name = "litellm-host" version = "0.1.0" edition.workspace = true license.workspace = true diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/host/src/event.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/event.rs rename to litellm-rust/crates/host/src/event.rs diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/host/src/host.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/host.rs rename to litellm-rust/crates/host/src/host.rs diff --git a/litellm-rust/crates/callbacks/src/lib.rs b/litellm-rust/crates/host/src/lib.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/lib.rs rename to litellm-rust/crates/host/src/lib.rs diff --git a/litellm-rust/crates/callbacks/src/machine.rs b/litellm-rust/crates/host/src/machine.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/machine.rs rename to litellm-rust/crates/host/src/machine.rs diff --git a/litellm-rust/crates/callbacks/src/route.rs b/litellm-rust/crates/host/src/route.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/route.rs rename to litellm-rust/crates/host/src/route.rs diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/host/src/run.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/run.rs rename to litellm-rust/crates/host/src/run.rs diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 4ca6c7cb2a5..d295e4407ba 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -15,7 +15,7 @@ litellm-auth.workspace = true litellm-auth-aws.workspace = true litellm-auth-azure.workspace = true litellm-auth-gcp.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true litellm-framing.workspace = true base64.workspace = true bytes.workspace = true diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index 7635bdd3d04..fdd568d83fd 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -3,7 +3,7 @@ use std::{sync::OnceLock, time::Duration}; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks::event::WireRequest; +use litellm_host::event::WireRequest; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 0c87aea1ca2..b590018bc4e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,7 @@ use litellm_core::messages::{ Error, route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, }; -use litellm_host_python::{HostOpError, RouteHost, from_py, lookup, to_py}; +use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py}; use litellm_llms::custom_httpx::transport::Error as TransportError; use pyo3::{ exceptions::{PyException, PyValueError}, @@ -136,12 +136,12 @@ impl RouteHost for MessagesRouteHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: MessagesOp, - ) -> Result> { + ) -> Result> { match op { MessagesOp::ProjectRequest => self .project(py, arguments) .map(|call| MessagesOpResult::Request(Box::new(call))) - .map_err(|error| HostOpError::Python(self.map_failure(py, error))), + .map_err(|error| InvokeError::Python(self.map_failure(py, error))), } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 19211671fd7..77c8d5d6641 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,6 +1,6 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{HostOpError, RouteHost, missing_state, to_py}; +use litellm_host_python::{InvokeError, RouteHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -113,9 +113,9 @@ impl RouteHost for OcrRouteHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: OcrOp, - ) -> Result> { + ) -> Result> { self.answer(py, arguments, op) - .map_err(|error| HostOpError::Python(self.map_failure(py, error))) + .map_err(|error| InvokeError::Python(self.map_failure(py, error))) } fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { From b7686d78b675a020bb15123c571438f647674431 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:40:02 -0700 Subject: [PATCH 29/33] refactor(rust): move RouteMachine out of core into litellm-host next to the Machine trait core/src/machine was route-neutral runtime code sitting among route surfaces, and the workspace had two modules named machine. It now lives in litellm-host beside the contract it implements, so core only holds routes. The OCR conversion from MachineFault moves to litellm-llms because the orphan rule no longer allows it in core --- litellm-rust/crates/core/src/lib.rs | 1 - litellm-rust/crates/core/src/messages/route.rs | 6 ++---- litellm-rust/crates/core/src/ocr/route.rs | 16 ++-------------- litellm-rust/crates/host/Cargo.toml | 2 +- litellm-rust/crates/host/src/lib.rs | 2 +- .../crates/{core => host}/src/machine/auth.rs | 5 ++--- .../host/src/{machine.rs => machine/mod.rs} | 6 ++++++ .../mod.rs => host/src/machine/route_machine.rs} | 10 ++++------ .../crates/llms/src/base_llm/ocr/error.rs | 11 +++++++++++ 9 files changed, 29 insertions(+), 30 deletions(-) rename litellm-rust/crates/{core => host}/src/machine/auth.rs (98%) rename litellm-rust/crates/host/src/{machine.rs => machine/mod.rs} (91%) rename litellm-rust/crates/{core/src/machine/mod.rs => host/src/machine/route_machine.rs} (97%) diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 58aef6cd629..e3e2fb48721 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -2,7 +2,6 @@ pub mod audio_transcription; pub mod chat_completions; pub mod constants; pub mod error; -pub mod machine; pub mod messages; pub mod ocr; pub mod responses; diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 549f848049f..3d607ba67f6 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -6,6 +6,7 @@ use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, + machine::{HostChannel, MachineFault, RouteMachine}, route::Route, }; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; @@ -18,10 +19,7 @@ use super::{ prepare::prepare_provider_request, types::MessagesRequest, }; -use crate::{ - constants::ANTHROPIC_MESSAGES_PROVIDER, - machine::{HostChannel, MachineFault, RouteMachine}, -}; +use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesOp { diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index d3711c34fa7..bfc8c5ca965 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex}; use litellm_auth::ResolvedCredential; use litellm_host::{ event::{CallEvent, RequestContext, WireRequest}, + machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, route::Route, }; use litellm_llms::{ @@ -11,10 +12,7 @@ use litellm_llms::{ }; use super::handler::perform_ocr_request; -use crate::{ - machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, - ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}, -}; +use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrOp { @@ -56,16 +54,6 @@ impl TokenRoute for Ocr { } } -impl From for Error { - fn from(fault: MachineFault) -> Self { - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "OCR host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("OCR {message}"), - MachineFault::Mismatch => "invalid OCR host operation result".into(), - }) - } -} - pub type OcrHost = HostChannel; pub type OcrMachine = RouteMachine; diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index ebabe5ccbdc..0c7c46192b5 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -8,7 +8,7 @@ repository.workspace = true [dependencies] litellm-auth.workspace = true serde_json.workspace = true +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] rstest.workspace = true -tokio = { workspace = true, features = ["macros"] } diff --git a/litellm-rust/crates/host/src/lib.rs b/litellm-rust/crates/host/src/lib.rs index 41b0983f0ce..65479c2380f 100644 --- a/litellm-rust/crates/host/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -1,7 +1,7 @@ //! The contract between a native call and the host runtime that drives it. //! //! A host is whatever sits on the far side of the language boundary: CPython today, -//! another runtime later. Core implements [`machine::Machine`] per route and never learns +//! another runtime later. Core runs each route on a [`machine::RouteMachine`] and never learns //! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers //! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent. diff --git a/litellm-rust/crates/core/src/machine/auth.rs b/litellm-rust/crates/host/src/machine/auth.rs similarity index 98% rename from litellm-rust/crates/core/src/machine/auth.rs rename to litellm-rust/crates/host/src/machine/auth.rs index cf91458ae43..ba7e242e766 100644 --- a/litellm-rust/crates/core/src/machine/auth.rs +++ b/litellm-rust/crates/host/src/machine/auth.rs @@ -1,9 +1,8 @@ use std::sync::Arc; -use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -use litellm_host::route::Route; - use super::{HostChannel, MachineFault}; +use crate::route::Route; +use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; /// A route whose host can mint credentials on the call's behalf. pub trait TokenRoute: Route { diff --git a/litellm-rust/crates/host/src/machine.rs b/litellm-rust/crates/host/src/machine/mod.rs similarity index 91% rename from litellm-rust/crates/host/src/machine.rs rename to litellm-rust/crates/host/src/machine/mod.rs index 2942913f095..2c26db61582 100644 --- a/litellm-rust/crates/host/src/machine.rs +++ b/litellm-rust/crates/host/src/machine/mod.rs @@ -1,6 +1,12 @@ +mod auth; +mod route_machine; + use std::future::Future; use std::pin::Pin; +pub use auth::{HostTokenProvider, TokenRoute}; +pub use route_machine::{ExecuteFuture, HostChannel, MachineFault, RouteMachine}; + use crate::host::{HostOp, HostResult}; use crate::route::Route; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/host/src/machine/route_machine.rs similarity index 97% rename from litellm-rust/crates/core/src/machine/mod.rs rename to litellm-rust/crates/host/src/machine/route_machine.rs index ffcefc663ab..38a0b8bc16a 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/host/src/machine/route_machine.rs @@ -2,18 +2,16 @@ //! place, and turns the host operations that future requests into [`Machine`] steps. No //! task is spawned; dropping the machine drops the in-flight call. -mod auth; - use std::{future::Future, pin::Pin}; -pub use auth::{HostTokenProvider, TokenRoute}; -use litellm_host::{ +use tokio::sync::{mpsc, oneshot}; + +use super::{HostFailure, Interrupted, Machine, MachineStep, Step}; +use crate::{ event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, - machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, }; -use tokio::sync::{mpsc, oneshot}; /// The machine's own failures, distinct from anything the provider call reports. #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index c3f481d7d44..3061a9fe2b2 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -102,6 +102,17 @@ pub enum Error { Headers(#[from] crate::custom_httpx::http_handler::HeaderError), } +impl From for Error { + fn from(fault: litellm_host::machine::MachineFault) -> Self { + use litellm_host::machine::MachineFault; + Self::InvalidRequest(match fault { + MachineFault::Abandoned => "OCR host driver was abandoned".into(), + MachineFault::Protocol(message) => format!("OCR {message}"), + MachineFault::Mismatch => "invalid OCR host operation result".into(), + }) + } +} + impl From for Error { fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self { Self::RequestField { From 41873e2bc796bce93b74ec9431ae01b403ed46a3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 22:41:17 +0000 Subject: [PATCH 30/33] fix(rust): box messages response output Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/core/src/messages/mod.rs | 2 +- litellm-rust/crates/core/src/messages/route.rs | 5 +++-- .../crates/python-bridge/src/routes/messages/host.rs | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index c07a83c7e43..289f79109dd 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -34,7 +34,7 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result Ok(message), + MessagesOutput::Message(message) => Ok(*message), MessagesOutput::Streamed => Err(Error::Unsupported( "streamed responses need a streaming host", )), diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 3d607ba67f6..838b56fcb4b 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -48,7 +48,7 @@ impl MessagesCall { } pub enum MessagesOutput { - Message(AnthropicMessagesResponse), + Message(Box), /// Every chunk already reached the host through `Deliver`. Streamed, } @@ -174,7 +174,8 @@ async fn execute(host: MessagesHost) -> Result { raw: RawResponse { body: text.clone() }, }) .await?; - decode_response(request.config, &request.model, &text).map(MessagesOutput::Message) + decode_response(request.config, &request.model, &text) + .map(|message| MessagesOutput::Message(Box::new(message))) } /// Hands each upstream chunk to the caller as it arrives. A caller that stops reading diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index b590018bc4e..c1b3f59df58 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -150,7 +150,7 @@ impl RouteHost for MessagesRouteHost { MessagesOutput::Message(message) => py .import("litellm.rust_bridge.messages.route_host")? .getattr("response")? - .call1((to_py(py, &message)?,)) + .call1((to_py(py, message.as_ref())?,)) .map(Bound::unbind), MessagesOutput::Streamed => Ok(py.None()), } From ba6b22cf56ead7aba23d530b342da6d342230581 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:57:50 -0700 Subject: [PATCH 31/33] test(rust): isolate callback registries per hypothesis example Replace the module-level LATEST_EDITS list with per-example callback registry isolation, and import litellm names with from-imports in the legacy callback shim so the module uses one import style. Co-Authored-By: Claude Opus 5 --- litellm/rust_bridge/legacy_callbacks.py | 18 +++--- tests/test_litellm_rust/conftest.py | 63 +++---------------- tests/test_litellm_rust/ocr/test_callbacks.py | 20 +++--- tests/test_litellm_rust/support/isolation.py | 58 +++++++++++++++++ 4 files changed, 88 insertions(+), 71 deletions(-) create mode 100644 tests/test_litellm_rust/support/isolation.py diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index bac40442ce5..30aa1d97bfc 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -65,13 +65,17 @@ def setup( def check_limits(kwargs: Mapping[str, object]) -> None: - import litellm + from litellm import ( + BudgetExceededError, + _current_cost, # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor + max_budget, + num_retries_per_request, + ) from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit - current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor - if litellm.max_budget and current_cost > litellm.max_budget: - raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) - if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + if max_budget and _current_cost > max_budget: + raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget) + if max_retries_per_request_hit(kwargs, num_retries_per_request): raise RuntimeError("Max retries per request hit!") @@ -281,9 +285,9 @@ def is_internal_call() -> bool: def credential_list() -> list[CredentialItem]: - import litellm + from litellm import credential_list as credentials - return litellm.credential_list + return credentials def warn_unknown_credential(name: str, loaded: int) -> None: diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 4387ea2e2fd..1b6fcfa00db 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,10 +1,9 @@ import asyncio import os -from collections.abc import AsyncIterator, Generator, Iterator +from collections.abc import AsyncIterator, Generator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack, contextmanager -from types import ModuleType -from typing import Final, cast +from contextlib import ExitStack +from typing import Final import pytest import pytest_asyncio @@ -18,62 +17,20 @@ from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivate _parse_env_bool, ) from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service -CALLBACK_ATTRIBUTES: Final = ( - "callbacks", - "input_callback", - "success_callback", - "failure_callback", - "_async_input_callback", - "_async_success_callback", - "_async_failure_callback", -) - - -def _list_attribute(container: ModuleType, attribute: str) -> list[object]: - value: Final = getattr(container, attribute) - if not isinstance(value, list): - raise AssertionError(f"{container.__name__}.{attribute} is not a list") - return cast(list[object], value) - - -@contextmanager -def _isolated_list(container: ModuleType, attribute: str) -> Iterator[None]: - source: Final = _list_attribute(container, attribute) - original: Final = list(source) - source.clear() # mutable-ok: test isolation mutates global registries by design - try: - yield - finally: - source.clear() - source.extend(original) - setattr(container, attribute, source) - - -@contextmanager -def _rebound(container: object, attribute: str, value: object) -> Iterator[None]: - original: Final[object] = getattr(container, attribute) - setattr(container, attribute, value) - try: - yield - finally: - setattr(container, attribute, original) - @pytest_asyncio.fixture(autouse=True, loop_scope="function") async def isolate_ocr_test_state() -> AsyncIterator[None]: with ExitStack() as stack: - for attribute in CALLBACK_ATTRIBUTES: - stack.enter_context(_isolated_list(litellm, attribute)) - stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor - stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry - stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache - stack.enter_context(_rebound(_CONFIGURATION, "override", None)) + stack.enter_context(isolated_callback_registries()) + stack.enter_context(rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache + stack.enter_context(rebound(_CONFIGURATION, "override", None)) executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") - stack.enter_context(_rebound(litellm_logging, "executor", executor)) - stack.enter_context(_rebound(utils, "executor", executor)) - stack.enter_context(_rebound(thread_pool_executor, "executor", executor)) + stack.enter_context(rebound(litellm_logging, "executor", executor)) + stack.enter_context(rebound(utils, "executor", executor)) + stack.enter_context(rebound(thread_pool_executor, "executor", executor)) try: yield finally: diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 45b99d19d90..ac4a1a11a80 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -3,6 +3,8 @@ import copy import gc import queue import threading +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import pytest @@ -13,6 +15,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, @@ -307,18 +310,13 @@ JSON_VALUES: Final = st.recursive( ) -LATEST_EDITS: Final[list[dict[str, object]]] = [] - - -class ApplyLatestEdits(CustomLogger): - """Registrations can outlive one hypothesis example, so every instance applies the current example's edits.""" - - def __init__(self, latest: list[dict[str, object]]) -> None: +class ApplyEdits(CustomLogger): + def __init__(self, edits: Mapping[str, object]) -> None: super().__init__() - self.latest = latest + self.edits: Final = edits def log_pre_api_call(self, model, messages, kwargs): - request_body(kwargs).update(copy.deepcopy(self.latest[-1])) + request_body(kwargs).update(copy.deepcopy(dict(self.edits))) @settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) @@ -327,9 +325,9 @@ def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_lef ocr_server: RecordingServer, edits: dict[str, object] ) -> None: ocr_server.expected_requests = None - LATEST_EDITS.append(edits) - call_native_ocr_with_callbacks(ocr_server, [ApplyLatestEdits(LATEST_EDITS)]) + with isolated_callback_registries(): + call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))]) assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} diff --git a/tests/test_litellm_rust/support/isolation.py b/tests/test_litellm_rust/support/isolation.py new file mode 100644 index 00000000000..f98ce4843a8 --- /dev/null +++ b/tests/test_litellm_rust/support/isolation.py @@ -0,0 +1,58 @@ +from collections.abc import Generator +from contextlib import ExitStack, contextmanager +from types import ModuleType +from typing import Final, cast + +import litellm +from litellm import utils +from litellm.litellm_core_utils import litellm_logging + +CALLBACK_ATTRIBUTES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _list_attribute(container: ModuleType, attribute: str) -> list[object]: + value: Final = getattr(container, attribute) + if not isinstance(value, list): + raise AssertionError(f"{container.__name__}.{attribute} is not a list") + return cast(list[object], value) + + +@contextmanager +def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]: + source: Final = _list_attribute(container, attribute) + original: Final = list(source) + source.clear() # mutable-ok: test isolation mutates global registries by design + try: + yield + finally: + source.clear() + source.extend(original) + setattr(container, attribute, source) + + +@contextmanager +def rebound(container: object, attribute: str, value: object) -> Generator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@contextmanager +def isolated_callback_registries() -> Generator[None]: + with ExitStack() as stack: + for attribute in CALLBACK_ATTRIBUTES: + stack.enter_context(_isolated_list(litellm, attribute)) + stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor + stack.enter_context(rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry + yield From 139445179a8a2fd1804cbeba1154e77b4e044fb6 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 18 Sep 2026 22:59:32 +0000 Subject: [PATCH 32/33] ci: remove the dead Agent Shin triage workflows and scripts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/_agent_shin_actions.py | 50 - .github/scripts/agent_shin_shared.py | 211 -- .github/scripts/close_low_quality_prs.py | 573 ----- .github/scripts/triage-requirements.txt | 282 --- .github/scripts/triage_with_llm.py | 1797 -------------- .github/workflows/close_low_quality_prs.yml | 92 - .../create_daily_oss_agent_shin_branch.yml | 28 - .github/workflows/triage_reconsider.yml | 172 -- .../test_github_close_low_quality_prs.py | 856 ------- tests/test_litellm/test_github_review_gate.py | 524 ---- .../test_github_triage_with_llm.py | 2134 ----------------- .../test_github_triage_workflows.py | 264 -- 12 files changed, 6983 deletions(-) delete mode 100644 .github/scripts/_agent_shin_actions.py delete mode 100644 .github/scripts/agent_shin_shared.py delete mode 100644 .github/scripts/close_low_quality_prs.py delete mode 100644 .github/scripts/triage-requirements.txt delete mode 100644 .github/scripts/triage_with_llm.py delete mode 100644 .github/workflows/close_low_quality_prs.yml delete mode 100644 .github/workflows/create_daily_oss_agent_shin_branch.yml delete mode 100644 .github/workflows/triage_reconsider.yml delete mode 100644 tests/test_litellm/test_github_close_low_quality_prs.py delete mode 100644 tests/test_litellm/test_github_review_gate.py delete mode 100644 tests/test_litellm/test_github_triage_with_llm.py delete mode 100644 tests/test_litellm/test_github_triage_workflows.py diff --git a/.github/scripts/_agent_shin_actions.py b/.github/scripts/_agent_shin_actions.py deleted file mode 100644 index b3d1ff055b3..00000000000 --- a/.github/scripts/_agent_shin_actions.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Dry-run wrapper(s) around Agent Shin GitHub mutations. - -The rollout scripts currently need only one mutation wrapped, so this module -exposes a single ``maybe_post_comment`` helper. It takes a ``dry_run: bool`` -keyword argument and the body is intentionally trivial: - - if dry_run: - print(...) # log what we would do, return - return - real_mutation(...) # otherwise, actually do it - -That shape means a dry-run preview differs from the real run in exactly one -line per side effect: the call site. So when you `python3 script.py` locally -without ``--close``, you can be confident the actions printed are the ones the -GitHub Action would have performed (modulo ordering on retry/error paths, -which are deliberately simple). Any further mutation a rollout script needs -should get the same ``maybe_*`` treatment instead of calling the raw -``triage_with_llm`` mutation directly. - -Importing from this module pulls in the real mutation from ``triage_with_llm`` -— call sites in the rollout scripts should NEVER import ``post_comment`` -directly; that would skip the dry-run gate and is the bug class this module -exists to prevent. -""" - -from __future__ import annotations - -import sys -import textwrap - -# Import the module itself rather than the bare names so monkeypatching -# `triage_with_llm.post_comment` (or any of the other mutations) in tests is -# reflected here — `from triage_with_llm import post_comment` would bind the -# original function to a local name and bypass the patch, defeating the whole -# point of these wrappers. -import triage_with_llm - - -def _log(line: str) -> None: - """Print a single dry-run line to stdout (one log statement per side effect).""" - print(line, file=sys.stdout, flush=True) - - -def maybe_post_comment(repo: str, number: int, body: str, *, dry_run: bool) -> None: - """Post a comment on ``repo#number`` — or, in dry-run, log what we would post.""" - if dry_run: - _log(f"[DRY RUN] comment {repo}#{number}:") - _log(textwrap.indent(body, " ")) - return - triage_with_llm.post_comment(repo, number, body) diff --git a/.github/scripts/agent_shin_shared.py b/.github/scripts/agent_shin_shared.py deleted file mode 100644 index 8f3dc3c2322..00000000000 --- a/.github/scripts/agent_shin_shared.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Constants and helpers shared by Agent Shin's triage scripts. - -Both `triage_with_llm.py` (the LLM-judge entrypoint) and -`close_low_quality_prs.py` (the daily Greptile-score sweep) need to -agree on the same notions of: - - * What counts as a Greptile-authored review comment - (``GREPTILE_BOT_LOGINS``) and how to extract a confidence score from - its body (``SCORE_PATTERN`` / :func:`extract_greptile_score`). - * How long the 2-hour grace window is (``GRACE_PERIOD_SECONDS``) and - the HTML marker stamped into a grace-warning comment so the *other* - script can see "Agent Shin already warned" and behave accordingly - (``GRACE_COMMENT_MARKER``). - * Who Agent Shin is on GitHub (``AGENT_SHIN_DEFAULT_BOT_LOGIN``). - * How GitHub-style ISO-8601 timestamps round-trip into timezone-aware - :class:`datetime.datetime` (:func:`parse_iso8601`). - -Keeping these in one module means a future change (new Greptile output -format, a longer grace window, a new allowlisted account) is a single edit -instead of two — the original split version had to call out in comments -that the two copies "must stay in sync" precisely because nothing -enforced it. -""" - -from __future__ import annotations - -import datetime as dt -import json -import os -import re -import subprocess -from typing import Iterable - -GREPTILE_BOT_LOGINS = frozenset({"greptile-apps", "greptile-apps[bot]"}) - -SCORE_PATTERN = re.compile( - r"confidence\s*score\s*[:\-]?\s*(\d+)\s*/\s*5", - re.IGNORECASE, -) - -GRACE_COMMENT_MARKER = "" - -# Hidden HTML marker stamped on every Agent Shin auto-close comment (the LLM -# judge's grace/review-gate close and the daily Greptile sweep's close). -# `was_closed_by_agent_shin` requires this marker — not just the closing actor — -# before `@agent-shin reconsider` may reopen, because the `github-actions[bot]` -# identity is shared with every other workflow in the repo and is not unique to -# Agent Shin. Both close paths must stamp it or the reconsider path silently -# rejects the contributor. -AGENT_SHIN_CLOSE_MARKER = "" - -# 2 hours between the grace warning and the auto-close. Short enough to -# dogfood the "fix it before it closes" loop in one sitting; bump back up -# (e.g. 86400 for a day) for the public rollout. -GRACE_PERIOD_SECONDS = 7200 - -AGENT_SHIN_DEFAULT_BOT_LOGIN = "github-actions[bot]" - - -def _logins(*names: str) -> frozenset[str]: - """Build a login set normalized for case-insensitive membership checks. - - Callers compare via ``login.lower() in ``, so the stored values - must be lowercase. Normalizing here lets the literals keep each - account's canonical GitHub casing (e.g. ``SwiftWinds``) for - readability without breaking the lookup. - """ - return frozenset(name.lower() for name in names) - - -# Dogfood rollout gate. While this set is non-empty, Agent Shin acts ONLY on -# PRs/issues authored by these logins and skips everyone else. For an -# allowlisted author the usual internal/external classification is bypassed, so -# an internal account (e.g. a maintainer's own work login) still gets triaged -# while the bot is being tested on a small set of accounts. Empty the set to -# lift the restriction and restore full triage for the public rollout. Logins -# are compared case-insensitively. -ALLOWLIST_LOGINS = _logins("mateo-berri", "SwiftWinds") - -# `gh {pr,issue} list` has no "fetch everything" flag — `--limit` is the only -# control and it defaults to 30. Pass a ceiling far above any realistic open -# backlog (low thousands today) so gh paginates the API until the queue is -# exhausted rather than silently truncating. The bulk sweeps MUST see the whole -# backlog: gh lists newest-first, so a low cap drops the *oldest* PRs/issues — -# exactly the stale ones a low-quality sweep is meant to catch. -GH_LIST_ALL_LIMIT = 100_000 - - -def extract_greptile_score(comments: Iterable[dict]) -> tuple[int, dict] | None: - """Return (score, comment) for the most recent Greptile-authored comment - that contains a "Confidence Score: X/5". Returns None if no such comment. - - "Most recent" is determined by the comment's `updated_at` (falling back to - `created_at`), so re-reviews override earlier passes. - """ - candidates: list[tuple[str, int, dict]] = [] - for comment in comments: - user = (comment.get("user") or {}).get("login", "") - if user not in GREPTILE_BOT_LOGINS: - continue - body = comment.get("body") or "" - match = SCORE_PATTERN.search(body) - if not match: - continue - score = int(match.group(1)) - timestamp = comment.get("updated_at") or comment.get("created_at") or "" - candidates.append((timestamp, score, comment)) - - if not candidates: - return None - - candidates.sort(key=lambda triple: triple[0]) - _, score, comment = candidates[-1] - return score, comment - - -def parse_iso8601(value: str) -> dt.datetime: - """Parse a GitHub ISO-8601 timestamp into a timezone-aware datetime.""" - return dt.datetime.fromisoformat(value.replace("Z", "+00:00")) - - -def gh(*args: str) -> str: - """Run a `gh` CLI command and return stdout. Raises on non-zero exit. - - Shared by both Agent Shin entrypoints so a future change here - (timeout handling, logging, retry on transient failures) only needs - to be made once. - """ - result = subprocess.run( - ["gh", *args], - capture_output=True, - text=True, - check=True, - ) - return result.stdout - - -def list_open_items(kind: str, *, repo: str | None, fields: str) -> list[dict]: - """Return EVERY open PR (``kind="pr"``) or issue (``kind="issue"``) in ``repo``. - - Wraps ``gh {pr,issue} list`` with ``--limit GH_LIST_ALL_LIMIT`` so the full - backlog is fetched instead of the default 30 (or any other arbitrary cap). - Both bulk sweeps — the daily Greptile closer and the one-shot rollout - heads-up — rely on this seeing the whole queue, including the oldest items. - - ``fields`` is the comma-separated ``--json`` field list the caller needs - (e.g. ``"number"`` for the rollout, the full set for the closer). - """ - if kind not in ("pr", "issue"): - raise ValueError(f"kind must be 'pr' or 'issue', got {kind!r}") - repo_args = ["--repo", repo] if repo else [] - raw = gh( - kind, - "list", - "--state", - "open", - "--limit", - str(GH_LIST_ALL_LIMIT), - "--json", - fields, - *repo_args, - ) - return json.loads(raw) - - -def seconds_since_latest_marker_comment( - comments: Iterable[dict], - *, - marker: str, - bot_login: str | None = None, - now: dt.datetime | None = None, -) -> float | None: - """Return seconds since the bot's most recent comment containing ``marker``. - - Filters comments by author so a contributor who quotes the HTML - marker (e.g. via GitHub's "Quote reply" feature, which preserves - HTML comments in the raw markdown of the quoted text) is not - mistaken for a bot warning — that would silently reset cooldown - timers and suppress legitimate notifications. - - ``bot_login`` defaults to the `AGENT_SHIN_BOT_LOGIN` env override or - ``AGENT_SHIN_DEFAULT_BOT_LOGIN`` so callers normally don't need to - pass it. ``now`` is injectable for tests / callers (like the daily - sweep) that want every age calculation pinned to one snapshot. - """ - expected_login = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - latest: dt.datetime | None = None - for comment in comments: - author = ((comment.get("user") or {}).get("login") or "").lower() - if author != expected_login: - continue - body = comment.get("body") or "" - if marker not in body: - continue - created = comment.get("created_at") - if not created: - continue - try: - ts = parse_iso8601(created) - except ValueError: - continue - if latest is None or ts > latest: - latest = ts - if latest is None: - return None - reference = now if now is not None else dt.datetime.now(dt.timezone.utc) - return (reference - latest).total_seconds() diff --git a/.github/scripts/close_low_quality_prs.py b/.github/scripts/close_low_quality_prs.py deleted file mode 100644 index 7b9bbb579e3..00000000000 --- a/.github/scripts/close_low_quality_prs.py +++ /dev/null @@ -1,573 +0,0 @@ -#!/usr/bin/env python3 -""" -Auto-close low-quality pull requests. - -Closes open PRs (including drafts, regardless of age) that satisfy ALL of: - 1. Have a Greptile (`greptile-apps`) review comment whose latest - "Confidence Score: X/5" is below the configured threshold (default: 4). - 2. Are authored by an external OSS contributor (internal BerriAI - contributors are exempt). - 3. Do not carry an opt-out label (default: "do not close"). - -`--min-age-days` is retained as an opt-in safety net for one-off backfill -runs (default: 0). The team's intent is that the count of open PRs equals -the count of PRs internal collaborators need to action on, so neither age -nor draft status acts as a free pass. - -For each match, the script posts an explanatory comment and closes the PR. -Because OSS contributors *cannot* reopen a PR closed by the bot/maintainer -(GitHub limitation), the close-comment instructs them to push their fixes -and **open a fresh PR**, or to comment `@agent-shin reconsider` on the -closed PR to have the LLM judge re-evaluate (and reopen on pass). - -Requires the `gh` CLI to be authenticated. - -Usage examples: - # Dry run (default) - prints what would be closed - python3 close_low_quality_prs.py - - # Actually close matching PRs - python3 close_low_quality_prs.py --close - - # Restrict to PRs at least N days old (one-off backfill safety net) - python3 close_low_quality_prs.py --min-age-days 7 --min-score 4 --close -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import json -import os -import subprocess -import sys -from typing import Iterable - -# Add this script's directory to `sys.path` so the sibling -# `agent_shin_shared` module is importable when the script is invoked -# directly (e.g. `python3 .github/scripts/close_low_quality_prs.py ...`). -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from agent_shin_shared import ( # noqa: E402 -- sys.path adjusted above - AGENT_SHIN_CLOSE_MARKER, - ALLOWLIST_LOGINS, - GRACE_COMMENT_MARKER, - GRACE_PERIOD_SECONDS, - GREPTILE_BOT_LOGINS, - SCORE_PATTERN, - extract_greptile_score, - gh, - list_open_items, - parse_iso8601, - seconds_since_latest_marker_comment, -) - -# `GREPTILE_BOT_LOGINS` and `SCORE_PATTERN` (Greptile's GitHub App login -# variants and the "Confidence Score: X/5" regex) are imported from -# `agent_shin_shared` so the LLM judge in `triage_with_llm.py` and this -# daily Greptile sweep read the score through the same set of logins -# and the same regex. - -# `author_association` values for internal BerriAI contributors who should be -# exempt from auto-triage. -INTERNAL_AUTHOR_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"}) - -# Default labels that exempt a PR from auto-close. Defined at module scope (not -# as a mutable argparse default) so that `--optout-label foo` REPLACES the -# defaults instead of appending to them — the argparse `action="append"` + -# `default=[...]` combination silently mutates the shared default list. -DEFAULT_OPTOUT_LABELS = ("do not close", "keep open", "wip") - -# `GRACE_COMMENT_MARKER` (HTML marker appended to grace-period warning -# comments — used by either script to recognize that a warning was -# already posted) and `GRACE_PERIOD_SECONDS` (length of the grace -# period between the warning and the actual auto-close, 2 hours) are -# imported from `agent_shin_shared` so the Agent Shin LLM judge and -# this daily Greptile sweep agree on the same marker and duration. - - -def fetch_open_prs(repo: str | None) -> list[dict]: - """Fetch all open PRs (number, createdAt, isDraft, labels, author). - - Includes drafts: `gh pr list --state open` returns both ready-for-review - and draft PRs by default. This is the desired behavior — drafts are not - a free pass; the internal-collaborator open-PR queue should reflect every - PR that needs human attention regardless of draft status. - """ - fields = "number,title,createdAt,isDraft,labels,author,url" - return list_open_items("pr", repo=repo, fields=fields) - - -def fetch_pr_author_association(pr_number: int, repo: str | None) -> str: - """Return the GitHub `author_association` for a PR, uppercase. - - Values: OWNER, MEMBER, COLLABORATOR, CONTRIBUTOR, FIRST_TIME_CONTRIBUTOR, - FIRST_TIMER, MANNEQUIN, NONE. Returns "" on lookup failure. - """ - endpoint = ( - f"repos/{repo}/pulls/{pr_number}" - if repo - else f"repos/{{owner}}/{{repo}}/pulls/{pr_number}" - ) - try: - data = json.loads(gh("api", endpoint)) - except subprocess.CalledProcessError: - return "" - return (data.get("author_association") or "").upper() - - -def is_external_pr_author(pr: dict, repo: str | None) -> bool: - """Return True if the PR author is an external OSS contributor. - - Internal = `OWNER` / `MEMBER` / `COLLABORATOR` association, or a bot login. - """ - login = ((pr.get("author") or {}).get("login") or "").lower() - if login.endswith("[bot]") or login in {"dependabot", "github-actions"}: - return False - association = fetch_pr_author_association(pr["number"], repo) - # Fail-safe: if the API lookup failed (empty string), treat the author as - # internal so we don't auto-close their PR. Auto-close is destructive, so - # an unknown association should never make a PR eligible for closing. - if not association or association in INTERNAL_AUTHOR_ASSOCIATIONS: - return False - return True - - -def fetch_pr_comments(pr_number: int, repo: str | None) -> list[dict]: - """Fetch issue-level comments on a PR (where Greptile posts its summary).""" - endpoint = ( - f"repos/{repo}/issues/{pr_number}/comments?per_page=100" - if repo - else f"repos/{{owner}}/{{repo}}/issues/{pr_number}/comments?per_page=100" - ) - raw = gh("api", "--paginate", endpoint) - comments: list[dict] = [] - for line in raw.strip().splitlines(): - line = line.strip() - if not line: - continue - try: - parsed = json.loads(line) - except json.JSONDecodeError: - # A malformed line should not blow up the whole sweep. Skip and - # carry on so the remaining PRs in this run still get evaluated. - continue - if isinstance(parsed, list): - comments.extend(parsed) - else: - comments.append(parsed) - return comments - - -def has_optout_label(pr: dict, optout_labels: set[str]) -> bool: - labels = {label.get("name", "").lower() for label in pr.get("labels", [])} - return bool(labels & {lbl.lower() for lbl in optout_labels}) - - -def seconds_since_last_grace_warning( - comments: Iterable[dict], - *, - bot_login: str | None = None, - now: dt.datetime | None = None, -) -> float | None: - """Return seconds since the bot's most recent grace-period warning, or - None if no such warning has ever been posted on this PR. - - Thin wrapper over - `agent_shin_shared.seconds_since_latest_marker_comment` — the - centralized helper handles the bot-author filter, marker match, - timestamp parsing, and `now` injection. Keeping this wrapper - preserves the closer's "already-fetched comments + injectable now" - interface so callers (and tests) don't need to change. - """ - return seconds_since_latest_marker_comment( - comments, - marker=GRACE_COMMENT_MARKER, - bot_login=bot_login, - now=now, - ) - - -def format_grace_warning_comment(score: int, threshold: int) -> str: - """Comment posted on the FIRST low-Greptile-score detection — gives - the contributor a 2-hour grace window before the auto-close fires on - the next daily cron run. - - Mirrors `format_grace_warning_pr_comment` in - `triage_with_llm.py` in spirit (2-hour grace + escape hatches), but - framed around Greptile's confidence score instead of the LLM judge's - rubric since the close trigger here is the Greptile signal. - """ - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository.\n" - "\n" - "Heads up: Greptile's most recent review scored this PR " - f"**{score}/5**, below our merge bar of **{threshold}/5**.\n" - "\n" - "If the score isn't lifted in the next **2 hours**, I'll auto-close this PR. That's " - "**not** us saying the change isn't worthwhile. We want the open-PR list to mirror " - "what a maintainer can act on *right now*, so contributors like you don't get lost in " - "a backlog. Take your time; everything below still works after the close.\n" - "\n" - "**During the grace period:** push fixes that address Greptile's feedback, then comment " - "`@greptileai` to request a fresh review. If " - f"the new score is **{threshold}/5 or higher**, the PR stays open and no further " - "action is needed on your side.\n" - "\n" - "**If the PR does get auto-closed in 2 hours, you still have an easy recovery path:**\n" - "\n" - "- Comment `@greptileai` to request a fresh review. **This still works even after " - f"the PR is closed**, and a score of {threshold}/5 or higher is one of the signals " - "that lifts the PR back into the review queue. A low Greptile score isn't a blocker.\n" - "- Comment `@agent-shin reconsider` after pushing fixes; I'll re-run the rubric and " - "reopen the PR if both gates (description rubric + Greptile score) now pass.\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -def post_grace_warning( - pr: dict, - score: int, - threshold: int, - repo: str | None, - dry_run: bool, -) -> None: - """Post the 2-hour grace-period warning comment on `pr`. - - The warning carries `GRACE_COMMENT_MARKER` so subsequent runs can - detect that the contributor has already been told about the - pending close. Does NOT close the PR — the close happens on the - next eligible run after `GRACE_PERIOD_SECONDS` elapses (handled - by `close_pr`). - """ - pr_number = pr["number"] - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would post grace warning to PR #{pr_number} " - f"(greptile={score}/5): {pr['title']}" - ) - return - - comment_body = format_grace_warning_comment(score, threshold) - gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args) - print(f" Posted grace warning on PR #{pr_number} (greptile={score}/5)") - - -def format_close_comment(score: int, threshold: int) -> str: - """Comment posted when a low-Greptile-score PR is auto-closed. - - Carries `AGENT_SHIN_CLOSE_MARKER` so the `@agent-shin reconsider` path - (guarded by `was_closed_by_agent_shin`) recognizes this as an Agent Shin - close and is allowed to reopen the PR once it passes again; without the - marker that recovery path the comment advertises silently rejects the - contributor. - """ - score_sentence = ( - f"Greptile's most recent review scored this PR **{score}/5**, below " - f"our merge bar of **{threshold}/5**, and the 2-hour grace period since " - "the warning has elapsed.\n\n" - ) - return ( - f"Closing as part of automated PR triage.\n\n" - f"{score_sentence}" - "We close low-confidence PRs aggressively to keep the review queue " - "manageable for maintainers and contributors alike. **This is not a " - "rejection of the idea.** To bring this back:\n\n" - "1. Push the fixes that address Greptile's feedback (continue using " - "your existing branch is fine).\n" - "2. **Open a new PR** with the updated branch. Greptile will review " - "it again, and if it scores " - f"**{threshold}/5 or higher** a maintainer will take another look.\n\n" - "_Why open a new PR instead of reopening this one?_ GitHub does not " - "let external contributors reopen a PR that was closed by a bot or " - "maintainer, so a fresh PR is the most reliable path forward. If you " - "would prefer this exact PR re-evaluated, comment " - "`@agent-shin reconsider` once you've pushed the fixes; Agent Shin " - "will re-run triage and reopen this PR if it now meets the bar. " - "You can also comment `@greptileai` to request a fresh Greptile " - "review; that works **even after the PR is closed**.\n\n" - "Thanks for contributing to LiteLLM. We know auto-closures can sting; " - "the goal is to keep the project healthy, not to dismiss your work." - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def close_pr( - pr: dict, - score: int, - threshold: int, - age_days: int, - repo: str | None, - dry_run: bool, - label: str | None, -) -> None: - """Post the explanatory comment and close the PR.""" - pr_number = pr["number"] - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would close PR #{pr_number} " - f"(age={age_days}d, greptile={score}/5): {pr['title']}" - ) - return - - comment_body = format_close_comment(score, threshold) - gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args) - - if label: - try: - gh("pr", "edit", str(pr_number), "--add-label", label, *repo_args) - except subprocess.CalledProcessError as exc: - stderr = (exc.stderr or "").strip() - print(f" warn: failed to add label '{label}' to #{pr_number}: {stderr}") - - gh("pr", "close", str(pr_number), *repo_args) - print(f" Closed PR #{pr_number} (greptile={score}/5, age={age_days}d)") - - -def evaluate_pr( - pr: dict, - now: dt.datetime, - min_age_days: int, - min_score: int, - repo: str | None, - optout_labels: set[str], - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> tuple[str, int | None, int | None]: - """Decide what to do with `pr` on this triage run. - - Returns (action, score_or_none, age_days_or_none) where action is one of: - "skip-too-young", "skip-optout-label", "skip-not-allowlisted", - "skip-internal", "skip-no-greptile-score", "skip-score-ok", - "warn-grace", "skip-in-grace-period", or "close". - - Drafts are NOT skipped — the goal is "open PR count == PRs internal - collaborators need to action on", and a draft that Greptile scored <4/5 - is still in that queue. Authors can opt out via the `wip` label (see - `DEFAULT_OPTOUT_LABELS`) if they need to keep a long-lived draft open. - - Grace-period semantics: the first time a PR fails the rubric, the - action is `warn-grace` — the caller should post a warning comment but - NOT close the PR. On a subsequent run, if the warning is still less - than `GRACE_PERIOD_SECONDS` old AND the PR still fails, the action is - `skip-in-grace-period`. Once the warning ages out and the rubric is - still failing, the action is `close`. - """ - if has_optout_label(pr, optout_labels): - return ("skip-optout-label", None, None) - - created = parse_iso8601(pr["createdAt"]) - age_days = (now - created).days - # `min_age_days` defaults to 0 (close as soon as Greptile scores low). - # Set a positive value via --min-age-days for one-off backfill runs that - # want to skip very-young PRs. - if min_age_days > 0 and age_days < min_age_days: - return ("skip-too-young", None, age_days) - - # While the allowlist is active it is the sole author gate: only those - # logins are acted on and the external-only restriction is bypassed for - # them. Otherwise auto-close only external OSS contributors — internal - # contributors (BerriAI org members) handle their own backlog. - login = ((pr.get("author") or {}).get("login") or "").lower() - if allowlist: - if login not in allowlist: - return ("skip-not-allowlisted", None, age_days) - elif not is_external_pr_author(pr, repo): - return ("skip-internal", None, age_days) - - comments = fetch_pr_comments(pr["number"], repo) - extraction = extract_greptile_score(comments) - if extraction is None: - return ("skip-no-greptile-score", None, age_days) - - score, _ = extraction - if score >= min_score: - return ("skip-score-ok", score, age_days) - - grace_age = seconds_since_last_grace_warning(comments, now=now) - if grace_age is None: - return ("warn-grace", score, age_days) - if grace_age < GRACE_PERIOD_SECONDS: - return ("skip-in-grace-period", score, age_days) - - return ("close", score, age_days) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--repo", - type=str, - default=None, - help="Repository (owner/repo). Auto-detected if omitted.", - ) - parser.add_argument( - "--min-age-days", - type=int, - default=0, - help=( - "Minimum age (in days) before a PR is eligible. Default 0 = " - "close as soon as Greptile flags it. Set a positive value for " - "one-off backfill runs that want to spare very-young PRs." - ), - ) - parser.add_argument( - "--min-score", - type=int, - default=4, - choices=range(1, 6), - help="Greptile score below which a PR is closed (default: 4 -> closes <4/5).", - ) - parser.add_argument( - "--optout-label", - action="append", - default=None, - help=( - "Label(s) that exempt a PR from auto-close. Repeat to add more. " - "Case-insensitive. When omitted, defaults to " - f"{list(DEFAULT_OPTOUT_LABELS)!r}; passing this flag REPLACES the " - "defaults (argparse `append` with a mutable default would append " - "instead, which we explicitly avoid)." - ), - ) - parser.add_argument( - "--close-label", - type=str, - default=None, - help=( - "Optional label to add to PRs that get auto-closed " - "(e.g. 'auto-closed-low-quality'). Must already exist on the repo." - ), - ) - parser.add_argument( - "--close", - action="store_true", - help="Actually close matching PRs (default is dry-run).", - ) - parser.add_argument( - "--limit", - type=int, - default=None, - help="Maximum number of PRs to close in one run (safety net).", - ) - args = parser.parse_args() - - dry_run = not args.close - if dry_run: - print("=== DRY RUN MODE (pass --close to actually close PRs) ===\n") - - print("Fetching open PRs...") - prs = fetch_open_prs(args.repo) - print(f"Found {len(prs)} open PRs.\n") - - now = dt.datetime.now(dt.timezone.utc) - optout_labels = set(args.optout_label or DEFAULT_OPTOUT_LABELS) - - closed = 0 - summary = { - "close": 0, - "warn-grace": 0, - "skip-in-grace-period": 0, - "skip-too-young": 0, - "skip-optout-label": 0, - "skip-not-allowlisted": 0, - "skip-internal": 0, - "skip-no-greptile-score": 0, - "skip-score-ok": 0, - } - - # `warned` tracks grace-warning comments posted in this run so the - # `--limit` safety net bounds *all* destructive write actions, not - # just closures. Without this cap, a backlog of PRs failing the - # threshold simultaneously could flood contributors with comments. - warned = 0 - for pr in sorted(prs, key=lambda p: p["createdAt"]): - try: - action, score, age_days = evaluate_pr( - pr, - now, - args.min_age_days, - args.min_score, - args.repo, - optout_labels, - ) - summary[action] = summary.get(action, 0) + 1 - - if action == "warn-grace": - assert score is not None - print( - f"#{pr['number']}: \"{pr['title']}\" " - f"(age={age_days}d, greptile={score}/5) -> warn-grace" - ) - post_grace_warning( - pr, - score=score, - threshold=args.min_score, - repo=args.repo, - dry_run=dry_run, - ) - if not dry_run: - warned += 1 - if args.limit is not None and (warned + closed) >= args.limit: - print( - f"\nReached --limit={args.limit} " - f"(closed={closed}, warned={warned}); stopping." - ) - break - continue - - if action != "close": - continue - - assert score is not None and age_days is not None - print( - f"#{pr['number']}: \"{pr['title']}\" " - f"(age={age_days}d, greptile={score}/5) -> close" - ) - close_pr( - pr, - score=score, - threshold=args.min_score, - age_days=age_days, - repo=args.repo, - dry_run=dry_run, - label=args.close_label, - ) - - if not dry_run: - closed += 1 - if args.limit is not None and (warned + closed) >= args.limit: - print( - f"\nReached --limit={args.limit} " - f"(closed={closed}, warned={warned}); stopping." - ) - break - except Exception as exc: # noqa: BLE001 - per-PR errors don't abort the sweep - summary["error"] = summary.get("error", 0) + 1 - print( - f"!! PR #{pr.get('number')}: {exc}", - file=sys.stderr, - ) - continue - - print("\n=== Summary ===") - for key, value in summary.items(): - print(f" {key:28s} {value}") - if dry_run: - print(f"\nTotal would close: {summary['close']}") - else: - print(f"\nTotal closed: {closed}") - print( - f"Total {'would warn (grace)' if dry_run else 'warned (grace)'}: " - f"{summary['warn-grace']}" - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/scripts/triage-requirements.txt b/.github/scripts/triage-requirements.txt deleted file mode 100644 index a18f05fbb95..00000000000 --- a/.github/scripts/triage-requirements.txt +++ /dev/null @@ -1,282 +0,0 @@ -# Hash-pinned dependency set for the Agent Shin triage scripts. -# Installed in privileged triage workflows, so every package is pinned to an -# exact version with SHA-256 hashes and installed with pip --require-hashes. -# -# Regenerate after bumping openai: -# echo 'openai==' \ -# | uv pip compile - --generate-hashes --python-version 3.12 \ -# --no-annotate --no-header -o .github/scripts/triage-requirements.txt - -annotated-types==0.7.0 \ - --hash=sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53 \ - --hash=sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89 -anyio==4.14.0 \ - --hash=sha256:b47c1f9ccf73e67021df785332508f99379c68fa7d0684e8e3492cb1d4b23f89 \ - --hash=sha256:dd9b7a2a9799ed6552fde617b2c5df02b7fdd7d88392fc48101e51bae46164d9 -certifi==2026.6.17 \ - --hash=sha256:024c88eeec92ca068db80f02b8b07c9cef7b9fe261d1d535abfd5abd6f6af432 \ - --hash=sha256:2227dcbaafe0d2f59279d1762ddddc37783ed4354594f194ffc31d20f41fc3db -distro==1.9.0 \ - --hash=sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed \ - --hash=sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2 -h11==0.16.0 \ - --hash=sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1 \ - --hash=sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86 -httpcore==1.0.9 \ - --hash=sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55 \ - --hash=sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8 -httpx==0.28.1 \ - --hash=sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc \ - --hash=sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad -idna==3.18 \ - --hash=sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2 \ - --hash=sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848 -jiter==0.15.0 \ - --hash=sha256:01a8222cf05ab1128e239421156c207949808acaaea2bdfd33130ae666786e86 \ - --hash=sha256:032396229564bca02440396bd327710719f724f5e7b7e9f7a8eb3faa4a2c2281 \ - --hash=sha256:04b400bbf8c9efb03d9bdd976475c919c1d85593b04b9fff7ae234065daf87ae \ - --hash=sha256:05906b93d72f03339e6bb7cf8dc10ebda64a0266126eed6beba79e20abcf5fd4 \ - --hash=sha256:066f8f33f18b2419cd8213b2436fa7fbc9c499f315971cfa3ce1f9820c001b1b \ - --hash=sha256:0ab068bce62a45aa3e7367eceaffb5dde60b7eb853be8dece45132e3d0ff4879 \ - --hash=sha256:0be6f5ad41a809f303f416d17cec92a7a725902fb9b4f3de3d19362ac0ef8554 \ - --hash=sha256:0e90a1c315a0226ec822d973817967f9223b7701546c8c2a7913e7ab0926294d \ - --hash=sha256:0f862193b8696249d22ec433e85fd2ab0ad9596bc3e45e6c0bc55e8aeba97be2 \ - --hash=sha256:1303d4d68a9b051ea90502402063ecf3807da00ad2affa19ca1ae3b90b3c5f67 \ - --hash=sha256:144f8e72cb53dab146347b91cceac01f5481237f2b93b4a339a1ee8f8878b67c \ - --hash=sha256:182226cbc930c9fab81bc2e41a4da672f89539906dadb05e75670ac07b94f71f \ - --hash=sha256:1c11465f97e2abf45a014b83b730222f8f1c5335e802c7055a67d50de6f1f4e3 \ - --hash=sha256:1c15024a3d892223b18f597c86d59387249dc396590844ce6b9f6131d1093bae \ - --hash=sha256:1d54fb5b31dea401a41af3f8a7d2512e9b6a6a005491e6166c7e4ffab9639a9c \ - --hash=sha256:25ffbe229aa8cd98c28879d8aa1a6e34ae77992ab984a65fba800859dab16269 \ - --hash=sha256:2a77aadd57cac1682e4401a72724d2796d89a4ba129b1a5812aa94ee480826eb \ - --hash=sha256:2ae901f3a55bfafdde31d289590fa25e3245735a2b1e8c7cc15871710a002871 \ - --hash=sha256:2b0074e2f56eb2dacca1689760fd2852a068f85a0547a157b82cb4cafeb6768b \ - --hash=sha256:2c8aea7781d2a372227871de4e1a1332aa96f5a89fd76c5e835dafdbad102887 \ - --hash=sha256:2c9cb907439d20bd0c7d7565ca01ee52234203208433749bae5b516907526928 \ - --hash=sha256:2fb6a5d26af81fc0f00f9360a891e05cf755e149bba391c4d563adc54812973d \ - --hash=sha256:2fd73e3da91a0a722d67165e849ce2cdc10de0e0d48738c142be8c6c5f310f4c \ - --hash=sha256:30ce1a5d16b5641dc935d50ef775af6a0871e3d14ab05d6fc54dff371b78e558 \ - --hash=sha256:30ce785d2adb8e32c3f7741442370a74834ec4c01f3c48f0750227a0b4ef27d6 \ - --hash=sha256:30f2218e6a9e5c18bc10fe6d41ac189c442c88eacf11bad9f28ef95a9bef00e6 \ - --hash=sha256:351a341c2105aa430b7047e30f1bf7975f6313b00165d3fc07be2edaf741f279 \ - --hash=sha256:37a10c377ce3a4a85f4a67f28b7afe093154cde77eaf248a72e856aa08b4d865 \ - --hash=sha256:392b8ab019e5502d08aff85c6272209c24bc2cbe706ea82a56368f524236614a \ - --hash=sha256:3e4540b8e74e4268811ac05db226a6a128ff572e7e0ce3f1163b693cadb184cd \ - --hash=sha256:40b2c7e92c44a84d748d21706c68dc6ff8161d80b59c99d774721a0d2317d7c7 \ - --hash=sha256:411fa4dfa5a7ae3d11491027ffb9beadec3996010a986862db70d91abba1c750 \ - --hash=sha256:4251acc80e2b7c9b7b8823456ea0fceeb0734dac2df7636d3c711b38476b5a76 \ - --hash=sha256:42bfb257930800cf43e7c62c832402c704ab60797c992faf88d20e903eac8f32 \ - --hash=sha256:4363818355dbc70ae1a8e9eaba9de350d93ede4ff6992b8f8eb8cbb6e5122d42 \ - --hash=sha256:4ab395feec8d249ec4044e228e98a7033f043426a265df439dc3698823f0a4e4 \ - --hash=sha256:50164d7610c00e7cd913a873fce30b6beeebf4b37e53983e33f22de4c900f6b8 \ - --hash=sha256:50e51156192722a9c58db112837d3f8ef96fb3c5ecc14e95f409134b08b158ec \ - --hash=sha256:510c8b3c17a0ed9ac69850c0438dada3c9b82d9c4d589fcb62002a5a9cf3a866 \ - --hash=sha256:5157de9f76eb4bc5ea74a1219366a25f945ad305641d74e04f59c54087091aa9 \ - --hash=sha256:54d5d6090cdc1b7c9e780dfb04949a990adb1e301a2fc0bbcee7de4638d33f9a \ - --hash=sha256:553fcac2ef2cb990877f9fc0833b8b629a3e6a5670b6b5fd58219b41a653ddc4 \ - --hash=sha256:5607e6013ed7e6b0ec9661e467b7ffde0aa7ab36833a04850f26fcf88ed4845b \ - --hash=sha256:5d6a60072b44c3c2b797a7ddcbcbbf2b34ea3cfd4721580fbfd2a09d9d9b84ba \ - --hash=sha256:5f30bae8bc1c2d613e28e5af3e8cceb09b742f1c8a8a5f839fb67afaffc03b61 \ - --hash=sha256:62ebd14e47e9aed9df4472afcb2663668ce4d74891cd54f86bf6e44029d6dc89 \ - --hash=sha256:631f13a3d04e97d4e083993b10f4b99530e3a10d953e2eb5e196b7dc7f812ce0 \ - --hash=sha256:6550fa135c7deb8ead6af49ed7ff648532ea8334a1447fe34a36315ef79c5c29 \ - --hash=sha256:66b1880df2d01e206e8339769d1c7c1753bcb653efd6289e203f6f24ebada0c0 \ - --hash=sha256:6eac374c5c975709b69c10f09afd199df74150172156ad10c8d4fd785b7da995 \ - --hash=sha256:71683c38c825452999b5717fcae07ea708e8c93003e808be4319c1b02e3d176e \ - --hash=sha256:7553333dd0930c104a5a0db8df72bf7219fe663d731383b576bb6ed6351c984d \ - --hash=sha256:75e8a04e91432dde9f1838373cf93d23726c79d3e908d319acf0e796f85592e7 \ - --hash=sha256:773b6eb282ce11ee19f05f6b2d4404fa308e5bbd353b0b80a0262caad6db2cd7 \ - --hash=sha256:774f93f65031856bf14ad9f59bdcab8b8cad501e5ceabd51ba3525f76937a25b \ - --hash=sha256:7c468136b8bd6bb18c8786e4236a1fa27362f24cb23450ba0cb204ab379b8e6f \ - --hash=sha256:7ce8902f939970048b233087082e7bb829db29375811c7ad50687b8624c6fd08 \ - --hash=sha256:7d3d6683288c11cbab50e865f2e2f13950179aa45410e30b2cfbd3fb7b0177bf \ - --hash=sha256:7f6163c0f10b055245f814dcc59f4818da60dfe72f3e72ab89fc24b6bd5e9c52 \ - --hash=sha256:8020c99ec13a7db2b6f96cbe82ef4721c88b426a4892f27478044af0284615ef \ - --hash=sha256:813dfbb17d65328bf86e5f0905dd277ba2265d3ca20556e86c0c7035b7182e5a \ - --hash=sha256:860a74063284a2ae9bfedd694f299cc2c68e2696c5f3d440cc9d18bb81b9dd04 \ - --hash=sha256:8c9004af7c8d67cce7f1aae1026fb55607f4aa600710d08ede3a3ce4aeefe7e0 \ - --hash=sha256:8d2c0c44d569ce0f2850f5c926f8caeb5f245fbc84475aeb36efccc2103e6dbd \ - --hash=sha256:8f7e9bc0f1135039b22ee6eab588d42df1ce55842b30740a352885eb267bd941 \ - --hash=sha256:90c5db5527c221249a876160663ab891ace358c17f7b9c93ec1478b7f0550e5c \ - --hash=sha256:9100ddbec09741cc66feb0fc6773f8bdbd0e3c345689368f260082ff85dcc0cd \ - --hash=sha256:913d02d29c9606643418d9ccfc3b72492ab25a6bf7889934e09a3490f8d3438b \ - --hash=sha256:980c256edb05b78a111b99c4de3b1d32e31634b867fd1fc2cf726e7b7bba9854 \ - --hash=sha256:9f924585cdacf631cd382b657966847bb537bf9ed0a6f9b991da5f05a631480f \ - --hash=sha256:a254e10b593624d230c365b6d616b22ca0ad65e63a16e6631c2b3466022e6ba8 \ - --hash=sha256:a2a438005b6f22d0273413484d6094d7c2c5d10ec1b3a3bf128e0d1d3ba53258 \ - --hash=sha256:a97261f1fccb8e50ecd2890a96e46efdc3f57c80a197324c6777827231eca712 \ - --hash=sha256:ab596fa3837e91e7e6a31b5f639988bfc6a35d1f915ac3932d946062219d588f \ - --hash=sha256:abbf258599526ad0326fe51e252e24f2bd6f24f1852681b4b78feda3808f1d18 \ - --hash=sha256:ac0d9ddea4350974be7a221fc25895f251a8fee748c889bdced2141c0fec1a49 \ - --hash=sha256:acf4ee4d1fc55917239fe72972fb292dd773055d05eb040d36f4326e02cc2c0e \ - --hash=sha256:ae1b0d82ac2d987f9ea512b1c9adfcc71a28de3dea3a6039b54d76cffda9901e \ - --hash=sha256:b15741f501469009ae0ae90b7147958a664a7dede40aa7ff174a8a4645f546d0 \ - --hash=sha256:b15d3ec9b0449c40e85319bdb4caa8b77ab526e74f5532ed94bec15e2f66822c \ - --hash=sha256:b3b3b775e33d3bfaec9899edc526ae97b0da0bf9d071a46124ba419149a414f8 \ - --hash=sha256:b6c0ffae686c39bf3737be60793783267628783ea42545632c10b291105aee45 \ - --hash=sha256:c210f8b35dc6f30aafd4b4365ca89b9d1189f21ab49b8e68fa6322a847aef138 \ - --hash=sha256:c2f6bb8b5216ab9e7873bc08b5d7bef2b8abbb578a3069bf1cd14a45d71d771d \ - --hash=sha256:c60e71b6d10cfc284c9bf36bd885e8d44c46f688ce50aa91b5edd90181dea687 \ - --hash=sha256:c6694a173ecabc12eb60efbc0b474464ead1951ff65cd8b1e72100715c64512b \ - --hash=sha256:c77496cb10bd7549690fbbab3e5ec05857b83e49276f4a9423a766ddd2afcd4c \ - --hash=sha256:c84c1b7be454b0c16f8499b4ebfbfd82ea5cca6527cceefcbbc06a7557b5ed2e \ - --hash=sha256:cc0bc345cf2df9d1c00ac443f50d543c1ccfa8b0422cb85b1ab70d681c0b255b \ - --hash=sha256:ceb8fc27d38793f9c97149be8302720c5b22e5c195a37bf2c45dc36c4600a512 \ - --hash=sha256:cf4bd113a69c0a740e27cb962ce10630c36d2b8f59d759a651b955ee9d18a823 \ - --hash=sha256:d1aa62e277fc1cbd80e6deacae6f4d983b41b3d7728e0645c5d741a6149bba45 \ - --hash=sha256:d1e7b1776f0797956c509e123d0952d10d293a9492dea9f288ab9570ec01d1a5 \ - --hash=sha256:d636d5095155afd364247f65070fab7beda13498d7ff4de331046e704ab9657f \ - --hash=sha256:d726e3ceeb337191324b49de298142f27c3ad10886341555d1d5315b5f252c6a \ - --hash=sha256:d72d8af5c1013656a8870c866660627d1a75bc185814ee022c8533caa1de88ae \ - --hash=sha256:d8d2955167274e15d79a7a020afdd9b39c990eb80b2d89fca695d92dcfdd38ec \ - --hash=sha256:d92a5cd21fdb083931d546c207aa29633787c5dc5b02daab2d32b843f88a2c53 \ - --hash=sha256:e58585a58209d72691ce2d62a9147445f5a87beb0bde97fde284c96ae392a3d1 \ - --hash=sha256:e7196e56f1cd69af1dbb07dff02dcfb260a50b45a82d409d92a06fedb32473b5 \ - --hash=sha256:eda3071db3346334beae1360b46da4606da57bf3528c167b3c38533afaf9f2c5 \ - --hash=sha256:edebcf7d1f601199084bb6e844d7dc67e03e04f6ac786b0332d616635c4ff7a4 \ - --hash=sha256:ef1fd24d9413f6209e00d3d5a453e67acfe004a25cc6c8e8484faed4311ab9e8 \ - --hash=sha256:f0b271b462769543716f92d3a4f90527df6ef5ed05ee95ec4137f513e21e1b77 \ - --hash=sha256:f18f85e4218d1b40f000f42a92239a7a61a902cd42c65e6c360dbd17dcb20894 \ - --hash=sha256:f1e1754960f38ec40613a07e5e372df67acb3b890fb383b6fb3de3e49ddbf3c7 \ - --hash=sha256:f2143ab06181d2b029eedcb6af3cebe95f11bbac62441781860f98ee9330a6a6 \ - --hash=sha256:f3d37768fce7f88dd2a8c6091f2325dea27d30d30d5c6e7a1c0f0af77723b708 \ - --hash=sha256:fa248c9eb220197d363f688818dac2fd4b2f0cd7d843ca7105d652034823427d -openai==2.33.0 \ - --hash=sha256:03ac37d70e8c9e3a8124214e3afa785e2cbc12e627fbd98177a086ef2fd87ad5 \ - --hash=sha256:f850c435e2a4685bba3295bd54912dd26315d9c1b7733068186134d6e0599f9a -pydantic==2.13.4 \ - --hash=sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba \ - --hash=sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6 -pydantic-core==2.46.4 \ - --hash=sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0 \ - --hash=sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262 \ - --hash=sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda \ - --hash=sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0 \ - --hash=sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e \ - --hash=sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b \ - --hash=sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594 \ - --hash=sha256:10e17cbb10a330363733efc4d7c4d0dd827ac0909b8f6a6542298fed1ea62f29 \ - --hash=sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2 \ - --hash=sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c \ - --hash=sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d \ - --hash=sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398 \ - --hash=sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d \ - --hash=sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3 \ - --hash=sha256:19e51f073cd3df251856a8a4189fbdf1de4012c3ebacfb1884f94f1eb406079f \ - --hash=sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb \ - --hash=sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7 \ - --hash=sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5 \ - --hash=sha256:228ee9bae8bef5b1e97ec58302f80357c37199e0d0a99174e138d28e6957b9d9 \ - --hash=sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462 \ - --hash=sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4 \ - --hash=sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b \ - --hash=sha256:2f84c03c8607173d16b5a854ec68a2f9079ae03237a54fb506d13af47e1d018d \ - --hash=sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df \ - --hash=sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2 \ - --hash=sha256:3447661d99f75a3683a4cf5c87da72f2161964611864dbbeac7fbb118bb4bfc0 \ - --hash=sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519 \ - --hash=sha256:395aebd9183f9d112f569aeb5b2214d1a10a33bec8456447f7fbdfa51d38d4cd \ - --hash=sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7 \ - --hash=sha256:3be77f45df024d789a672ae34f8b06fb346c4f9f46ea714956660ea4862e89ac \ - --hash=sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6 \ - --hash=sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565 \ - --hash=sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898 \ - --hash=sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb \ - --hash=sha256:432c179df7874eeb73307aad2df0755e1ae0efa61ff0ea89b93e194411ae3928 \ - --hash=sha256:4a05d69cba51d852c5c3e92758653245a50c0b646ced0cf05bd793ed592839d6 \ - --hash=sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3 \ - --hash=sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a \ - --hash=sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596 \ - --hash=sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987 \ - --hash=sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e \ - --hash=sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d \ - --hash=sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712 \ - --hash=sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008 \ - --hash=sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd \ - --hash=sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1 \ - --hash=sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be \ - --hash=sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea \ - --hash=sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292 \ - --hash=sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33 \ - --hash=sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3 \ - --hash=sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4 \ - --hash=sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b \ - --hash=sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826 \ - --hash=sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac \ - --hash=sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7 \ - --hash=sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d \ - --hash=sha256:8358a950c8909158e3df31538a7e4edc2d7265a7c54b47f0864d9e5bae9dcebf \ - --hash=sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4 \ - --hash=sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc \ - --hash=sha256:8b9bab013d1c7a79d3501ff86d0bc9c31bf587db4551677b96bec07df78c6b15 \ - --hash=sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3 \ - --hash=sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b \ - --hash=sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914 \ - --hash=sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04 \ - --hash=sha256:905a0ed8ea6f2d61c1738835f99b699348d7857379083e5fc497fa0c967a407c \ - --hash=sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b \ - --hash=sha256:91a06d2e259ecfbd8c901d70c3c507900458498142b3026a296b7de4d1322cc9 \ - --hash=sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce \ - --hash=sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4 \ - --hash=sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a \ - --hash=sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f \ - --hash=sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424 \ - --hash=sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894 \ - --hash=sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9 \ - --hash=sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76 \ - --hash=sha256:9f444c499b3eefd3a92e348059471ea0c3a6e303d9c1cec09fa748fd9f895201 \ - --hash=sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb \ - --hash=sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109 \ - --hash=sha256:a396dcc17e5a0b164dbe026896245a4fa9ff402edca1dff0be3d53a517f74de4 \ - --hash=sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848 \ - --hash=sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526 \ - --hash=sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0 \ - --hash=sha256:b078afbc25f3a1436c7a1d2cd3e322497ee99615ba97c563566fdf46aff1ee01 \ - --hash=sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458 \ - --hash=sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e \ - --hash=sha256:bb63e0198ca18aad131c089b9204c23079c3afa95487e561f4c522d519e55aba \ - --hash=sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a \ - --hash=sha256:c1747f85cee84c26985853c6f3d9bd3e75da5212912443fa111c113b9c246f39 \ - --hash=sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c \ - --hash=sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000 \ - --hash=sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b \ - --hash=sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf \ - --hash=sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4 \ - --hash=sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd \ - --hash=sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28 \ - --hash=sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9 \ - --hash=sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30 \ - --hash=sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983 \ - --hash=sha256:d80ee3d731373b24cebbc10d689ca4ee1875caf0d5703a245db18efd4dd37fc1 \ - --hash=sha256:d995260fdf4e1db774581b4900e0f832abe3c7c84996726bbc161b19c8f29e76 \ - --hash=sha256:da4b951fe36dc7c3a1ccb4e3cd1747c3542b8c9ceede8fc86cae054e764485f5 \ - --hash=sha256:daa27d92c36f24388fe3ad306b174781c747627f134452e4f128ea00ce1fe8c4 \ - --hash=sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7 \ - --hash=sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c \ - --hash=sha256:e68b7a074f65a2fd746c52a7ce6142ab7006074ac269ace0c25cd8ba171f8066 \ - --hash=sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3 \ - --hash=sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02 \ - --hash=sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89 \ - --hash=sha256:ea793e075b70290d89d8142074262885d3f7da19634845135751bd6344f73b50 \ - --hash=sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76 \ - --hash=sha256:f13a646d65d09fbf1bc6b3a9635d30095c8e7e5cc419ff35ecc563c5fd04cd49 \ - --hash=sha256:f47286a97f0bc9b8859519809077b91b2cefe4ae47fcbf5e466a009c1c5d742b \ - --hash=sha256:f747929cf940cddb5b3668a390056ddd5ba2e5010615ea2dcf4f9c4f3ab8791d \ - --hash=sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7 \ - --hash=sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4 \ - --hash=sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c \ - --hash=sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e \ - --hash=sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff \ - --hash=sha256:fd8b3d9fd264be37976686c7f65cd52a83f5e84f4bfd2adf9c1d469676bbb6ae -sniffio==1.3.1 \ - --hash=sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2 \ - --hash=sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc -tqdm==4.68.3 \ - --hash=sha256:00dfa48452b6b6cfae3dd9885636c23d3422d1ec97c66d96818cbd5e0821d482 \ - --hash=sha256:39832cc2def2789a6f29df83f172db7416cea70052c0907a57801c5f2fdccb03 -typing-extensions==4.15.0 \ - --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \ - --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548 -typing-inspection==0.4.2 \ - --hash=sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7 \ - --hash=sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464 diff --git a/.github/scripts/triage_with_llm.py b/.github/scripts/triage_with_llm.py deleted file mode 100644 index e23a012425a..00000000000 --- a/.github/scripts/triage_with_llm.py +++ /dev/null @@ -1,1797 +0,0 @@ -#!/usr/bin/env python3 -""" -Agent Shin — LLM-as-judge triage for external OSS pull requests and issues. - -Evaluates a single PR or issue against the contribution rubric and, when the -LLM judge marks it as failing, posts an explanatory comment + closes the -PR/issue. Re-triggers on `reopened` so contributors can iterate back in by -filling in the missing pieces and reopening. - -Internal BerriAI contributors (`author_association` in {OWNER, MEMBER, -COLLABORATOR}) and bot accounts are skipped entirely. - -Usage: - triage_with_llm.py --repo owner/repo --pr 1234 - triage_with_llm.py --repo owner/repo --issue 5678 - triage_with_llm.py --repo owner/repo --pr 1234 --close # actually close - triage_with_llm.py --repo owner/repo --pr 1234 --print-prompt # show prompt - -Defaults are SAFE: without `--close` the script writes a verdict to stdout (and, -when running in GitHub Actions, to $GITHUB_STEP_SUMMARY) but takes no GitHub -write actions. - -Environment: - GH_TOKEN / GITHUB_TOKEN - for `gh` CLI auth (auto-set in Actions) - OPENAI_API_KEY - required when --close is passed - OPENAI_BASE_URL - optional (route to any OpenAI-compatible API) - TRIAGE_MODEL - optional model override (default: gpt-5.4-mini) -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import json -import os -import re -import subprocess -import sys -import textwrap -import urllib.parse -from typing import Any, Iterable - -# Add this script's directory to `sys.path` so the sibling -# `agent_shin_shared` module is importable when the script is invoked -# directly (e.g. `python3 .github/scripts/triage_with_llm.py ...`) and -# also when the tests load this script via -# `importlib.util.spec_from_file_location`. -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from agent_shin_shared import ( # noqa: E402 -- sys.path adjusted above - AGENT_SHIN_CLOSE_MARKER, - AGENT_SHIN_DEFAULT_BOT_LOGIN, - ALLOWLIST_LOGINS, - GRACE_COMMENT_MARKER, - GRACE_PERIOD_SECONDS, - GREPTILE_BOT_LOGINS, - SCORE_PATTERN, - extract_greptile_score, - gh, - parse_iso8601, - seconds_since_latest_marker_comment, -) - -DEFAULT_MODEL = "gpt-5.4-mini" - -INTERNAL_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"}) - -# `AGENT_SHIN_DEFAULT_BOT_LOGIN` is imported from `agent_shin_shared`. -# When the workflow uses the default `secrets.GITHUB_TOKEN`, the -# closure / reopen event's `actor.login` is `github-actions[bot]`. The -# env override `AGENT_SHIN_BOT_LOGIN` exists for local debugging and for -# repos that wire Agent Shin to a PAT. - -# HTML marker appended to every reconsider verdict comment. We grep for this -# on subsequent reconsider triggers to enforce a short cooldown so that -# repeated `@agent-shin reconsider` comments don't burn CI/LLM budget. -# Using a unique HTML comment keeps the marker invisible to humans while -# being trivially greppable from a comments-list API response. -RECONSIDER_COMMENT_MARKER = "" - -# Minimum gap between two reconsider verdicts on the same PR/issue. Set to -# 10 minutes — long enough that a contributor can't trivially spam the -# trigger, short enough that a genuine "I just pushed a fix and reupdated -# the body" iteration loop isn't punished. -RECONSIDER_RATE_LIMIT_SECONDS = 600 - -# `GRACE_COMMENT_MARKER` (HTML marker on the grace-period warning comment -# posted on the first low-quality detection — used on subsequent triage -# runs to detect that a warning was already posted and measure how long -# ago it was posted) and `GRACE_PERIOD_SECONDS` (length of the grace -# period between the warning and the actual auto-close, 2 hours) are -# imported from `agent_shin_shared` so the daily Greptile sweep and the -# LLM judge agree on the same marker and duration. - -# --- Review-gate ("ready for review" label lifecycle) configuration ---------- -# The review gate keeps a single label in sync with whether a PR currently -# clears BOTH quality bars: the LLM rubric (clear problem + expected/actual + -# QA proof, or a linked issue) AND Greptile's most recent confidence score. -READY_FOR_REVIEW_LABEL = "ready for review" -DEFAULT_GRACE_DAYS = 1 # 24h before an un-passing, un-tagged PR is auto-closed -DEFAULT_MIN_GREPTILE_SCORE = 4 # Greptile < 4/5 counts as "not passing" - -# Hidden HTML-comment markers stamped into review-gate comments. They never -# render in the GitHub UI but let the gate detect its own prior actions so it -# (a) posts the within-grace "what's missing" notice at most once and (b) can -# tell a first-time pass ("ready for review") from a recovery after a -# regression ("all clear again"). -READY_MARKER = "" -REGRESSED_MARKER = "" -WITHIN_GRACE_MARKER = "" - -# `GREPTILE_BOT_LOGINS` (Greptile's GitHub App login variants — -# `greptile-apps[bot]` in REST API comments, `greptile-apps` in -# `gh pr view --json` output) and `SCORE_PATTERN` (regex matching lines -# like `Confidence Score: 3/5`) are imported from `agent_shin_shared` -# so the daily sweep and the review gate read the score through the -# same set of logins / patterns. - -# `AGENT_SHIN_CLOSE_MARKER` is imported from `agent_shin_shared` so this LLM -# judge and the daily Greptile sweep stamp the same marker on their close -# comments — `was_closed_by_agent_shin` keys the reconsider reopen path off it. - -# Model families that require `reasoning_effort` to be set, and that reject -# `temperature != 1` unless `reasoning_effort` is "none". For these models we -# pass `reasoning_effort="none"` so a `temperature=0` deterministic judgment -# is still accepted. See litellm/llms/openai/chat/gpt_5_transformation.py for -# the full set of constraints LiteLLM applies to these models. -GPT5_FAMILY_PREFIX = "gpt-5" - -# Regexes for picking off "obvious passes" without burning LLM tokens. -# -# Keep this list to GitHub's documented PR-closing keywords only -# (https://docs.github.com/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue). -# Casual mentions like "see #1234" or "ref #1234" are intentionally NOT -# auto-passed — they should fall through to the LLM judge, which has the -# stricter rubric "a bare issue number without a closing keyword counts only -# if it's clearly the related issue (not a passing mention)". -LINKED_ISSUE_PATTERN = re.compile( - r"\b(?:fixes|fix|fixed|closes|close|closed|resolves|resolve|resolved)\s+" - r"(?:#\d+|https?://github\.com/[\w.-]+/[\w.-]+/issues/\d+)", - re.IGNORECASE, -) -HTML_COMMENT_PATTERN = re.compile(r"", re.DOTALL) - - -# --------------------------------------------------------------------------- -# gh helpers -# -# `gh` is imported from `agent_shin_shared` so a future change (timeout, -# logging, retry) only needs to be made once. - - -def fetch_pr(repo: str, number: int) -> dict: - """Return the full GitHub REST representation of a PR.""" - return json.loads(gh("api", f"repos/{repo}/pulls/{number}")) - - -def fetch_issue(repo: str, number: int) -> dict: - """Return the full GitHub REST representation of an issue.""" - return json.loads(gh("api", f"repos/{repo}/issues/{number}")) - - -def post_comment(repo: str, number: int, body: str) -> None: - """Post an issue-style comment (works for both issues and PRs).""" - gh( - "api", - f"repos/{repo}/issues/{number}/comments", - "-X", - "POST", - "-f", - f"body={body}", - ) - - -def close_pr(repo: str, number: int) -> None: - """Close a pull request (state=closed).""" - gh( - "api", - f"repos/{repo}/pulls/{number}", - "-X", - "PATCH", - "-f", - "state=closed", - ) - - -def reopen_pr(repo: str, number: int) -> None: - """Reopen a previously-closed pull request (state=open). - - Used by the `@agent-shin reconsider` comment-trigger flow: the bot has - write access via GH_TOKEN, so it can reopen on the contributor's behalf - even though GitHub doesn't let the OSS author do it themselves. - """ - gh( - "api", - f"repos/{repo}/pulls/{number}", - "-X", - "PATCH", - "-f", - "state=open", - ) - - -def close_issue(repo: str, number: int, *, not_planned: bool = True) -> None: - """Close an issue, marking state_reason=not_planned by default.""" - args = [ - "api", - f"repos/{repo}/issues/{number}", - "-X", - "PATCH", - "-f", - "state=closed", - ] - if not_planned: - args.extend(["-f", "state_reason=not_planned"]) - gh(*args) - - -def reopen_issue(repo: str, number: int) -> None: - """Reopen a previously-closed issue (state=open, state_reason=reopened).""" - gh( - "api", - f"repos/{repo}/issues/{number}", - "-X", - "PATCH", - "-f", - "state=open", - "-f", - "state_reason=reopened", - ) - - -def add_label(repo: str, number: int, label: str) -> None: - """Add a label to a PR/issue (GitHub creates the label if it's missing).""" - gh( - "api", - f"repos/{repo}/issues/{number}/labels", - "-X", - "POST", - "-f", - f"labels[]={label}", - ) - - -def remove_label(repo: str, number: int, label: str) -> None: - """Remove a label from a PR/issue. A missing label (404) is not an error.""" - encoded = urllib.parse.quote(label, safe="") - try: - gh( - "api", - f"repos/{repo}/issues/{number}/labels/{encoded}", - "-X", - "DELETE", - ) - except subprocess.CalledProcessError as exc: - stderr = (exc.stderr or "").lower() - if "404" in stderr or "not found" in stderr: - return - raise - - -def _iter_paginated_json(*api_args: str) -> Any: - """Yield JSON objects from `gh api --paginate ... -q '.[]'`. - - `gh api --paginate` on a JSON-array endpoint concatenates pages into - one stream; `-q '.[]'` flattens that stream into newline-delimited - objects (jq-style). This keeps memory bounded for chatty endpoints - like issue events/comments on long-lived PRs. - """ - raw = gh("api", "--paginate", *api_args, "-q", ".[]") - for line in raw.splitlines(): - line = line.strip() - if not line: - continue - try: - yield json.loads(line) - except json.JSONDecodeError: - # A malformed line should not blow up the whole guard. Skip and - # carry on — at worst the guard fail-closes (returns False / - # None) and the caller treats it as "unknown". - continue - - -def fetch_last_close_event( - repo: str, number: int -) -> tuple[str | None, dt.datetime | None]: - """Return the actor login and timestamp of the most recent `closed` event. - - Either field may be None: actor when the events API returns nothing - (unusual for a closed item, but possible on transient errors), and - timestamp when the event lacks `created_at` or the value can't be - parsed. `was_closed_by_agent_shin` fail-closes on either. - """ - actor: str | None = None - closed_at: dt.datetime | None = None - for event in _iter_paginated_json(f"repos/{repo}/issues/{number}/events"): - if event.get("event") != "closed": - continue - actor = (event.get("actor") or {}).get("login") - created = event.get("created_at") - if not created: - closed_at = None - continue - try: - closed_at = parse_iso8601(created) - except ValueError: - closed_at = None - return actor, closed_at - - -# How much older than the latest `closed` event the Agent Shin marker -# comment is allowed to be while still counting as "this close was Agent -# Shin's". Agent Shin posts the close comment immediately before closing, -# so the marker timestamp is normally at most a few seconds before the -# close event; the buffer just absorbs clock skew between the comments -# API and the events API. -AGENT_SHIN_CLOSE_MARKER_SKEW_SECONDS = 300 - - -def was_closed_by_agent_shin( - repo: str, number: int, *, bot_login: str | None = None -) -> bool: - """Return True iff Agent Shin itself most-recently closed this PR/issue. - - This is the guard that stops `@agent-shin reconsider` from reopening an - item Agent Shin did not close — a maintainer closing for non-rubric - reasons (security, duplicate, design rejection), or a different workflow - (stale/duplicate sweeps) closing under the shared `github-actions[bot]` - identity. Three independent signals must all hold, because that identity - is not unique to Agent Shin and a marker comment from a prior - closed/reopened cycle would otherwise vouch for an unrelated close: - - 1. The most recent `closed` event's actor is the bot identity. - 2. Agent Shin left one of its auto-close comments, detected via - `AGENT_SHIN_CLOSE_MARKER`. The actor check alone can't tell an - Agent Shin close from any other `github-actions[bot]` close. - 3. That marker comment was posted at (or just before) the latest - close event, not on a previous close in an - Agent-Shin-close -> reconsider-reopen -> other-bot-reclose cycle. - - The check is intentionally fail-closed: any uncertainty about who closed - the item is treated as "not Agent Shin" so the destructive reopen path - stays gated. - """ - expected = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - actor, closed_at = fetch_last_close_event(repo, number) - if not actor or actor.lower() != expected or closed_at is None: - return False - marker_seconds = seconds_since_last_agent_shin_close( - repo, number, bot_login=bot_login - ) - if marker_seconds is None: - return False - close_age_seconds = (dt.datetime.now(dt.timezone.utc) - closed_at).total_seconds() - return marker_seconds <= close_age_seconds + AGENT_SHIN_CLOSE_MARKER_SKEW_SECONDS - - -def _seconds_since_latest_marker_comment( - repo: str, - number: int, - *, - marker: str, - bot_login: str | None = None, -) -> float | None: - """Return seconds since the bot's most recent comment with ``marker``. - - Fetches comments via `_iter_paginated_json` and delegates the - iteration / author-filter / timestamp logic to - `agent_shin_shared.seconds_since_latest_marker_comment` so the daily - Greptile sweep and the LLM judge use one source of truth for the - "bot already posted X" detection. The wall-clock `now` is resolved - against this module's `dt` so tests that freeze time via - `monkeypatch.setattr(triage_module, "dt", ...)` still apply. - """ - return seconds_since_latest_marker_comment( - _iter_paginated_json(f"repos/{repo}/issues/{number}/comments"), - marker=marker, - bot_login=bot_login, - now=dt.datetime.now(dt.timezone.utc), - ) - - -def seconds_since_last_reconsider_verdict( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since the bot's most recent reconsider verdict comment. - - Detects comments by matching the HTML marker `RECONSIDER_COMMENT_MARKER` - appended by `format_reopen_comment` and - `format_reconsider_still_failing_comment`. Returns None when the bot - has never posted a reconsider verdict on this PR/issue (or when the - only matching comments are missing a `created_at` timestamp, which - shouldn't happen on a real GitHub response). - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=RECONSIDER_COMMENT_MARKER, bot_login=bot_login - ) - - -def seconds_since_last_grace_warning( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since the bot's most recent grace-period warning. - - Detects warning comments by matching the HTML marker - `GRACE_COMMENT_MARKER` appended by `format_grace_warning_pr_comment` - and `format_grace_warning_issue_comment`. Returns None when no - grace warning has ever been posted on this PR/issue — that's the - "first low-quality detection" signal that drives the warning path. - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=GRACE_COMMENT_MARKER, bot_login=bot_login - ) - - -def seconds_since_last_agent_shin_close( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since Agent Shin's most recent auto-close comment. - - Detects close comments by matching `AGENT_SHIN_CLOSE_MARKER` (stamped by - `format_pr_close_comment` / `format_issue_close_comment`). Returns None - when Agent Shin has never closed this PR/issue — the signal - `was_closed_by_agent_shin` uses to keep the reconsider reopen path gated - against closures performed by other workflows sharing the bot identity. - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=AGENT_SHIN_CLOSE_MARKER, bot_login=bot_login - ) - - -# --------------------------------------------------------------------------- -# Author classification - - -def is_internal_contributor(item: dict) -> bool: - """Return True if the PR/issue author should be exempted from triage. - - Fail-safe: if `author_association` is missing or empty (which should never - happen on a successful GitHub REST response but is possible on schema - changes or partial responses), treat the author as INTERNAL so the - destructive close path never fires on an unknown contributor. This matches - the sibling `is_external_pr_author` in `close_low_quality_prs.py`. - """ - login = ((item.get("user") or {}).get("login") or "").lower() - if login.endswith("[bot]") or login in {"dependabot", "github-actions"}: - return True - association = (item.get("author_association") or "").upper() - if not association or association in INTERNAL_ASSOCIATIONS: - return True - return False - - -# --------------------------------------------------------------------------- -# Greptile score + age helpers (`extract_greptile_score`, `parse_iso8601`) -# live in `agent_shin_shared` — they're imported at the top of this module -# so both `triage_with_llm.py` and `close_low_quality_prs.py` share a -# single source of truth for the Confidence-Score regex and ISO-8601 -# parsing. - - -# --------------------------------------------------------------------------- -# Prompt construction - - -def strip_html_comments(text: str) -> str: - """Remove HTML comments — template placeholder text shouldn't fool the judge.""" - return HTML_COMMENT_PATTERN.sub("", text or "") - - -def has_linked_issue(text: str) -> bool: - """Heuristic: does this body link to an open issue (Fixes #123 etc.)?""" - return bool(LINKED_ISSUE_PATTERN.search(strip_html_comments(text or ""))) - - -def build_pr_prompt(*, title: str, body: str) -> str: - cleaned_body = strip_html_comments(body or "").strip() or "(empty)" - # Dedent the static template *before* interpolating dynamic fields so that - # multi-line bodies (whose 2nd+ lines start at column 0) don't defeat the - # common-indent computation in textwrap.dedent. - template = textwrap.dedent(""" - You are "Agent Shin", the OSS triage bot for the LiteLLM open-source - repository (BerriAI/litellm). Decide whether this external pull request - meets the project's contribution standards. - - A PR PASSES triage only if BOTH (1) AND (2) are satisfied. A linked - issue alone is NOT enough — it covers context, not proof. - - (1) CONTEXT — the PR provides AT LEAST ONE of: - (a) A link to a related GitHub issue. Acceptable forms: - "Fixes #1234", "Closes #1234", "Resolves #1234", - "Refs https://github.com/BerriAI/litellm/issues/1234". A - bare "#1234" without a closing keyword counts only if it - is clearly the related issue (not a passing mention). - (b) A clear problem description in the body (what bug or - missing feature this addresses, beyond the title) AND - expected vs. actual behavior (or, for features, "what's - possible now vs. with this PR"). - - (2) END-TO-END QA PROOF: the PR body contains AT LEAST ONE of: - (a) A screen recording / video showing the behavior before - and after the change (the bug reproducing, then the fix - working). For a brand-new feature with no meaningful - "before", a recording of it working end-to-end is fine. - (b) A screenshot (or before/after screenshots) showing the - fix or feature working. - (c) Specific commands that were actually run (curl, python, - a CLI invocation, etc.) PAIRED WITH their real - output, demonstrating the change works end-to-end against - the real system. Commands whose external dependencies - (LLM provider, DB, network) are mocked or stubbed do NOT - satisfy (2c); they are not end-to-end. - - `has_qa_proof` must be set to `true` only when (2a), (2b), - or a non-mocked (2c) is actually present in the body. If the - only "proof" is mocked tests, `has_qa_proof` is `false` and - the verdict is "fail". - - The following do NOT count as QA proof: - - Generic claims like "I tested it", "works locally", "all - tests pass", or a checked "I added tests" checkbox with no - output shown. - - A description of what tests exist or were added, without - their actual output in the PR body. - - `pytest` (or any test runner) executed against the - repository's own unit tests. Those mock the LLM provider, - DB, and network, so they are NOT end-to-end and never - satisfy (2), no matter how much passing output is pasted. - - A linked issue. The linked issue is context (1a), never - proof (2). - - FAIL the PR if EITHER (1) or (2) is missing. Do not bias toward PASS: - if QA proof is absent, the verdict is "fail" even when the rest of - the PR is well-written. - - Respond with a single JSON object, no prose: - - {{ - "verdict": "pass" | "fail", - "linked_issue": boolean, - "has_problem_description": boolean, - "has_expected_vs_actual": boolean, - "has_qa_proof": boolean, - "qa_proof_type": "video" | "screenshot" | "commands_with_output" | "none", - "missing": ["plain-english strings naming what is missing"], - "explanation": "1-2 sentence reasoning for the team to skim" - }} - - --- - PR title: {title} - - PR body: - --- - {cleaned_body} - --- - """).strip() - return template.format(title=title, cleaned_body=cleaned_body) - - -def build_issue_prompt(*, title: str, body: str) -> str: - cleaned_body = strip_html_comments(body or "").strip() or "(empty)" - # Dedent the static template *before* interpolating dynamic fields so that - # multi-line bodies (whose 2nd+ lines start at column 0) don't defeat the - # common-indent computation in textwrap.dedent. - template = textwrap.dedent(""" - You are "Agent Shin", the OSS triage bot for the LiteLLM open-source - repository (BerriAI/litellm). Decide whether this GitHub issue meets - the project's reporting standards. - - For a BUG REPORT the issue PASSES triage only when it contains BOTH: - (1) END-TO-END EVIDENCE OF THE BUG (the "before"; set - `has_repro=true` only when this is present): AT LEAST ONE of: - (a) A screen recording / video of the bug happening. - (b) A screenshot of the bug. - (c) The exact command(s) actually run (curl, python, a CLI - invocation, etc.) PAIRED WITH their real output, traceback, - or logs showing the failure against the real system. - Commands whose external dependencies (LLM provider, DB, - network) are mocked or stubbed do NOT count. - Prose-only "steps to reproduce" with no run output, video, or - screenshot do NOT satisfy (1). An unfilled template scaffold - (bare headings such as "Version or commit:" with nothing under - them, empty numbered lists) counts as absent, not as evidence. - (2) Expected vs. actual behavior (`has_expected_vs_actual`). - - FAIL the bug report if either (1) or (2) is missing. Do not bias - toward PASS: if the bug isn't demonstrated end-to-end, the verdict is - "fail" even when the report is well-written. - - For a FEATURE REQUEST the issue PASSES triage only when it contains - ALL of: - - A clear description of the proposed feature (what should LiteLLM do - that it does not today). - - Motivation / use case with a concrete example (config, API call, - UI flow, or scenario showing what's blocked today). - - END-TO-END EVIDENCE OF THE DEAD-END (set - `has_dead_end_evidence=true` only when this is present): a video, - a screenshot, or the exact command(s) actually run paired with - their real output, showing the point where the flow stops today. - Mocked or stubbed dependencies do NOT count, and an unfilled - template scaffold (bare headings, empty numbered lists) counts as - absent. - - For an issue that is neither a bug report nor a feature request (a - question, support request, or discussion), PASS as long as it has a - clear, specific ask and is not empty or template placeholder text. - - Respond with a single JSON object, no prose: - - {{ - "verdict": "pass" | "fail", - "kind": "bug" | "feature" | "other", - "has_repro": boolean, - "has_expected_vs_actual": boolean, - "has_motivation_example": boolean, - "has_dead_end_evidence": boolean, - "missing": ["plain-english strings naming what is missing"], - "explanation": "1-2 sentence reasoning for the team to skim" - }} - - --- - Issue title: {title} - - Issue body: - --- - {cleaned_body} - --- - """).strip() - return template.format(title=title, cleaned_body=cleaned_body) - - -# --------------------------------------------------------------------------- -# LLM call + verdict parsing - - -def call_llm_judge( - prompt: str, *, model: str, api_key: str, base_url: str | None -) -> str: - """Call an OpenAI-compatible chat completions endpoint. Returns raw text.""" - # Import inside the function so unit tests that monkey-patch this never - # need the openai package installed. - from openai import OpenAI - - client = ( - OpenAI(api_key=api_key, base_url=base_url) - if base_url - else OpenAI(api_key=api_key) - ) - kwargs: dict[str, Any] = { - "model": model, - "messages": [{"role": "user", "content": prompt}], - "temperature": 0, - "response_format": {"type": "json_object"}, - } - # gpt-5.x reasoning models reject `temperature != 1` unless - # `reasoning_effort` is explicitly "none". Set it via `extra_body` so this - # works across openai SDK versions regardless of whether the SDK natively - # types `reasoning_effort` as a top-level chat-completions param yet. - if model.lower().startswith(GPT5_FAMILY_PREFIX): - kwargs["extra_body"] = {"reasoning_effort": "none"} - response = client.chat.completions.create(**kwargs) - return response.choices[0].message.content or "" - - -def parse_verdict(raw: str) -> dict: - """Parse the LLM's JSON response. Tolerates ```json fences and stray text.""" - if not raw: - raise ValueError("empty LLM response") - text = raw.strip() - if text.startswith("```"): - text = re.sub(r"^```(?:json)?\s*", "", text) - text = re.sub(r"\s*```$", "", text) - try: - return json.loads(text) - except json.JSONDecodeError: - match = re.search(r"\{.*\}", text, re.DOTALL) - if not match: - raise ValueError(f"could not extract JSON from LLM response: {raw[:200]}") - return json.loads(match.group(0)) - - -# --------------------------------------------------------------------------- -# Comment composition - - -def _format_missing(missing: list[str]) -> str: - if not missing: - return "- (see explanation below)" - return "\n".join(f"- {m}" for m in missing) - - -# Rubric items the judge can mark present. The first element of each tuple is -# the verdict-JSON boolean field, the second is the human-readable label we -# render in the "what you got right" section of close / grace-warning comments. -_PR_PRESENT_LABELS: tuple[tuple[str, str], ...] = ( - ("linked_issue", "Linked a related GitHub issue"), - ("has_problem_description", "Clear problem description"), - ("has_expected_vs_actual", "Expected vs. actual behavior"), - ("has_qa_proof", "End-to-end QA proof"), -) - -# Issue rubric labels grouped by `kind`. The judge sets `kind` to one of -# {"bug", "feature", "other"}; when "other" we render both groups so we don't -# silently drop a present-flag the judge actually set to True. -_ISSUE_BUG_LABELS: tuple[tuple[str, str], ...] = ( - ( - "has_repro", - "End-to-end evidence of the bug (video, screenshot, or command + real output)", - ), - ("has_expected_vs_actual", "Expected vs. actual behavior"), -) -_ISSUE_FEATURE_LABELS: tuple[tuple[str, str], ...] = ( - ("has_motivation_example", "Motivation and concrete example"), - ( - "has_dead_end_evidence", - "End-to-end evidence of the dead-end (video, screenshot, or command + real output)", - ), -) - - -def _format_present_for_pr(verdict: dict) -> list[str]: - """Human-readable rubric items the judge confirmed are present on a PR. - - Drives the "what you got right" section in close / grace-warning comments. - The user gave explicit feedback: contributors should see what they nailed - *before* the list of gaps, so the comment doesn't read as pure rejection. - """ - return [label for field, label in _PR_PRESENT_LABELS if verdict.get(field)] - - -def _format_present_for_issue(verdict: dict) -> list[str]: - """Human-readable rubric items the judge confirmed are present on an issue. - - Branches on the judge's `kind` field. For `"other"` (or missing kind) we - render the union so a present-flag isn't dropped just because the judge - couldn't classify the issue cleanly. - """ - kind = (verdict.get("kind") or "").lower() - groups: list[tuple[tuple[str, str], ...]] = [] - if kind in ("bug", "other", ""): - groups.append(_ISSUE_BUG_LABELS) - if kind in ("feature", "other", ""): - groups.append(_ISSUE_FEATURE_LABELS) - out: list[str] = [] - for group in groups: - for field, label in group: - if verdict.get(field) and label not in out: - out.append(label) - return out - - -def _format_present_block(items: list[str]) -> str: - """Render the optional "what you got right" block. Empty string when the - judge didn't confirm anything as present — better to omit the section - entirely than to show "What you got right: (nothing)". - """ - if not items: - return "" - bullets = "\n".join(f"- ✅ {item}" for item in items) - return f"**What you got right:**\n\n{bullets}\n\n" - - -def format_pr_close_comment(verdict: dict) -> str: - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_pr(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the description against our " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md). " - "Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "**Closing this PR isn't a rejection of the change.** We want the open-PR list to " - "mirror what a maintainer can act on *right now*, so contributors don't get lost in a " - 'backlog. A closed PR is a soft "park this for later"; your work is still here, ' - "the diff is still here, and getting it reopened is one comment away. Take your time.\n" - "\n" - "**To bring this PR back:**\n" - "\n" - "- Update the description with the missing pieces, then comment `@agent-shin reconsider` " - "on this PR. I'll re-evaluate and reopen if it now passes.\n" - "- Or **Open a new PR** with the same fix and the updated description. GitHub doesn't " - "always let external contributors reopen a bot-closed PR, so a fresh PR is the most " - "reliable path back into the review queue.\n" - "- If Greptile's most recent score on this PR was below 4/5, comment `@greptileai` to " - "request a fresh review; that **still works even after the PR is closed**, and a " - "stronger score is one of the signals that lifts the PR back into the queue. A low " - "Greptile score isn't a blocker.\n" - "\n" - '**What "end-to-end QA proof" means**, since it\'s the most common gap: at least one ' - "of a short before/after screen recording / video (the bug reproducing, then the fix " - "working; for a brand-new feature, a recording of it working end-to-end), a screenshot " - "(or before/after screenshots) of it working, or the exact commands you ran paired " - "with their **real output** against the real system. Running `pytest` on the repo's " - "unit tests doesn't count; those mock the LLM provider, DB, and network, so they " - "aren't end-to-end. Output from a real, no-mocks integration run is what we look " - "for. A linked issue alone isn't enough either: it covers context, not proof. See " - "[the full rubric](https://docs.litellm.ai/blog/agent-shin-triage#the-rubric-for-pull-requests).\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, comment " - "`@agent-shin reconsider` or ping a maintainer; they'll override me.)_" - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def format_issue_close_comment(verdict: dict) -> str: - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_issue(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for filing this! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the issue against our reporting checklist. Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "**Closing this isn't us saying the bug isn't real or the request isn't useful.** We " - "want the open-issue list to mirror what a maintainer can act on *right now*, so " - "reports like yours don't get buried in a backlog. A closed issue is a soft \"park " - 'this for later"; your report is still here, and getting it reopened is one comment ' - "away. Take your time.\n" - "\n" - "**To bring this issue back:**\n" - "\n" - "1. Edit the issue description to add the missing pieces:\n" - " - For **bug reports**: end-to-end evidence of the bug (a screen recording / " - "video, a screenshot, or the exact commands you ran with their real output / " - "traceback) plus expected vs. actual behavior. Written steps with no run output, " - "video, or screenshot don't count, and mocked or stubbed runs don't count.\n" - " - For **feature requests**: a concrete description of what should change, a " - "use case and example (config / API call / UI flow), plus end-to-end evidence of " - "the dead-end (a video, a screenshot, or the exact commands you ran with their " - "real output showing where the flow stops today). Mocked or stubbed runs don't " - "count.\n" - "2. Comment `@agent-shin reconsider`. I'll re-run triage and reopen the issue if it " - "now meets the bar. (GitHub doesn't let external authors reopen an issue a maintainer " - "or bot closed, so the comment-based reconsider is the reliable path.)\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, comment " - "`@agent-shin reconsider` or ping a maintainer; they'll override me.)_" - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def format_grace_warning_pr_comment(verdict: dict) -> str: - """Comment posted on the FIRST low-quality detection — gives the - contributor a 2-hour grace window to fix the PR before the next - triage run actually closes it. - - This is the "before-close" warning. On the second triage run, if the - grace marker is older than `GRACE_PERIOD_SECONDS` AND the PR still - fails the rubric, the close path runs (which posts - `format_pr_close_comment` and closes the PR). - """ - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_pr(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the description against our " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md). " - "Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "If the description isn't updated in the next **2 hours**, I'll auto-close this PR. " - "That's **not** us saying we don't care about the change; we want the open-PR list to " - "mirror what a maintainer can act on *right now*, so contributors don't get lost in a " - 'backlog. A closed PR is a soft "park this for later," not a rejection. Take your ' - "time; everything below still works after the close.\n" - "\n" - "**During the grace period:** just update the PR description with the missing pieces. " - "No need to ping me; I'll re-check on the next sweep and skip the auto-close if it " - "now passes. See " - "[what counts as QA proof](https://docs.litellm.ai/blog/agent-shin-triage#the-rubric-for-pull-requests) " - "for the full rubric (a linked issue alone isn't enough; it covers context, not proof).\n" - "\n" - "**If the PR does get auto-closed in 2 hours, you still have easy recovery paths:**\n" - "\n" - "- Comment `@agent-shin reconsider` after updating the description. I'll re-evaluate " - "and reopen the PR if it now passes.\n" - "- Comment `@greptileai` to request a fresh Greptile review; that **still works even " - "after the PR is closed**, and a stronger score is one of the signals that lifts the " - "PR back into the queue. So a low Greptile score isn't a blocker either.\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a " - "maintainer; they'll override me.)_\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -def format_grace_warning_issue_comment(verdict: dict) -> str: - """Issue analogue of `format_grace_warning_pr_comment`.""" - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_issue(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for filing this! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the issue against our reporting checklist. Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "If the issue isn't updated in the next **2 hours**, I'll auto-close it. That's **not** us " - "saying the bug isn't real or the request isn't useful; we want the open-issue list " - "to mirror what a maintainer can act on *right now*, so reports like yours don't get " - 'buried in a backlog. A closed issue is a soft "park this for later," not a ' - "rejection. Take your time; reopening is one comment away.\n" - "\n" - "**During the grace period:** just edit the issue description with the missing " - "pieces. No need to ping me; I'll re-check on the next sweep and skip the auto-close " - "if it now passes.\n" - "\n" - "Missing pieces, depending on what this is:\n" - "\n" - "- For **bug reports**: end-to-end evidence of the bug (a screen recording / video, a " - "screenshot, or the exact commands you ran with their real output / traceback) plus " - "expected vs. actual behavior. Written steps with no run output don't count, and " - "mocked or stubbed runs don't count.\n" - "- For **feature requests**: a concrete description of what should change, a use " - "case and example (config / API call / UI flow), plus end-to-end evidence of the " - "dead-end (a video, a screenshot, or the exact commands you ran with their real " - "output showing where the flow stops today). Mocked or stubbed runs don't count.\n" - "\n" - "**If the issue does get auto-closed in 2 hours**, comment `@agent-shin reconsider` " - "and I'll re-evaluate. If it now meets the bar, I'll reopen the issue.\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a " - "maintainer; they'll override me.)_\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -# --------------------------------------------------------------------------- -# Step-summary helpers - - -def write_step_summary(content: str) -> None: - """When running inside GitHub Actions, append to the step summary file.""" - path = os.environ.get("GITHUB_STEP_SUMMARY") - if not path: - return - try: - with open(path, "a", encoding="utf-8") as handle: - handle.write(content) - if not content.endswith("\n"): - handle.write("\n") - except OSError as exc: - print(f"warn: failed to write step summary: {exc}", file=sys.stderr) - - -# --------------------------------------------------------------------------- -# Core orchestration - - -def format_reopen_comment(kind: str) -> str: - """Comment posted when Agent Shin reopens after a successful reconsider.""" - noun = "PR" if kind == "pr" else "issue" - # The trailing HTML marker is used by `seconds_since_last_reconsider_verdict` - # to enforce a cooldown between repeated `@agent-shin reconsider` triggers. - # Keep the marker on its own line so it doesn't disturb the rendered text. - return ( - f"♻️ **Re-evaluated and reopened.** Thanks for updating the {noun}!\n" - "\n" - "Agent Shin re-ran triage on the latest description and it now meets " - "the bar. A maintainer will take another look soon; please don't " - f"close this {noun} again unless asked to.\n" - "\n" - "_(If a maintainer ends up closing this for non-rubric reasons, that " - "decision stands; comment `@agent-shin reconsider` again only if you " - "have substantively new information.)_\n" - "\n" - f"{RECONSIDER_COMMENT_MARKER}" - ) - - -def format_reconsider_still_failing_comment(kind: str, verdict: dict) -> str: - """Comment posted when reconsider re-runs triage but the verdict is still fail.""" - missing_lines = _format_missing(verdict.get("missing") or []) - explanation = verdict.get("explanation") or "" - noun = "PR" if kind == "pr" else "issue" - # The trailing HTML marker is used by `seconds_since_last_reconsider_verdict` - # to enforce a cooldown between repeated `@agent-shin reconsider` triggers. - return ( - f"⏸️ **Re-evaluated; this {noun} still doesn't meet the rubric.**\n" - "\n" - "Agent Shin re-ran triage on the current description but is still " - "missing:\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "Update the description with the missing pieces and comment " - "`@agent-shin reconsider` again, or ping a maintainer if you think " - "I got this wrong.\n" - "\n" - "_(I'm an LLM and I'm not infallible.)_\n" - "\n" - f"{RECONSIDER_COMMENT_MARKER}" - ) - - -# --------------------------------------------------------------------------- -# Review gate — "ready for review" label lifecycle - -_UNSET = object() - - -def _combine_missing( - verdict: dict, greptile_score: int | None, min_score: int -) -> list[str]: - """Merge the LLM rubric's `missing` list with a Greptile-score shortfall.""" - missing = list(verdict.get("missing") or []) - if greptile_score is not None and greptile_score < min_score: - missing.insert( - 0, - f"Greptile's most recent review scored this PR {greptile_score}/5 " - f"(below the {min_score}/5 bar)", - ) - return missing or ["(see explanation below)"] - - -def _has_marker( - comments: Iterable[dict], marker: str, *, bot_login: str | None = None -) -> bool: - """Return True iff the bot itself posted a comment containing ``marker``. - - Filters by author so a contributor who quotes the marker (e.g. via - GitHub's "Quote reply" feature, which preserves HTML comments in - raw markdown) is not mistaken for a bot action — that would - silently suppress notifications or change which "recovered" wording - is selected. Matches the author-filter pattern used by the sibling - `_seconds_since_latest_marker_comment` helper. - """ - expected_login = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - for comment in comments: - author = ((comment.get("user") or {}).get("login") or "").lower() - if author != expected_login: - continue - if marker in (comment.get("body") or ""): - return True - return False - - -def format_ready_for_review_comment( - verdict: dict, - greptile_score: int | None, - min_greptile_score: int = DEFAULT_MIN_GREPTILE_SCORE, -) -> str: - """Posted the first time a PR clears the bar (label added).""" - score_line = ( - f" Greptile scored it **{greptile_score}/5**." - if greptile_score is not None - else "" - ) - explanation = verdict.get("explanation") or "" - return ( - "✅ **Triage passed, tagging `ready for review`.**\n" - "\n" - "Agent Shin checked this PR against the " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md) " - "and it clears the bar (a linked issue, or a clear problem description " - f"+ expected vs. actual + QA proof).{score_line}\n" - "\n" - f"> {explanation}\n" - "\n" - "A maintainer will take it from here. If a later re-check finds the PR " - f"has regressed (Greptile drops below {min_greptile_score}/5, " - "the QA proof is removed, etc.) I'll pull the tag and comment with " - "what's missing; fix it and the tag comes back automatically.\n" - f"{READY_MARKER}" - ) - - -def format_all_clear_comment(verdict: dict, greptile_score: int | None) -> str: - """Posted when a PR recovers after a regression (label re-added).""" - score_line = ( - f" Greptile is back to **{greptile_score}/5**." - if greptile_score is not None - else "" - ) - explanation = verdict.get("explanation") or "" - return ( - "✅ **All clear again, re-adding `ready for review`.**\n" - "\n" - "Thanks for addressing the earlier feedback. On re-check this PR meets " - f"the contribution bar once more.{score_line}\n" - "\n" - f"> {explanation}\n" - "\n" - "A maintainer will take another look.\n" - f"{READY_MARKER}" - ) - - -def format_regression_comment( - missing: list[str], explanation: str, grace_days: int -) -> str: - """Posted when a previously-tagged PR regresses (label removed, PR stays open). - - Discloses the same ``grace_days`` deadline the state machine enforces: - once that window elapses with the PR still failing, the close path fires. - Hiding the deadline behind a bare "stays open" would surprise contributors - with an auto-close they were never warned about. - """ - window = "24 hours" if grace_days == 1 else f"{grace_days} days" - return ( - "⚠️ **Removing the `ready for review` tag.**\n" - "\n" - "On a re-check this PR no longer meets the contribution bar. What's " - "missing now:\n" - "\n" - f"{_format_missing(missing)}\n" - "\n" - f"> {explanation}\n" - "\n" - f"The PR stays open for ~{window}; address the points above and Agent " - 'Shin will post an "all clear" comment and re-add the tag ' - "automatically. If the points still aren't addressed after that " - "window, the PR is auto-closed; that's not a rejection, and you can " - "comment `@agent-shin reconsider` to have it re-evaluated and reopened " - "once it passes.\n" - f"{REGRESSED_MARKER}" - ) - - -def format_within_grace_comment( - missing: list[str], explanation: str, grace_days: int -) -> str: - """Posted once while a failing PR is still inside its grace window.""" - window = "24 hours" if grace_days == 1 else f"{grace_days} days" - return ( - "🚅 Hi, thanks for the PR! This is **Agent Shin**, the automated triage " - "bot. This PR doesn't quite meet the contribution bar yet:\n" - "\n" - f"{_format_missing(missing)}\n" - "\n" - f"> {explanation}\n" - "\n" - f"You have ~{window} from when this PR was opened to add the missing " - "pieces; just update the description and I'll re-check on the next " - "sweep. Once it passes I'll tag it `ready for review`. If it does get " - "auto-closed, that's not a rejection; comment `@agent-shin reconsider` " - "and I'll re-evaluate and reopen if it now passes.\n" - f"{WITHIN_GRACE_MARKER}" - ) - - -def review_gate( - *, - repo: str, - number: int, - close: bool, - model: str, - judge: Any = None, - greptile_score: Any = _UNSET, - comments: Any = _UNSET, - now: dt.datetime | None = None, - grace_days: int = DEFAULT_GRACE_DAYS, - min_greptile_score: int = DEFAULT_MIN_GREPTILE_SCORE, - label: str = READY_FOR_REVIEW_LABEL, - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> dict: - """Reconcile the `ready for review` label with a PR's current quality. - - A PR is *passing* when it clears BOTH gates: the LLM rubric (linked issue, - or problem description + expected/actual + QA proof) AND Greptile's most - recent confidence score (>= ``min_greptile_score``; absence of a score is - not held against the PR). The gate then drives a small state machine, using - the label itself as the persisted state so comments fire only on - transitions (never on every scheduled run): - - passing, untagged -> add label + "ready for review" / "all clear" - passing, tagged -> noop-passing - not passing, tagged -> remove label + regression comment (stays open) - not passing, untagged, old -> close + comment (past the grace window) - not passing, untagged, new -> one-time "what's missing" notice (within grace) - - ``close`` gates every destructive side effect: with ``close=False`` the - function returns a ``would-*`` preview and touches nothing, mirroring the - dry-run contract of :func:`triage`. ``judge``/``greptile_score``/ - ``comments``/``now`` are injectable for tests; in production they are - resolved from the OpenAI judge, the PR's Greptile comment, the live comment - list, and the wall clock respectively. - """ - item = fetch_pr(repo, number) - - title = item.get("title") or "" - body = item.get("body") or "" - login = (item.get("user") or {}).get("login") or "" - association = item.get("author_association") or "" - state = item.get("state") or "" - # GitHub label names are case-insensitive; compare lowercased so a repo - # that already has e.g. "Ready for Review" is recognized as the same - # label as our READY_FOR_REVIEW_LABEL constant ("ready for review"). - labels_now = {(lbl.get("name") or "").lower() for lbl in (item.get("labels") or [])} - label_key = label.lower() - created_raw = item.get("created_at") or "" - - base_result = { - "kind": "pr", - "number": number, - "title": title, - "author": login, - "author_association": association, - "state": state, - "labeled": label_key in labels_now, - "review_gate": True, - } - - if state != "open": - return {**base_result, "action": "skip-not-open"} - - if allowlist: - if login.lower() not in allowlist: - return {**base_result, "action": "skip-not-allowlisted"} - elif is_internal_contributor(item): - return {**base_result, "action": "skip-internal-author"} - - # Resolve the comment list once — used for both the Greptile score and the - # marker-based dedup below. - if comments is _UNSET: - comments = list(_iter_paginated_json(f"repos/{repo}/issues/{number}/comments")) - - # --- rubric verdict: linked-issue short-circuit, else the LLM judge ------- - if has_linked_issue(body): - verdict = { - "verdict": "pass", - "linked_issue": True, - "missing": [], - "explanation": "Linked-issue regex matched; LLM was not called.", - } - rubric_pass = True - else: - prompt = build_pr_prompt(title=title, body=body) - if judge is None: - api_key = os.environ.get("OPENAI_API_KEY") - if not api_key: - return {**base_result, "action": "skip-no-llm-key"} - base_url = os.environ.get("OPENAI_BASE_URL") or None - - def judge(p: str) -> str: - return call_llm_judge( - p, model=model, api_key=api_key, base_url=base_url - ) - - try: - verdict = parse_verdict(judge(prompt)) - except Exception as exc: # noqa: BLE001 - judge errors must never act - return {**base_result, "action": "skip-llm-error", "error": str(exc)} - rubric_pass = (verdict.get("verdict") or "").lower() == "pass" - - # --- Greptile score ------------------------------------------------------- - if greptile_score is _UNSET: - extraction = extract_greptile_score(comments) - greptile_score = extraction[0] if extraction else None - greptile_ok = greptile_score is None or greptile_score >= min_greptile_score - passing = rubric_pass and greptile_ok - - # --- age ------------------------------------------------------------------ - age_days = None - if created_raw: - reference = now or dt.datetime.now(dt.timezone.utc) - age_days = (reference - parse_iso8601(created_raw)).days - - label_present = label_key in labels_now - explanation = verdict.get("explanation") or "" - # When the rubric short-circuited to pass (linked-issue regex) but - # Greptile dragged the PR below the bar, the synthetic verdict's - # explanation ("LLM was not called") would mislead a contributor reading - # the regression / close comment. Surface the real reason instead. - if rubric_pass and not greptile_ok: - explanation = ( - f"Greptile's most recent review scored this PR " - f"{greptile_score}/5 (below the {min_greptile_score}/5 bar)." - ) - verdict = {**verdict, "explanation": explanation} - base_result = { - **base_result, - "verdict": verdict, - "greptile_score": greptile_score, - "passing": passing, - "age_days": age_days, - } - - if passing: - if label_present: - return {**base_result, "action": "noop-passing"} - recovered = _has_marker(comments, REGRESSED_MARKER) - comment = ( - format_all_clear_comment(verdict, greptile_score) - if recovered - else format_ready_for_review_comment( - verdict, greptile_score, min_greptile_score - ) - ) - if not close: - return {**base_result, "action": "would-label-ready", "comment": comment} - post_comment(repo, number, comment) - add_label(repo, number, label) - return {**base_result, "action": "labeled-ready", "comment": comment} - - missing = _combine_missing(verdict, greptile_score, min_greptile_score) - - if label_present: - comment = format_regression_comment(missing, explanation, grace_days) - if not close: - return {**base_result, "action": "would-remove-label", "comment": comment} - remove_label(repo, number, label) - post_comment(repo, number, comment) - return {**base_result, "action": "label-removed-regressed", "comment": comment} - - # Not passing and not tagged. If the PR was previously tagged and then - # regressed (we removed the label and posted REGRESSED_MARKER), honor the - # "PR stays open — fix it and the tag comes back" promise from - # `format_regression_comment` and skip the close path. Without this guard, - # any PR older than `grace_days` would be closed on the next evaluation, - # giving the contributor no realistic window to address the regression. - # - # The promise has a deliberate expiration: once `grace_days` have elapsed - # since the regression notice, fall through to the close path so a PR that - # was abandoned post-regression doesn't sit open forever. - if _has_marker(comments, REGRESSED_MARKER): - reference = now or dt.datetime.now(dt.timezone.utc) - seconds_since_regression = seconds_since_latest_marker_comment( - comments, marker=REGRESSED_MARKER, now=reference - ) - grace_seconds = grace_days * 86400 - if seconds_since_regression is None or seconds_since_regression < grace_seconds: - return {**base_result, "action": "regressed-already-notified"} - - # Not passing and not tagged: close if past the grace window, else notify once. - if age_days is not None and age_days >= grace_days: - comment = format_pr_close_comment({**verdict, "missing": missing}) - if not close: - return {**base_result, "action": "would-close", "comment": comment} - post_comment(repo, number, comment) - close_pr(repo, number) - return {**base_result, "action": "closed", "comment": comment} - - if _has_marker(comments, WITHIN_GRACE_MARKER): - return {**base_result, "action": "within-grace-already-notified"} - comment = format_within_grace_comment(missing, explanation, grace_days) - if not close: - return { - **base_result, - "action": "would-notify-within-grace", - "comment": comment, - } - post_comment(repo, number, comment) - return {**base_result, "action": "within-grace-notified", "comment": comment} - - -def triage( - *, - repo: str, - kind: str, - number: int, - close: bool, - model: str, - judge: Any = None, - print_prompt: bool = False, - reconsider: bool = False, - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> dict: - """Triage a single PR or issue. Returns a result dict for logging/tests. - - `judge` is an optional callable `(prompt) -> str` for tests / dry-run with - a stub. In production, leave it None and the script uses `call_llm_judge`. - - When `reconsider=True`, the closed-state guard is skipped and a - fail-but-no-comment is replaced with a "still failing" comment + leave - closed; a pass triggers `reopen_pr`/`reopen_issue` plus a reopen comment. - Reconsider mode is intended for the `@agent-shin reconsider` comment - trigger. Like regular triage, `close=False` keeps reconsider in dry-run - (returns `would-reopen` / `would-reconsider-still-failing` so a local - operator can preview without write side effects); the workflow only - passes `--close` when `AGENT_SHIN_ENABLED=true`. - - Reconsider mode adds two extra safety guards on top of the regular - triage skip-internal-author check: - - 1. **Bot-closed guard.** Only reopens if the most recent close was - performed by the bot identity (default `github-actions[bot]`). - This stops a contributor from using `@agent-shin reconsider` to - override a maintainer's close for non-rubric reasons. - 2. **Rate-limit guard.** If the bot has already posted a reconsider - verdict on this PR/issue within `RECONSIDER_RATE_LIMIT_SECONDS`, - skip — repeated triggers from the same contributor shouldn't burn - CI minutes or LLM budget. - """ - fetcher = {"pr": fetch_pr, "issue": fetch_issue}[kind] - item = fetcher(repo, number) - - title = item.get("title") or "" - body = item.get("body") or "" - login = (item.get("user") or {}).get("login") or "" - association = item.get("author_association") or "" - state = item.get("state") or "" - - base_result = { - "kind": kind, - "number": number, - "title": title, - "author": login, - "author_association": association, - "state": state, - "reconsider": reconsider, - } - - # Reconsider only makes sense on a closed PR/issue. A "reconsider on an - # open PR" is a no-op (the regular triage flow already evaluates open - # PRs); return a clear skip so the workflow can short-circuit. - if reconsider: - if state != "closed": - return {**base_result, "action": "skip-not-closed"} - else: - if state != "open": - return {**base_result, "action": "skip-not-open"} - - if allowlist: - if login.lower() not in allowlist: - return {**base_result, "action": "skip-not-allowlisted"} - elif is_internal_contributor(item): - return {**base_result, "action": "skip-internal-author"} - - # Reconsider-only guards — these run BEFORE the LLM call so a - # maintainer-closed PR / rate-limited trigger never spends LLM budget. - if reconsider: - if not was_closed_by_agent_shin(repo, number): - return {**base_result, "action": "skip-not-bot-closed"} - age = seconds_since_last_reconsider_verdict(repo, number) - if age is not None and age < RECONSIDER_RATE_LIMIT_SECONDS: - return { - **base_result, - "action": "skip-rate-limited", - "rate_limit_age_seconds": age, - "rate_limit_window_seconds": RECONSIDER_RATE_LIMIT_SECONDS, - } - - if kind == "pr": - # Short-circuit: if body very clearly links a related issue, just pass. - if has_linked_issue(body): - base = { - **base_result, - "action": "pass-linked-issue", - "verdict": { - "verdict": "pass", - "linked_issue": True, - "explanation": "Linked-issue regex matched; LLM was not called.", - }, - } - if reconsider: - # Pass-on-reconsider -> reopen the PR with a friendly comment. - reopen_body = format_reopen_comment(kind) - if not close: - return { - **base, - "action": "would-reopen", - "comment": reopen_body, - } - post_comment(repo, number, reopen_body) - reopen_pr(repo, number) - return { - **base, - "action": "reopened", - "comment": reopen_body, - } - return base - prompt = build_pr_prompt(title=title, body=body) - else: - prompt = build_issue_prompt(title=title, body=body) - - if print_prompt: - return {**base_result, "action": "print-prompt", "prompt": prompt} - - if judge is None: - api_key = os.environ.get("OPENAI_API_KEY") - if not api_key: - # No key configured — never take a destructive action. Report skip. - return { - **base_result, - "action": "skip-no-llm-key", - "prompt_preview": prompt[:200], - } - base_url = os.environ.get("OPENAI_BASE_URL") or None - - def judge(p: str) -> str: - return call_llm_judge(p, model=model, api_key=api_key, base_url=base_url) - - try: - raw = judge(prompt) - verdict = parse_verdict(raw) - except Exception as exc: # noqa: BLE001 - judge errors must never close PRs - return {**base_result, "action": "skip-llm-error", "error": str(exc)} - - decision = (verdict.get("verdict") or "").lower() - - if reconsider: - # Reconsider: an explicit `pass` -> reopen + post reopen comment; - # anything else (fail, missing/malformed verdict, typo) -> leave - # closed + post a "still failing" comment so the contributor can - # iterate again. Reopen is destructive, so a flaky/empty verdict - # must not satisfy the gate. - # In dry-run (`close=False`) we return `would-*` actions instead - # of touching GitHub state, mirroring the regular triage flow's - # `would-close`. This lets a local operator preview the outcome - # of `python triage_with_llm.py --reconsider --pr N` without - # risking accidental comments or reopens. - if decision == "pass": - reopen_body = format_reopen_comment(kind) - if not close: - return { - **base_result, - "action": "would-reopen", - "verdict": verdict, - "comment": reopen_body, - } - post_comment(repo, number, reopen_body) - if kind == "pr": - reopen_pr(repo, number) - else: - reopen_issue(repo, number) - return { - **base_result, - "action": "reopened", - "verdict": verdict, - "comment": reopen_body, - } - still_failing = format_reconsider_still_failing_comment(kind, verdict) - if not close: - return { - **base_result, - "action": "would-reconsider-still-failing", - "verdict": verdict, - "comment": still_failing, - } - post_comment(repo, number, still_failing) - return { - **base_result, - "action": "reconsider-still-failing", - "verdict": verdict, - "comment": still_failing, - } - - if decision != "fail": - return {**base_result, "action": "pass-llm", "verdict": verdict} - - # Grace-period flow: on the first low-quality detection, post a warning - # comment instead of closing immediately. On a subsequent triage run - # (manual re-trigger, or the daily `close_low_quality_prs.py` cron - # finding the same PR in its own pass), if `GRACE_PERIOD_SECONDS` has - # elapsed since the warning AND the PR still fails the rubric, close. - grace_age = seconds_since_last_grace_warning(repo, number) - if grace_age is None: - warning_body = ( - format_grace_warning_pr_comment(verdict) - if kind == "pr" - else format_grace_warning_issue_comment(verdict) - ) - if not close: - return { - **base_result, - "action": "would-warn-grace", - "verdict": verdict, - "comment": warning_body, - } - post_comment(repo, number, warning_body) - return { - **base_result, - "action": "warned-grace", - "verdict": verdict, - "comment": warning_body, - } - if grace_age < GRACE_PERIOD_SECONDS: - return { - **base_result, - "action": "skip-in-grace-period", - "verdict": verdict, - "grace_age_seconds": grace_age, - "grace_period_seconds": GRACE_PERIOD_SECONDS, - } - - # The grace window has elapsed. `--close` still gates the destructive - # write so a dry-run preview never posts or closes — the workflow only - # passes `--close` when `AGENT_SHIN_ENABLED=true`, which keeps the bot - # inert by default. - if not close: - return {**base_result, "action": "would-close", "verdict": verdict} - - comment_body = ( - format_pr_close_comment(verdict) - if kind == "pr" - else format_issue_close_comment(verdict) - ) - post_comment(repo, number, comment_body) - if kind == "pr": - close_pr(repo, number) - else: - close_issue(repo, number) - - return { - **base_result, - "action": "closed", - "verdict": verdict, - "comment": comment_body, - } - - -# --------------------------------------------------------------------------- -# CLI - - -def render_summary(result: dict) -> str: - """Render a human-readable summary block (used for stdout + step summary).""" - lines = ["## Agent Shin verdict", ""] - lines.append( - f"- **{result['kind'].upper()} #{result['number']}**: {result.get('title', '')}" - ) - lines.append( - f"- **Author**: `{result.get('author', '')}` ({result.get('author_association', '')})" - ) - lines.append(f"- **State**: {result.get('state', '')}") - lines.append(f"- **Action**: `{result['action']}`") - verdict = result.get("verdict") - if verdict: - lines.append("") - lines.append("```json") - lines.append(json.dumps(verdict, indent=2)) - lines.append("```") - error = result.get("error") - if error: - lines.append("") - lines.append(f"_LLM error: {error}_") - comment = result.get("comment") - if comment: - lines.append("") - lines.append("### Posted comment:") - lines.append("") - lines.append("> " + comment.replace("\n", "\n> ")) - return "\n".join(lines) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--repo", required=True, help="Repository (owner/repo).") - target = parser.add_mutually_exclusive_group(required=True) - target.add_argument("--pr", type=int, help="Pull request number to triage.") - target.add_argument("--issue", type=int, help="Issue number to triage.") - parser.add_argument( - "--close", - action="store_true", - help="Actually post comment + close on fail (default: dry run).", - ) - parser.add_argument( - "--model", - # `os.environ.get("TRIAGE_MODEL", DEFAULT_MODEL)` would return "" when - # GitHub Actions exposes an unset repo variable as an empty-string env - # var, silently bypassing DEFAULT_MODEL and causing every call to fail - # as `skip-llm-error`. The `or` guard collapses empty -> default. - default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL, - help=f"OpenAI-compatible model name (default: {DEFAULT_MODEL}).", - ) - parser.add_argument( - "--print-prompt", - action="store_true", - help="Print the prompt that would be sent to the judge and exit.", - ) - parser.add_argument( - "--reconsider", - action="store_true", - help=( - "Re-run triage on a CLOSED PR/issue and reopen it on pass. " - "Used by the `@agent-shin reconsider` comment-trigger workflow. " - "Only invoke this from a workflow that has already gated on " - "AGENT_SHIN_ENABLED=true and verified the commenter is the " - "PR/issue author or an internal collaborator." - ), - ) - parser.add_argument( - "--review-gate", - action="store_true", - help=( - "Reconcile the `ready for review` label for an OPEN PR: tag on " - "pass, remove the tag + comment on regression, close after the " - "grace window if it never passed. PR-only." - ), - ) - parser.add_argument( - "--grace-days", - type=int, - default=DEFAULT_GRACE_DAYS, - help=( - "Review-gate only: hours/24 a failing, un-tagged PR may stay open " - f"before auto-close (default: {DEFAULT_GRACE_DAYS} = 24h)." - ), - ) - parser.add_argument( - "--min-greptile-score", - type=int, - default=DEFAULT_MIN_GREPTILE_SCORE, - choices=range(1, 6), - help=( - "Review-gate only: Greptile score below which a PR counts as not " - f"passing (default: {DEFAULT_MIN_GREPTILE_SCORE} -> <4/5 regresses)." - ), - ) - args = parser.parse_args() - - kind = "pr" if args.pr is not None else "issue" - number = args.pr if args.pr is not None else args.issue - - if args.review_gate: - if kind != "pr": - parser.error("--review-gate applies to pull requests only (use --pr).") - result = review_gate( - repo=args.repo, - number=number, - close=args.close, - model=args.model, - grace_days=args.grace_days, - min_greptile_score=args.min_greptile_score, - ) - else: - result = triage( - repo=args.repo, - kind=kind, - number=number, - close=args.close, - model=args.model, - print_prompt=args.print_prompt, - reconsider=args.reconsider, - ) - - if result.get("action") == "print-prompt": - print(result["prompt"]) - return 0 - - summary = render_summary(result) - print(summary) - write_step_summary(summary + "\n") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/close_low_quality_prs.yml b/.github/workflows/close_low_quality_prs.yml deleted file mode 100644 index 2401be84000..00000000000 --- a/.github/workflows/close_low_quality_prs.yml +++ /dev/null @@ -1,92 +0,0 @@ -name: Close Low-Quality PRs - -# Auto-close any open PR (including drafts, regardless of age) authored by an -# external OSS contributor that Greptile reviewed with a confidence score -# below 4/5. Closures are explained in a comment that tells the contributor -# to push fixes and open a fresh PR (since OSS authors cannot reopen a PR -# closed by a bot/maintainer) or comment `@agent-shin reconsider` to have -# Agent Shin re-evaluate. -# -# Manual one-off run: -# gh workflow run "Close Low-Quality PRs" -f close=true -# -# Dry-run preview (no PRs are touched): -# gh workflow run "Close Low-Quality PRs" -f close=false - -on: - schedule: - # Daily at 09:00 UTC. Pairs well with the stale-issue workflow at midnight. - - cron: "0 9 * * *" - workflow_dispatch: - inputs: - close: - description: "Actually close matching PRs (false = dry run)." - required: false - default: "false" - type: choice - options: - - "true" - - "false" - min_age_days: - description: "Minimum PR age in days (default 0 = no age filter)." - required: false - default: "0" - min_score: - description: "Greptile score below which a PR is closed (1-5)." - required: false - default: "4" - limit: - description: "Maximum number of PRs to close in a single run." - required: false - default: "25" - -permissions: - contents: read - pull-requests: write - issues: write - -jobs: - close-low-quality-prs: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - steps: - - name: Checkout triage script - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Run low-quality PR closer - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # Scheduled runs are ALWAYS dry-run, even when AGENT_SHIN_ENABLED is - # "true", so the team can QA the closer's verdicts in step summaries - # before any contributor sees a PR closed. Real closures only happen - # on manual workflow_dispatch with close=true (and the variable set). - CLOSE_FLAG: ${{ github.event.inputs.close || 'false' }} - AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }} - MIN_AGE_DAYS: ${{ github.event.inputs.min_age_days || '0' }} - MIN_SCORE: ${{ github.event.inputs.min_score || '4' }} - LIMIT: ${{ github.event.inputs.limit || '25' }} - run: | - set -euo pipefail - ARGS=( - --repo "${{ github.repository }}" - --min-age-days "${MIN_AGE_DAYS}" - --min-score "${MIN_SCORE}" - --limit "${LIMIT}" - ) - if [ "${AGENT_SHIN_ENABLED:-false}" != "true" ]; then - echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> forcing dry-run regardless of close input." - elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${CLOSE_FLAG}" = "true" ]; then - ARGS+=(--close) - echo "::notice::Running in close-on-fail mode." - else - echo "::notice::AGENT_SHIN_ENABLED is true but this trigger is dry-run (scheduled event or close=false)." - fi - python3 .github/scripts/close_low_quality_prs.py "${ARGS[@]}" diff --git a/.github/workflows/create_daily_oss_agent_shin_branch.yml b/.github/workflows/create_daily_oss_agent_shin_branch.yml deleted file mode 100644 index 9baf9f142f6..00000000000 --- a/.github/workflows/create_daily_oss_agent_shin_branch.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: Create Daily oss-agent-shin Branch - -on: - schedule: - - cron: "0 0 * * *" # Runs every day at midnight UTC - workflow_dispatch: # Allow manual trigger - -jobs: - create-oss-agent-shin-branch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - permissions: - contents: write - - steps: - - name: Create daily oss-agent-shin branch - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')" - echo "Creating branch: $BRANCH_NAME" - if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then - echo "Branch $BRANCH_NAME already exists. Skipping creation." - exit 0 - fi - MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha') - gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent - echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA" diff --git a/.github/workflows/triage_reconsider.yml b/.github/workflows/triage_reconsider.yml deleted file mode 100644 index f35f681d09a..00000000000 --- a/.github/workflows/triage_reconsider.yml +++ /dev/null @@ -1,172 +0,0 @@ -name: Agent Shin — reconsider - -# Comment-trigger workflow: when the PR/issue author (or an internal -# collaborator) comments `@agent-shin reconsider` on a CLOSED PR/issue, -# Agent Shin re-runs LLM-judge triage on the current title+body and: -# -# - on PASS: posts a "re-evaluated and reopened" comment + reopens. -# - on FAIL: posts a "still missing X" comment and leaves it closed, -# so the contributor can iterate again. -# -# This exists because GitHub does NOT let an external (non-write-access) -# OSS contributor reopen a PR/issue closed by a bot or maintainer. Without -# this comment trigger, a contributor whose PR Agent Shin auto-closed -# would have no path back into the review queue except opening a fresh PR -# (which loses the original PR's history). The bot, on the other hand, -# has write access via GH_TOKEN and can reopen on their behalf. -# -# DRY-RUN BY DEFAULT — gated on `vars.AGENT_SHIN_ENABLED == 'true'` just -# like the other Agent Shin workflows. The workflow also gates on the -# commenter being either the PR/issue author or an internal collaborator -# (OWNER/MEMBER/COLLABORATOR) so random commenters cannot DOS the LLM -# judge or force a reopen. - -on: - issue_comment: - types: [created] - -permissions: - contents: read - issues: write - pull-requests: write - -jobs: - reconsider: - if: | - github.repository == 'BerriAI/litellm' - && contains(github.event.comment.body, '@agent-shin reconsider') - runs-on: ubuntu-latest - steps: - - name: Authorize commenter - # Only the PR/issue author OR an internal collaborator may trigger - # a reconsider. Outside random commenters could otherwise spam the - # phrase to burn LLM budget or, if a fail-open bug were ever - # introduced, force a reopen on someone else's behalf. - # - # We expose the authorization decision as a step output and gate - # every subsequent (potentially destructive) step on it. A `run:` - # step with `exit 0` would NOT stop the job — only `if:` gating - # on a known-true output is safe here. - id: auth - env: - COMMENTER: ${{ github.event.comment.user.login }} - AUTHOR: ${{ github.event.issue.user.login }} - ASSOCIATION: ${{ github.event.comment.author_association }} - run: | - set -euo pipefail - if [ "${COMMENTER}" = "${AUTHOR}" ]; then - echo "::notice::Authorized: commenter is the PR/issue author." - echo "authorized=true" >> "$GITHUB_OUTPUT" - exit 0 - fi - case "${ASSOCIATION}" in - OWNER|MEMBER|COLLABORATOR) - echo "::notice::Authorized: commenter is an internal collaborator (${ASSOCIATION})." - echo "authorized=true" >> "$GITHUB_OUTPUT" - ;; - *) - echo "::notice::Commenter '${COMMENTER}' (${ASSOCIATION}) is not authorized to trigger reconsider; skipping subsequent steps." - echo "authorized=false" >> "$GITHUB_OUTPUT" - ;; - esac - - - name: React 👀 to acknowledge the reconsider - # Add an eyes reaction to the triggering comment the moment we accept - # it, so the contributor gets instant feedback that the bot saw their - # `@agent-shin reconsider` before the slower triage steps run. Gated on - # AGENT_SHIN_ENABLED so dry-run leaves no visible trace. Best-effort: - # a reactions API hiccup must never fail the actual reconsider. - if: steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMMENT_ID: ${{ github.event.comment.id }} - run: | - set -euo pipefail - gh api --method POST \ - -H "Accept: application/vnd.github+json" \ - "repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \ - -f content=eyes \ - || echo "::warning::failed to add 👀 reaction (non-fatal)" - - - name: Checkout triage script - if: steps.auth.outputs.authorized == 'true' - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - if: steps.auth.outputs.authorized == 'true' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Install LLM client - if: steps.auth.outputs.authorized == 'true' - run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt - - - name: Run Agent Shin reconsider - if: steps.auth.outputs.authorized == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # Only expose the LLM key when the bot is enabled, so a PR/issue - # author can't force paid LLM calls by spamming `@agent-shin - # reconsider` while the bot is still in dry-run. The Python script - # calls the LLM whenever this var is set (regardless of `--close`); - # stripping `--close` doesn't suppress the API call, only the - # destructive side effects. Mirror the gating used by every other - # Agent Shin workflow (triage_pr_with_llm.yml, review_gate.yml, ...). - OPENAI_API_KEY: ${{ vars.AGENT_SHIN_ENABLED == 'true' && secrets.OPENAI_API_KEY || '' }} - OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }} - TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }} - AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }} - # `issue_comment` events fire for both issues and PR comments. - # `issue.pull_request` is set iff this is a PR comment, so we use - # its presence to decide whether to invoke `--pr N` or `--issue N`. - IS_PR: ${{ github.event.issue.pull_request != null }} - NUMBER: ${{ github.event.issue.number }} - run: | - set -euo pipefail - if [ "${IS_PR}" = "true" ]; then - ARGS=(--repo "${{ github.repository }}" --pr "${NUMBER}" --reconsider) - else - ARGS=(--repo "${{ github.repository }}" --issue "${NUMBER}" --reconsider) - fi - # Reconsider's destructive actions (post comment + reopen) are - # gated on `--close`, mirroring the regular triage workflows. - # When AGENT_SHIN_ENABLED is not the EXACT string "true", we - # still run the script so its verdict + would-X action lands in - # the step summary for QA — but without `--close`, the script - # returns `would-reopen` / `would-reconsider-still-failing` - # instead of touching GitHub state. - # - # Use the positive `= "true"` gate (not `!= "true" -> exit`) so - # the workflow guardrails in - # tests/test_litellm/test_github_triage_workflows.py see the - # canonical fail-safe enable pattern. Unknown values like - # "True", "yes", "1", or typos fall through to the dry-run - # branch, which is the safe default. - if [ "${AGENT_SHIN_ENABLED:-false}" = "true" ]; then - ARGS+=(--close) - echo "::notice::Agent Shin reconsider ENABLED — running real triage (close=true)." - else - echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> reconsider stays in dry-run (no comment, no reopen)." - fi - python3 .github/scripts/triage_with_llm.py "${ARGS[@]}" - - - name: React 👍 when the reconsider finishes - # Once the reconsider run has completed successfully, add a thumbs-up so - # the contributor sees the bot is done (the 👀 stays, signalling - # seen -> handled). `success()` keeps this from firing if the run - # errored, and the AGENT_SHIN_ENABLED gate keeps dry-run inert. - if: success() && steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMMENT_ID: ${{ github.event.comment.id }} - run: | - set -euo pipefail - gh api --method POST \ - -H "Accept: application/vnd.github+json" \ - "repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \ - -f content=+1 \ - || echo "::warning::failed to add 👍 reaction (non-fatal)" diff --git a/tests/test_litellm/test_github_close_low_quality_prs.py b/tests/test_litellm/test_github_close_low_quality_prs.py deleted file mode 100644 index 2a891ca72f5..00000000000 --- a/tests/test_litellm/test_github_close_low_quality_prs.py +++ /dev/null @@ -1,856 +0,0 @@ -"""Unit tests for `.github/scripts/close_low_quality_prs.py`. - -These exercise the pure logic (score extraction and per-PR evaluation) without -hitting GitHub. Network/CLI calls are stubbed via monkeypatch. -""" - -from __future__ import annotations - -import datetime as dt -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / ".github" - / "scripts" - / "close_low_quality_prs.py" -) - - -@pytest.fixture(scope="module") -def closer_module(): - """Load the script as a module via its file path (it lives outside the package).""" - spec = importlib.util.spec_from_file_location("close_low_quality_prs", SCRIPT_PATH) - assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" - module = importlib.util.module_from_spec(spec) - sys.modules["close_low_quality_prs"] = module - spec.loader.exec_module(module) - return module - - -def _greptile_comment( - body: str, - updated_at: str = "2026-05-10T00:00:00Z", - login: str = "greptile-apps[bot]", -) -> dict: - return { - "user": {"login": login}, - "body": body, - "created_at": updated_at, - "updated_at": updated_at, - } - - -class TestExtractGreptileScore: - def test_should_extract_score_from_html_header(self, closer_module): - comments = [ - _greptile_comment("

Confidence Score: 3/5

\nSome body text.") - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 3 - - def test_should_accept_both_greptile_login_variants(self, closer_module): - # REST API form ("greptile-apps[bot]") and GraphQL form ("greptile-apps") - for login in ("greptile-apps", "greptile-apps[bot]"): - comments = [ - _greptile_comment("

Confidence Score: 2/5

", login=login) - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None, f"failed to detect score for login={login}" - score, _ = result - assert score == 2 - - def test_should_extract_score_from_plain_text(self, closer_module): - comments = [_greptile_comment("Confidence Score: 5/5 — looks good!")] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 5 - - def test_should_tolerate_whitespace_and_case(self, closer_module): - comments = [_greptile_comment("**confidence score : 2 / 5**")] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 2 - - def test_should_pick_most_recent_comment_when_rereview_happens(self, closer_module): - comments = [ - _greptile_comment( - "Confidence Score: 2/5", updated_at="2026-05-01T00:00:00Z" - ), - _greptile_comment( - "Confidence Score: 5/5", updated_at="2026-05-12T00:00:00Z" - ), - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 5 - - def test_should_ignore_non_greptile_authors(self, closer_module): - comments = [ - { - "user": {"login": "some-human"}, - "body": "Confidence Score: 1/5", - "created_at": "2026-05-12T00:00:00Z", - "updated_at": "2026-05-12T00:00:00Z", - } - ] - assert closer_module.extract_greptile_score(comments) is None - - def test_should_return_none_when_no_score_present(self, closer_module): - comments = [_greptile_comment("Greptile summary without a score.")] - assert closer_module.extract_greptile_score(comments) is None - - def test_should_return_none_for_empty_comments(self, closer_module): - assert closer_module.extract_greptile_score([]) is None - - -class TestEvaluatePr: - @pytest.fixture(autouse=True) - def _now(self): - return dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - - def _make_pr( - self, - *, - number: int = 1, - created_days_ago: int = 10, - is_draft: bool = False, - labels: list[str] | None = None, - author_login: str = "mateo-berri", - ) -> dict: - created = dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - dt.timedelta( - days=created_days_ago - ) - return { - "number": number, - "title": f"PR #{number}", - "createdAt": created.isoformat().replace("+00:00", "Z"), - "isDraft": is_draft, - "labels": [{"name": lbl} for lbl in (labels or [])], - "author": {"login": author_login}, - "url": f"https://example.com/pr/{number}", - } - - @pytest.fixture(autouse=True) - def _external_author(self, closer_module, monkeypatch): - """Treat every test PR as external unless overridden.""" - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: True - ) - - def test_should_warn_drafts_when_score_low_first_time( - self, closer_module, _now, monkeypatch - ): - # Drafts are NOT a free pass — the open-PR queue should reflect any - # PR that needs human attention regardless of draft status. Authors - # who need a long-lived draft can use the `wip` opt-out label. - # First run: warn the contributor (1-day grace), don't close yet. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 2/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(is_draft=True, created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 2 and age == 0 - - def test_should_warn_brand_new_pr_when_min_age_zero( - self, closer_module, _now, monkeypatch - ): - # `min_age_days=0` means no age filter — a freshly-opened PR is - # eligible the moment Greptile scores it below threshold. The - # first detection still goes through the warn-grace step rather - # than closing immediately, giving the contributor 2 hours to - # respond before the next run actually closes the PR. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 1/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 1 and age == 0 - - def test_should_skip_optout_label_case_insensitive( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for opt-outs"), - ) - action, _, _ = closer_module.evaluate_pr( - self._make_pr(labels=["WIP"]), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels={"wip"}, - ) - assert action == "skip-optout-label" - - def test_should_skip_too_young_when_min_age_set( - self, closer_module, _now, monkeypatch - ): - # The min-age-days flag is now opt-in (default 0). When a maintainer - # explicitly passes a positive value (e.g. for a backfill run that - # wants to spare brand-new PRs), the skip-too-young path still works. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for young PRs"), - ) - action, _, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=2), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-too-young" - assert age == 2 - - def test_should_not_skip_when_min_age_is_zero( - self, closer_module, _now, monkeypatch - ): - # With the new default min_age_days=0, even a 0-day-old PR is - # evaluated. This test pins that behavior so future refactors don't - # silently restore an age filter. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 5/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-score-ok" - assert score == 5 and age == 0 - - def test_should_skip_when_greptile_has_not_reviewed( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr(closer_module, "fetch_pr_comments", lambda *a, **kw: []) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-no-greptile-score" - assert score is None and age == 10 - - def test_should_skip_when_score_meets_threshold( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 4/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-score-ok" - assert score == 4 and age == 10 - - def test_should_warn_when_old_and_low_score_no_prior_warning( - self, closer_module, _now, monkeypatch - ): - # Even an old PR that still has no grace warning gets one on the - # first eligible run — the daily cron is the natural cadence, so - # an existing-but-never-warned PR enters the grace flow normally. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 3/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 3 and age == 10 - - def test_should_close_when_grace_warning_aged_out_and_score_still_low( - self, closer_module, _now, monkeypatch - ): - # Day-1 the closer posted a warning. Day-2 the PR still scores <4 - # AND the warning is older than `GRACE_PERIOD_SECONDS`, so the - # action flips to `close`. This is the "grace expired" path. - old_warning = { - "user": {"login": "github-actions[bot]"}, - "body": ( - "you have 2 hours to fix this\n\n" + closer_module.GRACE_COMMENT_MARKER - ), - "created_at": ( - _now - dt.timedelta(seconds=closer_module.GRACE_PERIOD_SECONDS + 60) - ) - .isoformat() - .replace("+00:00", "Z"), - "updated_at": "2026-05-15T00:00:00Z", - } - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment( - "

Confidence Score: 1/5

", - updated_at="2026-05-15T00:00:00Z", - ), - old_warning, - ], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=14), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "close" - assert score == 1 - - def test_should_skip_when_grace_warning_within_window( - self, closer_module, _now, monkeypatch - ): - # Within the 2-hour grace window the closer must NOT close the - # PR even if the score is still low. The warning is only an hour - # old; give the contributor time to push fixes before destruction. - recent_warning = { - "user": {"login": "github-actions[bot]"}, - "body": "warning text\n\n" + closer_module.GRACE_COMMENT_MARKER, - "created_at": (_now - dt.timedelta(hours=1)) - .isoformat() - .replace("+00:00", "Z"), - } - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment("Confidence Score: 2/5"), - recent_warning, - ], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-in-grace-period" - assert score == 2 - - def test_should_warn_grace_for_swiftwinds_not_close_immediately( - self, closer_module, _now, monkeypatch - ): - # Regression: SwiftWinds (the dogfood account) used to be in a - # now-removed `IMMEDIATE_CLOSE_LOGINS` bypass that closed on first - # detection. It must now follow the SAME grace path as every other - # external author: warn first, close only after the window elapses. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 1/5")], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0, author_login="SwiftWinds"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 1 - - def test_should_skip_internal_authors(self, closer_module, _now, monkeypatch): - # Override the fixture for this one test. - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for internal"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=14, author_login="krrishdholakia"), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - allowlist=frozenset(), - ) - assert action == "skip-internal" - assert score is None - - -class TestMainOptoutLabelDefault: - """`--optout-label` must REPLACE the canonical defaults, not append.""" - - def _patch_no_op(self, closer_module, monkeypatch): - monkeypatch.setattr(closer_module, "fetch_open_prs", lambda repo: []) - # `optout_labels` is captured indirectly via evaluate_pr; sniff the - # set passed in by stubbing evaluate_pr. - captured: dict = {} - - def fake_evaluate(pr, now, min_age_days, min_score, repo, optout_labels): - captured["optout_labels"] = set(optout_labels) - return ("skip-internal", None, None) - - monkeypatch.setattr(closer_module, "evaluate_pr", fake_evaluate) - return captured - - def test_should_use_canonical_defaults_when_flag_omitted( - self, closer_module, monkeypatch - ): - captured = self._patch_no_op(closer_module, monkeypatch) - # No PRs -> capture won't fire; instead inject one synthetic PR via - # fetch_open_prs so evaluate_pr is invoked at least once. - monkeypatch.setattr( - closer_module, - "fetch_open_prs", - lambda repo: [ - { - "number": 1, - "title": "p", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": True, - "labels": [], - "author": {"login": "x"}, - } - ], - ) - monkeypatch.setattr(sys, "argv", ["close_low_quality_prs.py"]) - rc = closer_module.main() - assert rc == 0 - assert captured["optout_labels"] == set(closer_module.DEFAULT_OPTOUT_LABELS) - - def test_should_replace_defaults_when_flag_provided( - self, closer_module, monkeypatch - ): - captured = self._patch_no_op(closer_module, monkeypatch) - monkeypatch.setattr( - closer_module, - "fetch_open_prs", - lambda repo: [ - { - "number": 1, - "title": "p", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": True, - "labels": [], - "author": {"login": "x"}, - } - ], - ) - monkeypatch.setattr( - sys, - "argv", - [ - "close_low_quality_prs.py", - "--optout-label", - "hold", - "--optout-label", - "needs-discussion", - ], - ) - rc = closer_module.main() - assert rc == 0 - # Crucially, none of the canonical defaults leak in. - assert captured["optout_labels"] == {"hold", "needs-discussion"} - for default in closer_module.DEFAULT_OPTOUT_LABELS: - assert default not in captured["optout_labels"], default - - -class TestSecondsSinceLastGraceWarning: - """Grace-period detection: only counts comments by the bot identity - that contain the shared `GRACE_COMMENT_MARKER`.""" - - def _make_marker_comment( - self, - closer_module, - *, - login: str = "github-actions[bot]", - created_at: str = "2026-05-16T00:00:00Z", - include_marker: bool = True, - ) -> dict: - body = "warning text" - if include_marker: - body += "\n\n" + closer_module.GRACE_COMMENT_MARKER - return { - "user": {"login": login}, - "body": body, - "created_at": created_at, - } - - def test_should_return_none_when_no_marker_comment(self, closer_module): - comments = [ - { - "user": {"login": "github-actions[bot]"}, - "body": "Some other bot comment", - "created_at": "2026-05-16T00:00:00Z", - } - ] - assert closer_module.seconds_since_last_grace_warning(comments) is None - - def test_should_return_none_for_empty(self, closer_module): - assert closer_module.seconds_since_last_grace_warning([]) is None - - def test_should_ignore_non_bot_comments_with_marker(self, closer_module): - # If a curious user quotes the marker in a comment, we must NOT - # treat it as a bot warning. The grace timer would then never fire. - comments = [ - self._make_marker_comment(closer_module, login="random-user"), - ] - assert closer_module.seconds_since_last_grace_warning(comments) is None - - def test_should_pick_latest_marker_comment(self, closer_module): - # When multiple grace warnings exist (e.g. a re-open cycle), use - # the most recent one to compute the age. - comments = [ - self._make_marker_comment(closer_module, created_at="2026-05-15T00:00:00Z"), - self._make_marker_comment(closer_module, created_at="2026-05-16T23:00:00Z"), - ] - now = dt.datetime(2026, 5, 17, 0, 0, 0, tzinfo=dt.timezone.utc) - age = closer_module.seconds_since_last_grace_warning(comments, now=now) - # 1h = 3600s - assert age == 3600.0 - - -class TestGraceWarningCommentText: - """Pin the user-facing language in the grace warning comment so the - grace-window and `@greptileai still works after close` promises - don't get accidentally dropped in a future refactor. - """ - - def test_should_state_grace_window(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - # The user's PR explicitly said "specify in the comment" — pin - # that the grace window appears in the comment. - assert "2 hours" in body - - def test_should_mention_agent_shin_reconsider(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert "@agent-shin reconsider" in body - - def test_should_promise_greptileai_works_after_close(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_should_carry_grace_marker(self, closer_module): - # The marker is what `seconds_since_last_grace_warning` greps for - # to detect a prior warning — dropping it would silently break - # the cooldown. - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert closer_module.GRACE_COMMENT_MARKER in body - - def test_close_comment_should_mention_greptileai_post_close(self, closer_module): - # The close comment should ALSO point at the @greptileai post-close - # re-review path so contributors see the same options whether they - # read the warning or only catch the close comment. - body = closer_module.format_close_comment(score=2, threshold=4) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_close_comment_should_advertise_reconsider(self, closer_module): - body = closer_module.format_close_comment(score=2, threshold=4) - assert "@agent-shin reconsider" in body - - def test_close_comment_should_carry_agent_shin_close_marker(self, closer_module): - # The close comment advertises `@agent-shin reconsider`, and the - # reconsider reopen guard (`was_closed_by_agent_shin`) only treats a - # PR as Agent-Shin-closed when the close comment carries this marker. - # Dropping it silently breaks the advertised recovery path for every - # PR closed by this daily sweep. - body = closer_module.format_close_comment(score=2, threshold=4) - assert closer_module.AGENT_SHIN_CLOSE_MARKER in body - - def test_close_comment_should_state_score_and_threshold(self, closer_module): - body = closer_module.format_close_comment(score=1, threshold=4) - assert "1/5" in body - assert "4/5" in body - - -class TestHasOptoutLabel: - def test_should_match_label_case_insensitively(self, closer_module): - pr = {"labels": [{"name": "Do Not Close"}, {"name": "bug"}]} - assert closer_module.has_optout_label(pr, {"do not close"}) is True - - def test_should_return_false_when_no_match(self, closer_module): - pr = {"labels": [{"name": "bug"}, {"name": "enhancement"}]} - assert closer_module.has_optout_label(pr, {"wip", "keep open"}) is False - - def test_should_handle_missing_labels(self, closer_module): - assert closer_module.has_optout_label({}, {"wip"}) is False - - -class TestListOpenItemsNoCap: - """The bulk sweeps must fetch the ENTIRE open backlog. - - Regression guard for the old hard-coded ``--limit 1000``: gh lists - newest-first, so a low cap silently dropped the *oldest* PRs/issues — - exactly the stale ones a low-quality sweep exists to catch. - """ - - @staticmethod - def _shared(closer_module): - # `closer_module` loading puts `.github/scripts` on sys.path and - # imports agent_shin_shared, so it's already in sys.modules. - import agent_shin_shared - - return agent_shin_shared - - def _capture_gh_args(self, closer_module, monkeypatch, *, returns="[]"): - shared = self._shared(closer_module) - captured: dict = {} - - def fake_gh(*args): - captured["args"] = args - return returns - - # `list_open_items` looks up `gh` in agent_shin_shared's namespace. - monkeypatch.setattr(shared, "gh", fake_gh) - return shared, captured - - def test_list_open_items_passes_no_cap_limit_not_1000( - self, closer_module, monkeypatch - ): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("pr", repo="o/r", fields="number,title") - args = captured["args"] - assert "--limit" in args - limit_value = args[args.index("--limit") + 1] - assert limit_value == str(shared.GH_LIST_ALL_LIMIT) - assert limit_value != "1000" - # A meaningful ceiling: comfortably above any realistic open backlog. - assert shared.GH_LIST_ALL_LIMIT >= 100_000 - - def test_list_open_items_uses_dedicated_command_state_and_fields( - self, closer_module, monkeypatch - ): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("issue", repo="o/r", fields="number") - args = captured["args"] - assert args[0] == "issue" and args[1] == "list" - assert args[args.index("--state") + 1] == "open" - assert args[args.index("--json") + 1] == "number" - assert tuple(args[-2:]) == ("--repo", "o/r") - - def test_list_open_items_omits_repo_when_none(self, closer_module, monkeypatch): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("pr", repo=None, fields="number") - assert "--repo" not in captured["args"] - - def test_list_open_items_parses_json_array(self, closer_module, monkeypatch): - shared, _ = self._capture_gh_args( - closer_module, monkeypatch, returns='[{"number": 1}, {"number": 2}]' - ) - items = shared.list_open_items("pr", repo=None, fields="number") - assert [i["number"] for i in items] == [1, 2] - - def test_list_open_items_rejects_unknown_kind(self, closer_module): - shared = self._shared(closer_module) - with pytest.raises(ValueError, match="kind must be 'pr' or 'issue', got 'both"): - shared.list_open_items("both", repo="o/r", fields="number") - - def test_fetch_open_prs_delegates_with_no_cap(self, closer_module, monkeypatch): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - closer_module.fetch_open_prs("o/r") - args = captured["args"] - assert args[0] == "pr" - assert args[args.index("--limit") + 1] == str(shared.GH_LIST_ALL_LIMIT) - # Still requests every field downstream evaluate_pr / labels logic needs. - assert "createdAt" in args[args.index("--json") + 1] - - -class TestEvaluatePrAllowlist: - """While the dogfood allowlist is active `evaluate_pr` only acts on the - named accounts and bypasses the external-only restriction for them. - Emptying it restores the internal-author skip.""" - - @pytest.fixture(autouse=True) - def _now(self): - return dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - - def _make_pr(self, *, author_login: str, created_days_ago: int = 10) -> dict: - created = dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - dt.timedelta( - days=created_days_ago - ) - return { - "number": 1, - "title": "PR #1", - "createdAt": created.isoformat().replace("+00:00", "Z"), - "isDraft": False, - "labels": [], - "author": {"login": author_login}, - "url": "https://example.com/pr/1", - } - - def test_should_skip_author_not_on_allowlist( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("must not fetch comments for non-allowlisted"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="random-oss-dev"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-not-allowlisted" - assert score is None - - def test_should_act_on_allowlisted_internal_author( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 2/5")], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="mateo-berri", created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 2 - - def test_empty_allowlist_restores_internal_skip( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("must not fetch comments for internal"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="krrishdholakia"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - allowlist=frozenset(), - ) - assert action == "skip-internal" - - def test_allowlist_constant_is_the_two_dogfood_accounts(self, closer_module): - assert closer_module.ALLOWLIST_LOGINS == frozenset( - {"mateo-berri", "swiftwinds"} - ) - - -class TestDryRunGateOnClose: - """Regression: the daily sweep is dry-run unless `--close` is passed - (the workflow only adds it when `AGENT_SHIN_ENABLED=true`). A closeable - PR (low score, grace window elapsed) must be DETECTED and reported as - "would close", but the dry run must never make a real GitHub mutation, - so merging Agent Shin stays inert by default.""" - - def _closeable_pr(self) -> dict: - return { - "number": 7, - "title": "thin PR", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": False, - "labels": [], - "author": {"login": "SwiftWinds"}, - "url": "https://example.com/pr/7", - } - - def test_dry_run_sweep_detects_but_does_not_close( - self, closer_module, monkeypatch, capsys - ): - aged_out_warning = { - "user": {"login": "github-actions[bot]"}, - "body": "warned\n\n" + closer_module.GRACE_COMMENT_MARKER, - # Far enough in the past that it's aged out regardless of - # GRACE_PERIOD_SECONDS, since main() pins `now` to real time. - "created_at": "2020-01-01T00:00:00Z", - } - monkeypatch.setattr( - closer_module, "fetch_open_prs", lambda repo: [self._closeable_pr()] - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment("Confidence Score: 1/5"), - aged_out_warning, - ], - ) - # Any real GitHub mutation during a dry run is the bug under test. - monkeypatch.setattr( - closer_module, - "gh", - lambda *a, **kw: pytest.fail(f"dry run must not call gh: {a}"), - ) - monkeypatch.setattr(sys, "argv", ["close_low_quality_prs.py"]) - - rc = closer_module.main() - - assert rc == 0 - # The PR is detected as closeable, just not acted on. - assert "Total would close: 1" in capsys.readouterr().out diff --git a/tests/test_litellm/test_github_review_gate.py b/tests/test_litellm/test_github_review_gate.py deleted file mode 100644 index 001fa8f43f5..00000000000 --- a/tests/test_litellm/test_github_review_gate.py +++ /dev/null @@ -1,524 +0,0 @@ -"""Unit tests for the `ready for review` label lifecycle (Agent Shin review gate). - -Exercises `triage_with_llm.review_gate`, the state machine that keeps the -`ready for review` label in sync with whether a PR clears both the LLM rubric -and Greptile's confidence score: - - * pass (untagged) -> add label + "ready for review" comment - * pass (untagged, recovered) -> add label + "all clear again" comment - * pass (already tagged) -> noop - * regress (tagged) -> remove label + "what's missing" comment, stays open - * fail (untagged, within 24h)-> one-time "what's missing" notice - * fail (untagged, >24h) -> close + comment - * dry run (close=False) -> would-* previews, no side effects -""" - -from __future__ import annotations - -import datetime as dt -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] / ".github" / "scripts" / "triage_with_llm.py" -) - -NOW = dt.datetime(2026, 5, 24, 12, 0, 0, tzinfo=dt.timezone.utc) -JUST_NOW = "2026-05-24T11:00:00Z" # 1h old -> within 24h grace -TWO_DAYS_AGO = "2026-05-22T11:00:00Z" # >24h old -> past grace - - -@pytest.fixture(scope="module") -def triage_module(): - spec = importlib.util.spec_from_file_location("triage_with_llm", SCRIPT_PATH) - assert spec and spec.loader - module = importlib.util.module_from_spec(spec) - sys.modules["triage_with_llm"] = module - spec.loader.exec_module(module) - return module - - -class _Recorder: - """Captures every gh mutation review_gate could fire, and fails loudly - on the ones a given scenario forbids.""" - - def __init__(self, triage_module, monkeypatch): - self.comments: list[str] = [] - self.added: list[str] = [] - self.removed: list[str] = [] - self.closed: list[int] = [] - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: self.comments.append(body), - ) - monkeypatch.setattr( - triage_module, - "add_label", - lambda repo, n, label: self.added.append(label), - ) - monkeypatch.setattr( - triage_module, - "remove_label", - lambda repo, n, label: self.removed.append(label), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda repo, n: self.closed.append(n), - ) - - -def _make_pr(**overrides): - base = { - "number": 7, - "title": "feat: do a thing", - "body": "some body without a linked issue or QA proof", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - "labels": [], - "created_at": JUST_NOW, - } - base.update(overrides) - return base - - -def _pass(prompt): - return '{"verdict": "pass", "missing": [], "explanation": "looks good"}' - - -def _fail(prompt): - return ( - '{"verdict": "fail", "missing": ["QA proof", "expected vs. actual"],' - ' "explanation": "thin description"}' - ) - - -def _gate(triage_module, **kwargs): - """Call review_gate with safe defaults for the injectable hooks.""" - params = dict( - repo="o/r", - number=7, - close=True, - model="m", - judge=_pass, - greptile_score=None, - comments=[], - now=NOW, - ) - params.update(kwargs) - return triage_module.review_gate(**params) - - -class TestReviewGatePass: - def test_pass_untagged_adds_label_and_ready_comment( - self, triage_module, monkeypatch - ): - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=5) - - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - assert rec.removed == [] and rec.closed == [] - assert len(rec.comments) == 1 - assert "ready for review" in rec.comments[0].lower() - assert triage_module.READY_MARKER in rec.comments[0] - assert "5/5" in rec.comments[0] - - def test_pass_already_tagged_is_noop(self, triage_module, monkeypatch): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=5) - - assert result["action"] == "noop-passing" - assert rec.added == [] and rec.removed == [] and rec.comments == [] - - def test_pass_after_prior_regression_uses_all_clear_wording( - self, triage_module, monkeypatch - ): - # A regression marker in history -> this is a recovery, not a first pass. - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - } - ] - - result = _gate(triage_module, judge=_pass, greptile_score=5, comments=prior) - - assert result["action"] == "labeled-ready" - assert "all clear" in rec.comments[0].lower() - - def test_linked_issue_passes_without_calling_judge( - self, triage_module, monkeypatch - ): - pr = _make_pr(body="Fixes #4321\n\nbody") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate( - triage_module, - judge=lambda p: pytest.fail("LLM must not be called for linked issue"), - greptile_score=5, - ) - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - - -class TestReviewGateRegression: - def test_regression_removes_label_and_keeps_pr_open( - self, triage_module, monkeypatch - ): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=5) - - assert result["action"] == "label-removed-regressed" - assert rec.removed == [triage_module.READY_FOR_REVIEW_LABEL] - assert rec.closed == [] # regression NEVER closes the PR - assert triage_module.REGRESSED_MARKER in rec.comments[0] - assert "QA proof" in rec.comments[0] - # The state machine closes a still-failing PR `grace_days` after this - # notice (default 24h); the comment must disclose that deadline rather - # than implying the PR stays open indefinitely. - assert "24 hours" in rec.comments[0] - assert "auto-closed" in rec.comments[0] - - def test_regression_comment_discloses_grace_deadline(self, triage_module): - one_day = triage_module.format_regression_comment( - ["QA proof"], "needs work", grace_days=1 - ) - assert "24 hours" in one_day - assert "auto-closed" in one_day - - three_days = triage_module.format_regression_comment( - ["QA proof"], "needs work", grace_days=3 - ) - assert "3 days" in three_days - assert "auto-closed" in three_days - - def test_greptile_drop_alone_triggers_regression(self, triage_module, monkeypatch): - # Rubric still passes, but Greptile fell to 2/5 -> not passing. - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=2) - - assert result["action"] == "label-removed-regressed" - assert rec.removed == [triage_module.READY_FOR_REVIEW_LABEL] - assert "2/5" in rec.comments[0] - - def test_greptile_score_read_from_comments_when_not_injected( - self, triage_module, monkeypatch - ): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - greptile = [ - { - "user": {"login": "greptile-apps[bot]"}, - "body": "Confidence Score: 2/5", - "created_at": "2026-05-24T10:00:00Z", - } - ] - - result = _gate( - triage_module, - judge=_pass, - greptile_score=triage_module._UNSET, - comments=greptile, - ) - assert result["action"] == "label-removed-regressed" - assert "2/5" in rec.comments[0] - - -class TestReviewGateGraceAndClose: - def test_within_grace_posts_one_time_notice(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, "fetch_pr", lambda repo, n: _make_pr(created_at=JUST_NOW) - ) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=None) - - assert result["action"] == "within-grace-notified" - assert rec.closed == [] and rec.added == [] and rec.removed == [] - assert triage_module.WITHIN_GRACE_MARKER in rec.comments[0] - assert "QA proof" in rec.comments[0] - - def test_within_grace_does_not_double_notify(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, "fetch_pr", lambda repo, n: _make_pr(created_at=JUST_NOW) - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.WITHIN_GRACE_MARKER, - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "within-grace-already-notified" - assert rec.comments == [] - - def test_past_grace_closes_with_comment(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=None) - - assert result["action"] == "closed" - assert rec.closed == [7] - assert len(rec.comments) == 1 - # The close comment must carry the reconsider provenance marker so - # `was_closed_by_agent_shin` can later recognize this as an Agent Shin - # close (and not some other workflow's `github-actions[bot]` close). - assert triage_module.AGENT_SHIN_CLOSE_MARKER in rec.comments[0] - - def test_recent_regression_marker_blocks_close(self, triage_module, monkeypatch): - """A failing PR with a fresh regression notice must NOT be closed — - the contributor needs a window to address the regression.""" - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - # Posted just an hour before NOW -> well inside grace_days. - "created_at": "2026-05-24T11:00:00Z", - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "regressed-already-notified" - assert rec.closed == [] and rec.comments == [] - - def test_stale_regression_marker_allows_close(self, triage_module, monkeypatch): - """Once grace_days have elapsed since the regression notice, the - review gate must let the close path fire — otherwise PRs that were - regressed and then abandoned stay open forever.""" - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - # Posted 30 days before NOW -> well past the default 1-day grace. - "created_at": "2026-04-24T11:00:00Z", - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "closed" - assert rec.closed == [7] - assert len(rec.comments) == 1 - - def test_linked_issue_with_greptile_fail_uses_greptile_explanation( - self, triage_module, monkeypatch - ): - """When the rubric short-circuits to pass (linked-issue regex) but - Greptile dragged the PR under the bar, the close comment's - explanation must describe the Greptile shortfall, not the - misleading "LLM was not called" rubric placeholder.""" - pr = _make_pr(body="Fixes #4321\n\nbody", created_at=TWO_DAYS_AGO) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate( - triage_module, - judge=lambda p: pytest.fail("LLM must not be called for linked issue"), - greptile_score=2, - ) - - assert result["action"] == "closed" - assert len(rec.comments) == 1 - body = rec.comments[0] - assert "LLM was not called" not in body - assert "Greptile" in body and "2/5" in body - - -class TestReviewGateDryRun: - @pytest.mark.parametrize( - "scenario,labels,judge,score,created,expected", - [ - ("pass", [], _pass, 5, JUST_NOW, "would-label-ready"), - ( - "regress", - [{"name": "ready for review"}], - _fail, - 5, - JUST_NOW, - "would-remove-label", - ), - ("within-grace", [], _fail, None, JUST_NOW, "would-notify-within-grace"), - ("past-grace", [], _fail, None, TWO_DAYS_AGO, "would-close"), - ], - ) - def test_dry_run_previews_without_side_effects( - self, - triage_module, - monkeypatch, - scenario, - labels, - judge, - score, - created, - expected, - ): - pr = _make_pr(labels=labels, created_at=created) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, close=False, judge=judge, greptile_score=score) - - assert result["action"] == expected - # Dry run touches nothing. - assert rec.added == [] and rec.removed == [] and rec.closed == [] - assert rec.comments == [] - assert "comment" in result # preview body still surfaced - - -class TestReviewGateGuards: - def test_skips_internal_author(self, triage_module, monkeypatch): - pr = _make_pr(author_association="MEMBER", user={"login": "krrish"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate( - triage_module, - judge=lambda p: pytest.fail("no LLM for internal"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_skips_closed_pr(self, triage_module, monkeypatch): - pr = _make_pr(state="closed") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate(triage_module, judge=lambda p: pytest.fail("no LLM for closed")) - assert result["action"] == "skip-not-open" - - def test_llm_error_is_non_destructive(self, triage_module, monkeypatch): - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - - def boom(prompt): - raise RuntimeError("api down") - - result = _gate(triage_module, judge=boom, greptile_score=None) - - assert result["action"] == "skip-llm-error" - assert rec.closed == [] and rec.added == [] and rec.removed == [] - - def test_full_recovery_cycle(self, triage_module, monkeypatch): - """pass -> regress -> recover, threading labels/comments like GitHub would.""" - state = {"labels": [], "comments": []} - - def fake_fetch(repo, n): - return _make_pr(labels=list(state["labels"]), created_at=JUST_NOW) - - monkeypatch.setattr(triage_module, "fetch_pr", fake_fetch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: state["comments"].append( - {"user": {"login": "github-actions[bot]"}, "body": body} - ), - ) - monkeypatch.setattr( - triage_module, - "add_label", - lambda repo, n, label: state["labels"].append({"name": label}), - ) - monkeypatch.setattr( - triage_module, - "remove_label", - lambda repo, n, label: state["labels"].clear(), - ) - monkeypatch.setattr( - triage_module, "close_pr", lambda repo, n: pytest.fail("must not close") - ) - - # 1) passes -> tagged - r1 = _gate( - triage_module, judge=_pass, greptile_score=5, comments=state["comments"] - ) - assert r1["action"] == "labeled-ready" - assert any(lbl["name"] == "ready for review" for lbl in state["labels"]) - - # 2) regresses -> tag removed, comment posted, PR still open - r2 = _gate( - triage_module, judge=_fail, greptile_score=2, comments=state["comments"] - ) - assert r2["action"] == "label-removed-regressed" - assert state["labels"] == [] - - # 3) fixed again -> "all clear" + tag back - r3 = _gate( - triage_module, judge=_pass, greptile_score=5, comments=state["comments"] - ) - assert r3["action"] == "labeled-ready" - assert any(lbl["name"] == "ready for review" for lbl in state["labels"]) - assert "all clear" in state["comments"][-1]["body"].lower() - - -class TestReviewGateAllowlist: - """While the dogfood allowlist is active it is the sole author gate: - only the named accounts pass, and for them the internal-author exemption - is bypassed. Emptying it restores the normal internal-author skip.""" - - def test_should_skip_author_not_on_allowlist(self, triage_module, monkeypatch): - pr = _make_pr(user={"login": "random-oss-dev"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - result = _gate( - triage_module, judge=lambda p: pytest.fail("no LLM for non-allowlisted") - ) - assert result["action"] == "skip-not-allowlisted" - assert rec.added == [] and rec.comments == [] and rec.closed == [] - - def test_should_act_on_allowlisted_internal_author( - self, triage_module, monkeypatch - ): - pr = _make_pr(author_association="MEMBER", user={"login": "mateo-berri"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - result = _gate(triage_module, judge=_pass, greptile_score=5) - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - - def test_empty_allowlist_restores_internal_skip(self, triage_module, monkeypatch): - pr = _make_pr(author_association="MEMBER", user={"login": "krrish"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate( - triage_module, - judge=lambda p: pytest.fail("no LLM for internal"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py deleted file mode 100644 index ddffb978b48..00000000000 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ /dev/null @@ -1,2134 +0,0 @@ -"""Unit tests for `.github/scripts/triage_with_llm.py` (Agent Shin).""" - -from __future__ import annotations - -import importlib.util -import json -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] / ".github" / "scripts" / "triage_with_llm.py" -) - - -@pytest.fixture(scope="module") -def triage_module(): - spec = importlib.util.spec_from_file_location("triage_with_llm", SCRIPT_PATH) - assert spec and spec.loader - module = importlib.util.module_from_spec(spec) - sys.modules["triage_with_llm"] = module - spec.loader.exec_module(module) - return module - - -class TestIsInternalContributor: - @pytest.mark.parametrize("association", ["OWNER", "MEMBER", "COLLABORATOR"]) - def test_should_mark_org_associations_as_internal(self, triage_module, association): - item = { - "author_association": association, - "user": {"login": "krrishdholakia"}, - } - assert triage_module.is_internal_contributor(item) is True - - @pytest.mark.parametrize( - "association", - ["CONTRIBUTOR", "FIRST_TIME_CONTRIBUTOR", "FIRST_TIMER", "NONE"], - ) - def test_should_mark_outside_associations_as_external( - self, triage_module, association - ): - item = { - "author_association": association, - "user": {"login": "random-oss-dev"}, - } - assert triage_module.is_internal_contributor(item) is False - - @pytest.mark.parametrize( - "item", - [ - {"author_association": "", "user": {"login": "random-oss-dev"}}, - {"user": {"login": "random-oss-dev"}}, # association field absent - ], - ) - def test_should_fail_safe_when_author_association_is_missing( - self, triage_module, item - ): - # Fail-safe: an empty/missing association must never make a PR - # eligible for the destructive close path. Treat as internal (skip). - assert triage_module.is_internal_contributor(item) is True - - @pytest.mark.parametrize( - "login", - ["dependabot[bot]", "greptile-apps[bot]", "dependabot", "github-actions"], - ) - def test_should_skip_bot_accounts_regardless_of_association( - self, triage_module, login - ): - item = {"author_association": "NONE", "user": {"login": login}} - assert triage_module.is_internal_contributor(item) is True - - -class TestHasLinkedIssue: - @pytest.mark.parametrize( - "body", - [ - "Fixes #1234", - "closes #1", - "Resolves #99", - "fix #42 — this addresses the regression", - "Closes https://github.com/BerriAI/litellm/issues/27000", - "Resolved https://github.com/BerriAI/litellm/issues/27001", - ], - ) - def test_should_detect_common_link_phrases(self, triage_module, body): - assert triage_module.has_linked_issue(body) is True - - @pytest.mark.parametrize( - "body", - [ - "", - "Some change", - # Casual mentions must NOT auto-pass — they should fall through to - # the LLM judge so the stricter "not a passing mention" rule fires. - "See #1234", - "see #1234 for context", - "ref #1234", - "Refs https://github.com/BerriAI/litellm/issues/27000", - "this addresses #1234", - ], - ) - def test_should_not_auto_pass_casual_mentions(self, triage_module, body): - assert triage_module.has_linked_issue(body) is False - - def test_should_not_detect_when_only_html_comment_template(self, triage_module): - body = "" - assert triage_module.has_linked_issue(body) is False - - -class TestStripHtmlComments: - def test_should_remove_single_line_comments(self, triage_module): - text = "before after" - assert "placeholder" not in triage_module.strip_html_comments(text) - - def test_should_remove_multiline_comments(self, triage_module): - text = "kept\n\nkept2" - cleaned = triage_module.strip_html_comments(text) - assert "Fixes #1" not in cleaned - assert "kept" in cleaned and "kept2" in cleaned - - def test_should_handle_none(self, triage_module): - assert triage_module.strip_html_comments(None) == "" - - -class TestCloseCommentText: - """Pin the user-facing language in close comments so changes are intentional.""" - - def test_pr_close_comment_should_recommend_new_pr_primarily(self, triage_module): - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"} - ) - # Primary path: open a new PR (because OSS authors can't reopen a - # bot-closed PR). Secondary path: `@agent-shin reconsider`. - assert "Open a new PR" in body - assert "@agent-shin reconsider" in body - # Old advice that no longer works for OSS contributors must NOT - # appear (they can't reopen a PR closed by a bot/maintainer). - assert "Reopen the PR" not in body - - def test_reopen_comment_should_carry_reconsider_marker(self, triage_module): - # The marker is what the rate-limit guard greps for to detect a - # prior reconsider verdict on the same PR. If the marker ever - # gets dropped from this comment, the cooldown silently breaks - # and a contributor can spam `@agent-shin reconsider` to burn - # LLM budget. - body = triage_module.format_reopen_comment("pr") - assert triage_module.RECONSIDER_COMMENT_MARKER in body - - def test_still_failing_comment_should_carry_reconsider_marker(self, triage_module): - body = triage_module.format_reconsider_still_failing_comment( - "pr", - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"}, - ) - assert triage_module.RECONSIDER_COMMENT_MARKER in body - - def test_pr_close_comment_should_not_promise_automatic_reopen_on_open( - self, triage_module - ): - # The previous comment said "I'll re-evaluate automatically" — that - # only worked because the author could reopen, which they often - # can't. The new wording must point them at the comment trigger or - # a new PR instead. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "I'll re-evaluate automatically" not in body - - def test_issue_close_comment_should_use_reconsider_trigger(self, triage_module): - # OSS authors have read access, which only lets them reopen issues - # they closed themselves; they CANNOT reopen an issue a maintainer or - # bot closed. So the recovery path is `@agent-shin reconsider` (the - # bot reopens), exactly like the PR path. If this regresses to "reopen - # it yourself", contributors hit a dead end on bot-closed issues. - body = triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": ["repro"], "explanation": "thin"} - ) - assert "@agent-shin reconsider" in body - - def test_pr_close_comment_should_link_blog_explainer(self, triage_module): - # The blog post is the canonical public explanation of what the bot - # checks and why. Every action-required bot comment must link to it - # so contributors landing on a bot-closed PR can self-serve context - # without pinging a maintainer. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "https://docs.litellm.ai/blog/agent-shin-triage" in body - - def test_issue_close_comment_should_link_blog_explainer(self, triage_module): - body = triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "https://docs.litellm.ai/blog/agent-shin-triage" in body - - def test_pr_close_comment_should_flag_mocked_tests_as_insufficient_proof( - self, triage_module - ): - # The PR rubric was tightened to require end-to-end QA proof and - # explicitly exclude mocked-dependency unit tests. The user-facing - # close comment must say so — otherwise contributors will keep - # re-submitting "pytest passed (mocks)" runs and getting closed - # again with no explanation of why. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "end-to-end qa proof" in body.lower() - assert "mock" in body.lower() - - def test_issue_recovery_comments_should_name_feature_dead_end_evidence( - self, triage_module - ): - # The feature-request pass bar demands end-to-end evidence of the - # dead-end, so the close and grace-warning recovery bullets must ask - # for it too — otherwise a requester follows those exact instructions - # (description + use case only) and fails `reconsider` again with no - # hint of what else was needed. - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - for body in ( - triage_module.format_issue_close_comment(verdict), - triage_module.format_grace_warning_issue_comment(verdict), - ): - normalized = " ".join(body.split()) - assert "end-to-end evidence of the dead-end" in normalized - assert "showing where the flow stops today" in normalized - - def test_all_agent_shin_comments_should_use_bullet_train_emoji(self, triage_module): - # The bullet train (🚅) is Agent Shin's symbol, matching the LiteLLM - # logo; the previous wave (👋) was generic and didn't match the bot's - # identity. Every action-required comment the bot can post must use the - # bullet train so the contributor recognizes who's writing without - # reading the signoff. - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - comments = { - "pr_close": triage_module.format_pr_close_comment(verdict), - "issue_close": triage_module.format_issue_close_comment(verdict), - "pr_grace": triage_module.format_grace_warning_pr_comment(verdict), - "issue_grace": triage_module.format_grace_warning_issue_comment(verdict), - "within_grace": triage_module.format_within_grace_comment( - [], "", grace_days=1 - ), - } - for name, body in comments.items(): - assert "🚅" in body, f"{name} comment is missing the bullet train emoji" - assert "👋" not in body, f"{name} comment still uses the old wave emoji" - - def test_pr_close_comment_should_show_what_pr_got_right(self, triage_module): - # The user explicitly asked for a "things you got right" section so - # the comment doesn't read as pure rejection. When the judge confirms - # a field is present (e.g. linked_issue), the bullet for it MUST - # appear in the close comment. - body = triage_module.format_pr_close_comment( - { - "verdict": "fail", - "linked_issue": True, - "has_problem_description": True, - "has_expected_vs_actual": False, - "has_qa_proof": False, - "missing": ["QA proof"], - "explanation": "no proof", - } - ) - assert "What you got right" in body - # The two present fields surface as ✅ bullets; the two absent - # fields do not get a ✅ bullet (the QA-proof rubric block still - # mentions the concept, but only the affirmed fields get checkmarks). - assert "- ✅ Linked a related GitHub issue" in body - assert "- ✅ Clear problem description" in body - assert "- ✅ Expected vs. actual behavior" not in body - assert "- ✅ End-to-end QA proof" not in body - - def test_pr_close_comment_should_omit_present_section_when_nothing_present( - self, triage_module - ): - # If the judge says nothing is present (every flag False), the - # "what you got right" block is skipped entirely — better to omit - # than to render "What you got right: (nothing)". - body = triage_module.format_pr_close_comment( - { - "verdict": "fail", - "linked_issue": False, - "has_problem_description": False, - "has_expected_vs_actual": False, - "has_qa_proof": False, - "missing": [], - "explanation": "", - } - ) - assert "What you got right" not in body - - def test_issue_close_comment_should_show_what_issue_got_right(self, triage_module): - # `has_expected_vs_actual` is present, the end-to-end bug evidence is - # not: the "what you got right" block must surface the former and omit - # the latter (no "✅ (nothing)"-style noise for absent items). - body = triage_module.format_issue_close_comment( - { - "verdict": "fail", - "kind": "bug", - "has_repro": False, - "has_expected_vs_actual": True, - "missing": ["end-to-end evidence of the bug"], - "explanation": "no repro shown", - } - ) - assert "What you got right" in body - assert "Expected vs. actual behavior" in body - assert "- ✅ End-to-end evidence of the bug" not in body - - def test_issue_close_comment_should_credit_feature_dead_end_evidence( - self, triage_module - ): - # A feature requester who pasted their dead-end run but skipped the - # motivation must see the evidence credited and only the motivation - # listed as a gap — without a dedicated verdict field the praise - # block could never acknowledge the work they did do. - body = triage_module.format_issue_close_comment( - { - "verdict": "fail", - "kind": "feature", - "has_motivation_example": False, - "has_dead_end_evidence": True, - "missing": ["motivation / use case"], - "explanation": "no use case given", - } - ) - assert "What you got right" in body - assert "- ✅ End-to-end evidence of the dead-end" in body - assert "- ✅ Motivation and concrete example" not in body - - def test_close_comments_should_use_softer_park_for_later_framing( - self, triage_module - ): - # User feedback: the messaging shouldn't feel like punishment. The - # comment must explicitly frame close as a "park this for later," not - # a rejection, and ground that in the queue-hygiene reason. - for body in ( - triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ): - assert "park this for later" in body - assert ( - "not a rejection" in body - or "isn't a rejection" in body - or ("isn't us saying" in body) - ) - - def test_only_close_comments_carry_the_agent_shin_close_marker(self, triage_module): - # The reconsider reopen guard keys off AGENT_SHIN_CLOSE_MARKER to tell - # an Agent Shin close from a same-identity close by another workflow. - # That only works if the marker is stamped on the close comments and - # NOT on the grace warnings (which don't close anything). - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - marker = triage_module.AGENT_SHIN_CLOSE_MARKER - assert marker in triage_module.format_pr_close_comment(verdict) - assert marker in triage_module.format_issue_close_comment(verdict) - assert marker not in triage_module.format_grace_warning_pr_comment(verdict) - assert marker not in triage_module.format_grace_warning_issue_comment(verdict) - - -class TestWasClosedByAgentShin: - """Bot-closed guard: only Agent Shin's own closures are reopen candidates.""" - - @staticmethod - def _stub_close_event( - triage_module, - monkeypatch, - *, - actor: str | None, - closed_at: object = "now", - ): - """Stub the most recent `closed` event used by the guard. - - `actor` is the login that closed the item. `closed_at` defaults - to "now" so the marker comment (stubbed at 42s ago) reads as - recent enough relative to the close; tests can pass a concrete - ``datetime`` to simulate older closes (e.g. the stale-marker - regression scenario). - """ - import datetime as real_dt - - if closed_at == "now": - closed_at = real_dt.datetime.now(real_dt.timezone.utc) - monkeypatch.setattr( - triage_module, - "fetch_last_close_event", - lambda repo, n: (actor, closed_at), - ) - - @staticmethod - def _stub_close_marker_present( - triage_module, monkeypatch, *, present: bool, age_seconds: float = 42.0 - ): - """Stub the Agent Shin close-comment marker lookup. - - `was_closed_by_agent_shin` requires the closing actor AND a - recent Agent Shin close comment; these tests pin the latter so - they exercise the actor half in isolation. - """ - monkeypatch.setattr( - triage_module, - "seconds_since_last_agent_shin_close", - lambda *a, **kw: age_seconds if present else None, - ) - - def test_should_return_true_when_bot_closed_and_close_comment_present( - self, triage_module, monkeypatch - ): - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is True - - def test_should_return_false_when_bot_closed_but_no_agent_shin_comment( - self, triage_module, monkeypatch - ): - # The `github-actions[bot]` identity is shared across workflows. A - # stale/duplicate sweep closing under that identity must NOT let - # @agent-shin reconsider reopen the item: without an Agent Shin close - # comment the guard fails closed. - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - self._stub_close_marker_present(triage_module, monkeypatch, present=False) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_return_false_when_last_close_actor_is_maintainer( - self, triage_module, monkeypatch - ): - # A maintainer closed it (e.g. duplicate, security, design). The - # bot must refuse to reopen on @agent-shin reconsider even if an - # earlier Agent Shin close comment is still on the thread. - self._stub_close_event(triage_module, monkeypatch, actor="krrishdholakia") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_fail_closed_when_no_close_event(self, triage_module, monkeypatch): - # If the events API returns nothing (network blip, repo permission - # quirk), the guard must fail-closed: refuse to reopen rather than - # assume the bot did it. - self._stub_close_event(triage_module, monkeypatch, actor=None, closed_at=None) - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_fail_closed_when_close_event_has_no_timestamp( - self, triage_module, monkeypatch - ): - # Without a usable close timestamp the guard cannot prove the - # marker comment belongs to the latest close; fail-closed. - self._stub_close_event( - triage_module, monkeypatch, actor="github-actions[bot]", closed_at=None - ) - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_return_false_when_marker_predates_latest_close( - self, triage_module, monkeypatch - ): - # Regression for the stale-marker bug: Agent Shin closed once - # (marker stamped), reconsider reopened, and a different workflow - # later closed under the same bot identity without stamping the - # marker. The old marker is still on the thread but does NOT - # belong to the latest close, so reconsider must not reopen. - import datetime as real_dt - - now = real_dt.datetime.now(real_dt.timezone.utc) - # Latest close happened a minute ago. - self._stub_close_event( - triage_module, - monkeypatch, - actor="github-actions[bot]", - closed_at=now - real_dt.timedelta(seconds=60), - ) - # The most recent Agent Shin marker is from an hour ago (a prior - # closed/reopened cycle), which is well outside the skew window. - self._stub_close_marker_present( - triage_module, monkeypatch, present=True, age_seconds=3600.0 - ) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_respect_bot_login_override_via_env( - self, triage_module, monkeypatch - ): - # Operators wiring Agent Shin to a PAT (instead of GITHUB_TOKEN) - # can override the expected bot login via env. The guard must - # respect the override so non-default deployments still work. - monkeypatch.setenv("AGENT_SHIN_BOT_LOGIN", "my-bot") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - self._stub_close_event(triage_module, monkeypatch, actor="my-bot") - assert triage_module.was_closed_by_agent_shin("o/r", 1) is True - # Default "github-actions[bot]" should NOT match when env is set. - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - -class TestSecondsSinceLastAgentShinClose: - """Close-provenance lookup: detects the bot's own auto-close marker.""" - - def _make_comment(self, *, login: str, body: str) -> dict: - return { - "user": {"login": login}, - "body": body, - "created_at": "2026-05-18T05:00:00Z", - } - - def test_should_return_none_when_bot_never_closed(self, triage_module, monkeypatch): - # Comments exist, but none is an Agent Shin close — e.g. only a grace - # warning, or a close by another workflow with no Agent Shin comment. - comments = [ - self._make_comment(login="outside-dev", body="any update?"), - self._make_comment( - login="github-actions[bot]", - body=triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is None - - def test_should_detect_bot_close_comment(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body=triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is not None - - def test_should_ignore_non_bot_comment_quoting_marker( - self, triage_module, monkeypatch - ): - # A contributor quoting the hidden marker (GitHub "Quote reply" - # preserves HTML comments) must not be mistaken for a bot close. - comments = [ - self._make_comment( - login="curious-user", - body=f"what is this? {triage_module.AGENT_SHIN_CLOSE_MARKER}", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is None - - -class TestSecondsSinceLastReconsiderVerdict: - """Rate-limit guard: detects the bot's own reconsider verdict marker.""" - - def _make_comment( - self, *, login: str, body: str, created_at: str | None = "2026-05-18T05:00:00Z" - ) -> dict: - comment: dict = {"user": {"login": login}, "body": body} - if created_at is not None: - comment["created_at"] = created_at - return comment - - def test_should_return_none_when_no_bot_reconsider_comments( - self, triage_module, monkeypatch - ): - # An issue with chatter from other users but no bot reconsider - # verdict must not be rate-limited. - comments = [ - self._make_comment(login="outside-dev", body="ping?"), - self._make_comment( - login="github-actions[bot]", body="some other bot message" - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - def test_should_pick_latest_bot_reconsider_marker(self, triage_module, monkeypatch): - # When multiple reconsider verdicts exist, return the AGE of the - # most recent one. Using a frozen reference helps pin the math. - comments = [ - self._make_comment( - login="github-actions[bot]", - body="old verdict " + triage_module.RECONSIDER_COMMENT_MARKER, - created_at="2026-05-18T04:00:00Z", - ), - self._make_comment( - login="github-actions[bot]", - body="newer verdict " + triage_module.RECONSIDER_COMMENT_MARKER, - created_at="2026-05-18T04:55:00Z", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - - # Freeze "now" via a tiny shim on the module's `dt` import. - import datetime as real_dt - - class FrozenDateTime(real_dt.datetime): - @classmethod - def now(cls, tz=None): - return real_dt.datetime(2026, 5, 18, 5, 0, 0, tzinfo=tz) - - frozen_module = type(triage_module.dt)("datetime") - frozen_module.datetime = FrozenDateTime - frozen_module.timezone = real_dt.timezone - monkeypatch.setattr(triage_module, "dt", frozen_module) - - age = triage_module.seconds_since_last_reconsider_verdict("o/r", 1) - # newer verdict is 5 minutes (300 seconds) before "now" - assert age == 300.0 - - def test_should_ignore_non_bot_comments_with_marker( - self, triage_module, monkeypatch - ): - # A user comment that happens to quote the marker (e.g. in - # a "what does this hidden marker do?" question) must NOT count. - # The rate-limit guard only trusts comments authored by the bot. - comments = [ - self._make_comment( - login="curious-user", - body=f"Saw this marker: {triage_module.RECONSIDER_COMMENT_MARKER}", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - def test_should_ignore_bot_comments_without_marker( - self, triage_module, monkeypatch - ): - # The bot posts other things too (Agent Shin close comments, - # CI status, etc.) — only the reconsider-verdict marker should - # arm the cooldown. - comments = [ - self._make_comment( - login="github-actions[bot]", - body="Agent Shin closed this PR (no marker)", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - -class TestParseVerdict: - def test_should_parse_plain_json(self, triage_module): - raw = '{"verdict": "pass", "missing": []}' - assert triage_module.parse_verdict(raw)["verdict"] == "pass" - - def test_should_strip_markdown_fence(self, triage_module): - raw = '```json\n{"verdict": "fail", "missing": ["foo"]}\n```' - result = triage_module.parse_verdict(raw) - assert result["verdict"] == "fail" - assert result["missing"] == ["foo"] - - def test_should_extract_embedded_json_from_prose(self, triage_module): - raw = 'Here you go: {"verdict": "pass", "missing": []}\nThanks.' - assert triage_module.parse_verdict(raw)["verdict"] == "pass" - - def test_should_raise_for_unparseable_text(self, triage_module): - with pytest.raises(ValueError, match='could not extract JSON from LLM response: not even close to'): - triage_module.parse_verdict("not even close to json") - - def test_should_raise_for_empty(self, triage_module): - with pytest.raises(ValueError, match='empty LLM response'): - triage_module.parse_verdict("") - - -class TestBuildPrompts: - def test_should_include_pr_title_and_body(self, triage_module): - prompt = triage_module.build_pr_prompt( - title="Add foo", body=" Real body" - ) - assert "Add foo" in prompt - assert "Real body" in prompt - assert "comment" not in prompt # HTML comments are stripped - - def test_should_show_empty_marker_for_empty_pr_body(self, triage_module): - prompt = triage_module.build_pr_prompt(title="t", body="") - assert "(empty)" in prompt - - def test_should_include_issue_title_and_body(self, triage_module): - prompt = triage_module.build_issue_prompt(title="Bug", body="repro here") - assert "Bug" in prompt - assert "repro here" in prompt - - def test_issue_bug_rubric_requires_end_to_end_evidence_and_drops_pass_bias( - self, triage_module - ): - # The bug bar was tightened: a report needs the "before" half shown - # end-to-end (video / screenshot / real command output), prose-only - # repro steps no longer pass, and the old "bias toward PASS" leniency - # is gone. If any of these regress, the judge silently goes soft on - # undemonstrated bug reports again. - prompt = triage_module.build_issue_prompt(title="t", body="x") - normalized = " ".join(prompt.split()) - assert "Bias toward PASS when the issue has structure" not in normalized - assert "END-TO-END EVIDENCE OF THE BUG" in normalized - assert "Do not bias toward PASS" in normalized - # The three accepted forms of the "before" demonstration must be named. - assert "screen recording / video" in normalized - assert "screenshot of the bug" in normalized - assert "mocked or stubbed" in normalized - # Prose-only steps are explicitly insufficient now. - assert "steps to reproduce" in normalized - # An unedited issue-form scaffold must not read as evidence: the proof - # field ships with visible headings, so the judge has to be told that - # bare headings with nothing under them count as absent. - assert "unfilled template scaffold" in normalized - assert "counts as absent, not as evidence" in normalized - - def test_issue_feature_rubric_requires_evidence_of_the_dead_end( - self, triage_module - ): - # The feature form asks the requester to walk the ideal flow against a - # live proxy and paste output up to the step that dead-ends, so the - # judge has to demand that evidence, and must not accept an unedited - # scaffold of bare headings as if it were a real attempt. - prompt = triage_module.build_issue_prompt(title="t", body="x") - normalized = " ".join(prompt.split()) - assert "END-TO-END EVIDENCE OF THE DEAD-END" in normalized - assert "showing the point where the flow stops today" in normalized - assert "unfilled template scaffold" in normalized - # The evidence has its own verdict field so feature requesters who - # provided it get credited in "What you got right", exactly like - # `has_repro` credits bug evidence. - assert "`has_dead_end_evidence=true` only when this is present" in normalized - assert '"has_dead_end_evidence": boolean' in normalized - - def test_should_not_crash_when_pr_body_contains_curly_braces(self, triage_module): - """User-supplied content with `{` / `}` must NOT be re-parsed by - `str.format()`. `format` only scans the template literal for - replacement fields; values being substituted in are inserted as - plain strings, so a body like `{"foo": "bar"}` or `{unmatched` - cannot blow up the script. Pinning this here so a future - "improvement" to the templating doesn't reintroduce a crash on - every PR that quotes JSON. - """ - for body in ( - 'Here is some JSON: {"foo": "bar", "n": 1}', - "Half a brace { left dangling, and a stray }", - "Format-spec-looking thing: {0}, {name:>10}, {!r}", - "Nested {a: {b: c}} braces", - ): - pr_prompt = triage_module.build_pr_prompt(title="t", body=body) - issue_prompt = triage_module.build_issue_prompt(title="t", body=body) - assert body in pr_prompt - assert body in issue_prompt - - def test_should_not_crash_when_pr_title_contains_curly_braces(self, triage_module): - title = "Fix bug in {0:>10} format-spec handling" - pr_prompt = triage_module.build_pr_prompt(title=title, body="x") - issue_prompt = triage_module.build_issue_prompt(title=title, body="x") - assert title in pr_prompt - assert title in issue_prompt - - def test_should_preserve_template_indentation_with_multiline_body( - self, triage_module - ): - """`textwrap.dedent` runs on the static template *before* user - content is interpolated, so a multi-line body (whose 2nd+ lines - start at column 0) cannot defeat the common-indent computation - and leave 8-space indentation on every template line. Pin the - dedented shape so the rendered prompt stays consistent for the - LLM judge. - """ - body = "first line\nsecond line at column 0\nthird line at column 0" - for builder in ( - triage_module.build_pr_prompt, - triage_module.build_issue_prompt, - ): - prompt = builder(title="t", body=body) - # Template lines should NOT carry the 8 leading spaces from - # the source-file indentation of the triple-quoted string. - assert " You are " not in prompt - assert 'You are "Agent Shin"' in prompt - assert body in prompt - - -class TestMainModelDefault: - """`--model` falls back to DEFAULT_MODEL even when TRIAGE_MODEL is empty.""" - - def _stub_triage(self, triage_module, monkeypatch): - captured: dict = {} - - def fake_triage(**kwargs): - captured.update(kwargs) - return { - "kind": kwargs["kind"], - "number": kwargs["number"], - "title": "", - "author": "x", - "author_association": "NONE", - "state": "open", - "action": "skip-no-llm-key", - } - - monkeypatch.setattr(triage_module, "triage", fake_triage) - return captured - - def test_should_fall_back_to_default_when_triage_model_env_empty( - self, triage_module, monkeypatch - ): - captured = self._stub_triage(triage_module, monkeypatch) - monkeypatch.setenv("TRIAGE_MODEL", "") - monkeypatch.setattr( - sys, - "argv", - ["triage_with_llm.py", "--repo", "o/r", "--pr", "1"], - ) - rc = triage_module.main() - assert rc == 0 - assert captured["model"] == triage_module.DEFAULT_MODEL - - def test_should_respect_explicit_triage_model_env(self, triage_module, monkeypatch): - captured = self._stub_triage(triage_module, monkeypatch) - monkeypatch.setenv("TRIAGE_MODEL", "gpt-4o-mini") - monkeypatch.setattr( - sys, - "argv", - ["triage_with_llm.py", "--repo", "o/r", "--pr", "1"], - ) - rc = triage_module.main() - assert rc == 0 - assert captured["model"] == "gpt-4o-mini" - - -class TestCallLlmJudge: - """call_llm_judge sets gpt-5 specific kwargs correctly.""" - - def _stub_openai(self, monkeypatch, captured: dict): - """Install a fake `openai.OpenAI` client into sys.modules. - - The fake client records the kwargs passed to chat.completions.create - and returns a minimal response object whose .choices[0].message.content - is "ok". - """ - import types - - class FakeMessage: - content = '{"verdict": "pass"}' - - class FakeChoice: - message = FakeMessage() - - class FakeResponse: - choices = [FakeChoice()] - - class FakeCompletions: - def create(self, **kwargs): - captured.update(kwargs) - return FakeResponse() - - class FakeChat: - completions = FakeCompletions() - - class FakeClient: - def __init__(self, api_key, base_url=None): - captured["__client_kwargs__"] = { - "api_key": api_key, - "base_url": base_url, - } - self.chat = FakeChat() - - fake_module = types.ModuleType("openai") - fake_module.OpenAI = FakeClient - monkeypatch.setitem(sys.modules, "openai", fake_module) - - def test_should_set_reasoning_effort_none_for_gpt5_family( - self, triage_module, monkeypatch - ): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model="gpt-5.4-mini", api_key="sk-test", base_url=None - ) - assert captured["model"] == "gpt-5.4-mini" - assert captured["temperature"] == 0 - assert captured["extra_body"] == {"reasoning_effort": "none"} - - def test_should_set_reasoning_effort_for_capitalized_or_dated_gpt5( - self, triage_module, monkeypatch - ): - for model in ("GPT-5.4-mini", "gpt-5.4-mini-2026-03-17", "gpt-5"): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model=model, api_key="sk-test", base_url=None - ) - assert captured["extra_body"] == {"reasoning_effort": "none"}, model - - def test_should_omit_reasoning_effort_for_non_gpt5( - self, triage_module, monkeypatch - ): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model="gpt-4o-mini", api_key="sk-test", base_url=None - ) - assert "extra_body" not in captured - - def test_should_pass_base_url_when_provided(self, triage_module, monkeypatch): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "p", - model="gpt-5.4-mini", - api_key="sk-test", - base_url="https://proxy.example.com/v1", - ) - assert ( - captured["__client_kwargs__"]["base_url"] == "https://proxy.example.com/v1" - ) - - -class TestTriageOrchestration: - """End-to-end-ish tests that mock both gh fetchers and the LLM.""" - - def _make_pr(self, **overrides): - base = { - "number": 1, - "title": "PR title", - "body": "PR body", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - base.update(overrides) - return base - - def test_should_skip_internal_author(self, triage_module, monkeypatch): - pr = self._make_pr( - author_association="MEMBER", user={"login": "krrishdholakia"} - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - - def boom(*a, **kw): - pytest.fail("LLM should not be called for internal authors") - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=boom, - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_should_skip_closed_pr(self, triage_module, monkeypatch): - pr = self._make_pr(state="closed") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("should not run on closed PRs"), - ) - assert result["action"] == "skip-not-open" - - def test_should_short_circuit_on_linked_issue(self, triage_module, monkeypatch): - pr = self._make_pr(body="Fixes #1234\n\nFoo bar") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM should not be called"), - ) - assert result["action"] == "pass-linked-issue" - assert result["verdict"]["verdict"] == "pass" - - def test_should_not_short_circuit_on_casual_mention( - self, triage_module, monkeypatch - ): - # "See #1234" is a passing mention, not a closing keyword. The LLM - # must get a chance to apply the stricter rubric. With no prior - # grace warning, the first failing verdict triggers the warning - # path (`would-warn-grace` in dry-run). - pr = self._make_pr(body="See #1234 for context. No QA proof here.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - called = {"judge": False} - - def judge(prompt): - called["judge"] = True - return json.dumps( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin."} - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=judge, - ) - assert called["judge"] is True - assert result["action"] == "would-warn-grace" - - def test_should_return_pass_llm_when_judge_passes(self, triage_module, monkeypatch): - pr = self._make_pr(body="Long body, no linked issue.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - captured = {} - - def judge(prompt): - captured["prompt"] = prompt - return json.dumps({"verdict": "pass", "missing": [], "explanation": "ok"}) - - result = triage_module.triage( - repo="o/r", kind="pr", number=1, close=True, model="m", judge=judge - ) - assert result["action"] == "pass-llm" - assert "Long body" in captured["prompt"] - - def test_should_return_would_close_in_dry_run_after_grace_aged_out( - self, triage_module, monkeypatch - ): - # When the grace warning has already aged out (>= GRACE_PERIOD_SECONDS) - # AND the rubric still fails, the dry-run preview returns - # `would-close` so a step-summary writer can render the close - # comment without touching GitHub state. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_aged_out(triage_module, monkeypatch) - - def fake_post(*a, **kw): - pytest.fail("should not post comments in dry-run") - - def fake_close(*a, **kw): - pytest.fail("should not close in dry-run") - - monkeypatch.setattr(triage_module, "post_comment", fake_post) - monkeypatch.setattr(triage_module, "close_pr", fake_close) - - verdict = { - "verdict": "fail", - "missing": ["problem description", "QA proof"], - "explanation": "Body is one sentence.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "would-close" - assert result["verdict"]["missing"] == ["problem description", "QA proof"] - - def test_should_post_comment_and_close_after_grace_window( - self, triage_module, monkeypatch - ): - # The "real close" path: --close passed AND the grace warning has - # aged out AND the rubric still fails. The bot posts the close - # comment and closes the PR. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_aged_out(triage_module, monkeypatch) - posted = {} - closed = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"repo": repo, "n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda repo, n: closed.update({"repo": repo, "n": n}), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "closed" - assert posted["n"] == 42 and closed["n"] == 42 - assert "Agent Shin" in posted["body"] - assert "QA proof" in posted["body"] - - def test_should_skip_on_llm_error_in_close_mode(self, triage_module, monkeypatch): - pr = self._make_pr(body="something.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment on LLM error"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on LLM error"), - ) - - def broken_judge(prompt): - raise RuntimeError("upstream 500") - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=broken_judge, - ) - assert result["action"] == "skip-llm-error" - assert "upstream 500" in result["error"] - - def test_should_skip_open_pr_in_reconsider_mode(self, triage_module, monkeypatch): - # Reconsider only makes sense on a CLOSED PR — running it on an open - # one is a no-op (the regular triage flow already evaluated it). - pr = self._make_pr(state="open") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: pytest.fail("should not run on open PR in reconsider"), - reconsider=True, - ) - assert result["action"] == "skip-not-closed" - - @staticmethod - def _stub_reconsider_guards(triage_module, monkeypatch): - """Default reconsider-guard stubs: pretend bot closed + no cooldown. - - The new safety guards (`was_closed_by_agent_shin`, - `seconds_since_last_reconsider_verdict`) hit the GitHub API in - production. Tests that exercise the reconsider happy path stub - them to "yes the bot closed it, no recent reconsider comment" - so the test stays focused on its actual assertion. - """ - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: None, - ) - - @staticmethod - def _stub_grace_aged_out(triage_module, monkeypatch): - """Pretend the grace warning has aged out. - - For tests that exercise the post-grace close path. Set the age - to twice the grace window so a future tweak to - `GRACE_PERIOD_SECONDS` doesn't accidentally make the stub fall - back inside the window. - """ - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: triage_module.GRACE_PERIOD_SECONDS * 2, - ) - - @staticmethod - def _stub_grace_no_warning(triage_module, monkeypatch): - """Pretend no grace warning has been posted yet (first detection).""" - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: None, - ) - - def test_should_reopen_on_reconsider_pass(self, triage_module, monkeypatch): - # Reconsider on a closed PR with a passing verdict -> reopen + post a - # friendly "re-evaluated" comment. close=True is the production path - # (the workflow only adds --close when AGENT_SHIN_ENABLED=true). - pr = self._make_pr( - state="closed", body="Updated body with QA proof + screenshots." - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - # close_pr / close_issue MUST NOT fire in reconsider mode. - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on reconsider pass"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok now"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 42 - assert posted["n"] == 42 - assert "reopened" in posted["body"].lower() - - def test_should_dry_run_reconsider_pass_when_close_false( - self, triage_module, monkeypatch - ): - # Reconsider must honor `close=False` (dry-run) just like the - # regular triage flow. A local invocation of - # `python triage_with_llm.py --reconsider --pr N` (no --close) - # must NOT post a comment or reopen the PR — it should return - # `would-reopen` so the operator can preview the outcome. - pr = self._make_pr( - state="closed", body="Updated body with QA proof + screenshots." - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post comment in dry-run reconsider"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen PR in dry-run reconsider"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=False, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok now"} - ), - reconsider=True, - ) - assert result["action"] == "would-reopen" - # The previewed comment body is still returned so a step-summary - # writer can render exactly what would have been posted. - assert "reopened" in result["comment"].lower() - - def test_should_post_still_failing_on_reconsider_fail( - self, triage_module, monkeypatch - ): - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - # Neither reopen nor close should fire when reconsider verdict is fail. - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen on fail"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close again on reconsider fail"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Still no QA proof.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - reconsider=True, - ) - assert result["action"] == "reconsider-still-failing" - assert posted["n"] == 42 - assert "QA proof" in posted["body"] - - def test_should_not_reopen_on_reconsider_with_ambiguous_verdict( - self, triage_module, monkeypatch - ): - # Regression: only an explicit `pass` verdict reopens. Missing, - # empty, or unexpected verdict strings ("failed", "", garbage) - # must fall through to the still-failing branch rather than - # reopen a PR the rubric did not actually clear. - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen on ambiguous verdict"), - ) - - for ambiguous in ("", "failed", "needs-info", "unknown"): - posted.clear() - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p, v=ambiguous: json.dumps( - {"verdict": v, "missing": [], "explanation": "weird"} - ), - reconsider=True, - ) - assert result["action"] == "reconsider-still-failing", ambiguous - assert "body" in posted, ambiguous - - def test_should_dry_run_reconsider_fail_when_close_false( - self, triage_module, monkeypatch - ): - # Mirror dry-run behavior for the FAIL branch — `close=False` - # must NOT post the "still failing" comment. - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail( - "must not post still-failing comment in dry-run" - ), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Still no QA proof.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - reconsider=True, - ) - assert result["action"] == "would-reconsider-still-failing" - assert "QA proof" in result["comment"] - - def test_should_reopen_on_reconsider_with_linked_issue_short_circuit( - self, triage_module, monkeypatch - ): - # The linked-issue short-circuit also has to honor reconsider mode: - # if the contributor edited the body to add `Fixes #1234`, the regex - # path should reopen the PR without calling the LLM. - pr = self._make_pr(state="closed", body="Fixes #1234\n\nAddresses the bug.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=55, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run when linked-issue matches"), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 55 - assert "reopened" in posted["body"].lower() - - def test_should_dry_run_reconsider_with_linked_issue_when_close_false( - self, triage_module, monkeypatch - ): - # Linked-issue short-circuit must ALSO honor dry-run. - pr = self._make_pr(state="closed", body="Fixes #1234\n\nAddresses the bug.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post in dry-run"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen in dry-run"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=55, - close=False, - model="m", - judge=lambda p: pytest.fail("LLM must not run when linked-issue matches"), - reconsider=True, - ) - assert result["action"] == "would-reopen" - - def test_should_skip_internal_in_reconsider_mode(self, triage_module, monkeypatch): - # Internal authors are exempt from triage in both regular and - # reconsider mode — Agent Shin should never reopen one of their PRs - # automatically, in case a maintainer closed it intentionally. - pr = self._make_pr( - state="closed", - author_association="MEMBER", - user={"login": "krrishdholakia"}, - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen for internal author"), - ) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: pytest.fail("LLM must not run for internal author"), - reconsider=True, - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_should_skip_reconsider_when_not_bot_closed( - self, triage_module, monkeypatch - ): - # SECURITY: `@agent-shin reconsider` must NOT reopen a PR/issue - # that a MAINTAINER closed for non-rubric reasons (e.g. duplicate, - # design rejection, security report). Only PRs closed by the bot - # itself should ever be candidates for the reconsider reopen path. - pr = self._make_pr(state="closed", body="something.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: False - ) - # Even though there's no rate-limit conflict, the bot-closed guard - # alone is sufficient to block. The LLM judge must never run on a - # maintainer-closed PR. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: None, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment on maintainer-closed PR"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen maintainer-closed PR"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run before bot-closed guard"), - reconsider=True, - ) - assert result["action"] == "skip-not-bot-closed" - - def test_should_rate_limit_repeated_reconsider_triggers( - self, triage_module, monkeypatch - ): - # COST CONTROL: each `@agent-shin reconsider` event burns CI - # minutes + an OpenAI API call. If the bot already posted a - # reconsider verdict within the cooldown window - # (RECONSIDER_RATE_LIMIT_SECONDS), refuse to run again. This - # bounds the damage from a contributor spamming the trigger. - pr = self._make_pr(state="closed", body="something with new edits.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - # Pretend the bot posted a reconsider verdict 1 second ago. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: 1.0, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment during cooldown"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen during cooldown"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run during cooldown"), - reconsider=True, - ) - assert result["action"] == "skip-rate-limited" - assert result["rate_limit_age_seconds"] == 1.0 - assert ( - result["rate_limit_window_seconds"] - == triage_module.RECONSIDER_RATE_LIMIT_SECONDS - ) - - def test_should_allow_reconsider_after_cooldown_window( - self, triage_module, monkeypatch - ): - # The cooldown is a window, not a one-shot lock — once - # RECONSIDER_RATE_LIMIT_SECONDS has elapsed since the last bot - # verdict, a fresh `@agent-shin reconsider` is allowed through. - pr = self._make_pr(state="closed", body="updated with screenshots now.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - # Last reconsider was 1 hour ago — well outside the 10-min window. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: 3600.0, - ) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 1 - - def test_should_reopen_issue_on_reconsider_pass(self, triage_module, monkeypatch): - issue = { - "number": 7, - "title": "Bug: now with repro", - "body": "## Repro\n```bash\ncurl ...\n```\n\nExpected X, got Y.", - "state": "closed", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - monkeypatch.setattr(triage_module, "fetch_issue", lambda repo, n: issue) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_issue", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="issue", - number=7, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "now reproducible"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 7 - assert "reopened" in posted["body"].lower() - - def test_should_triage_issues_kind(self, triage_module, monkeypatch): - issue = { - "number": 7, - "title": "Bug: X is broken", - "body": "no detail", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - monkeypatch.setattr(triage_module, "fetch_issue", lambda repo, n: issue) - # Grace already aged out -> close path. (Issues use the same - # GRACE_COMMENT_MARKER detection as PRs.) - self._stub_grace_aged_out(triage_module, monkeypatch) - closed = {} - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update(body=body), - ) - monkeypatch.setattr( - triage_module, "close_issue", lambda repo, n: closed.update(n=n) - ) - - verdict = { - "verdict": "fail", - "kind": "bug", - "has_repro": False, - "missing": ["reproduction", "expected vs. actual"], - "explanation": "No repro provided.", - } - result = triage_module.triage( - repo="o/r", - kind="issue", - number=7, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "closed" - assert closed["n"] == 7 - assert "reproduction" in posted["body"] - - # ---- Grace-period flow ------------------------------------------------ - - def test_should_post_grace_warning_on_first_failing_run_in_close_mode( - self, triage_module, monkeypatch - ): - # First low-quality detection -> bot posts a warning comment with - # the GRACE_COMMENT_MARKER. The PR must NOT be closed yet. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on first detection"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "warned-grace" - assert posted["n"] == 42 - # Pin the user-facing language pieces the user explicitly asked for. - assert "2 hours" in posted["body"] - assert "@agent-shin reconsider" in posted["body"] - assert "@greptileai" in posted["body"] - assert "even after the PR is closed" in posted["body"] - assert triage_module.GRACE_COMMENT_MARKER in posted["body"] - - def test_should_skip_close_inside_grace_window(self, triage_module, monkeypatch): - # A warning was posted recently; do nothing on this run regardless - # of close=True. The next run after `GRACE_PERIOD_SECONDS` elapses - # is the one that flips to actual close. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: 60.0, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment during grace window"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close during grace window"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "skip-in-grace-period" - assert result["grace_age_seconds"] == 60.0 - assert result["grace_period_seconds"] == triage_module.GRACE_PERIOD_SECONDS - - def test_should_dry_run_grace_warning_when_close_false( - self, triage_module, monkeypatch - ): - # In dry-run mode the FIRST failing detection returns - # `would-warn-grace` (with the previewed comment body) and never - # touches GitHub state. Lets a local operator preview the - # warning before flipping --close on. - pr = self._make_pr(body="thin") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post in dry-run grace warn"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "thin", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "would-warn-grace" - assert "2 hours" in result["comment"] - - def test_should_warn_grace_for_swiftwinds_not_close_instantly( - self, triage_module, monkeypatch - ): - # Regression: SwiftWinds (the dogfood account) used to be in a - # now-removed `IMMEDIATE_CLOSE_LOGINS` bypass that skipped the grace - # window and closed on first detection. It must follow the SAME - # grace path as every other author: warn first, close only after the - # window elapses. A re-added instant-close bypass would call - # close_pr here and fail the test. - pr = self._make_pr(body="just a sentence.", user={"login": "SwiftWinds"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail( - "SwiftWinds must not close on first detection; it gets the grace window" - ), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=99, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "warned-grace" - assert "2 hours" in posted["body"] - - -class TestGraceWarningCommentText: - """Pin the user-facing promises in the grace warning so a future - refactor can't silently drop them.""" - - def test_pr_grace_warning_should_state_grace_window(self, triage_module): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"} - ) - # The user explicitly asked: "specify in the comment" the grace window. - assert "2 hours" in body - - def test_pr_grace_warning_should_mention_reconsider_during_grace( - self, triage_module - ): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "@agent-shin reconsider" in body - - def test_pr_grace_warning_should_promise_greptileai_works_post_close( - self, triage_module - ): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - # Per user: comment should state @greptileai works even after close. - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_pr_grace_warning_should_carry_grace_marker(self, triage_module): - # The marker is what `seconds_since_last_grace_warning` greps for - # on subsequent runs to detect that a warning has been posted. - # Dropping it would silently break the close-after-grace path. - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert triage_module.GRACE_COMMENT_MARKER in body - - def test_issue_grace_warning_should_carry_grace_marker(self, triage_module): - body = triage_module.format_grace_warning_issue_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert triage_module.GRACE_COMMENT_MARKER in body - assert "2 hours" in body - # OSS authors can't reopen a bot-closed issue, so recovery is - # `@agent-shin reconsider` (the bot reopens), like the PR path. - assert "@agent-shin reconsider" in body - - def test_pr_close_comment_should_promise_greptileai_works_post_close( - self, triage_module - ): - # The standard close comment must ALSO point at @greptileai so - # contributors see the same options whether they read the warning - # or only catch the close comment. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_pr_grace_warning_should_not_prompt_reconsider_during_grace_window( - self, triage_module - ): - # Per user feedback: during the 24h grace window, the contributor - # should just update the PR description. Asking them to also comment - # "@agent-shin reconsider" right away adds a step they don't need — - # the bot re-checks automatically on the next sweep. The reconsider - # trigger is reserved for the post-close recovery path. - # - # We pin this by checking that the grace section explicitly tells - # the contributor they don't need to ping the bot during the grace - # window. The presence of "@agent-shin reconsider" elsewhere in the - # comment (as the post-close path) is fine and required by other - # tests. - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "No need to ping" in body or "no need to ping" in body - - def test_grace_warnings_should_show_what_got_right(self, triage_module): - # The "What you got right" section must appear in the grace warning - # too, not only the close comment — the contributor sees the warning - # first and that's their best chance to know what to keep. - pr_body = triage_module.format_grace_warning_pr_comment( - { - "verdict": "fail", - "linked_issue": True, - "has_problem_description": True, - "has_expected_vs_actual": True, - "has_qa_proof": False, - "missing": ["QA proof"], - "explanation": "thin", - } - ) - assert "What you got right" in pr_body - assert "Linked a related GitHub issue" in pr_body - - issue_body = triage_module.format_grace_warning_issue_comment( - { - "verdict": "fail", - "kind": "feature", - "has_motivation_example": True, - "missing": ["concrete description"], - "explanation": "vague", - } - ) - assert "What you got right" in issue_body - assert "Motivation and concrete example" in issue_body - - def test_grace_warnings_should_use_softer_park_for_later_framing( - self, triage_module - ): - # Same softer-framing pin as the close comment, but for the warning - # — the contributor's first contact with the bot must not read as a - # hard deadline / ultimatum. - for body in ( - triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - triage_module.format_grace_warning_issue_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ): - assert "park this for later" in body - assert ( - "not a rejection" in body - or "isn't a rejection" in body - or ("isn't us saying" in body) - ) - - -class TestSecondsSinceLastGraceWarning: - """Mirror of TestSecondsSinceLastReconsiderVerdict for the new helper. - Both helpers share `_seconds_since_latest_marker_comment` underneath - so the parsing logic is exercised either way; these tests pin the - grace-marker-specific behavior.""" - - def _make_comment( - self, - *, - login: str, - body: str, - created_at: str | None = "2026-05-18T05:00:00Z", - ) -> dict: - comment: dict = {"user": {"login": login}, "body": body} - if created_at is not None: - comment["created_at"] = created_at - return comment - - def test_should_return_none_when_no_grace_marker(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body="Some other bot message", - ), - self._make_comment(login="random-user", body="ping?"), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_grace_warning("o/r", 1) is None - - def test_should_ignore_non_bot_comments_with_marker( - self, triage_module, monkeypatch - ): - # A user who quotes the marker in a question must NOT be treated - # as the bot warning; otherwise the close-after-grace path would - # never fire because the timer keeps resetting. - comments = [ - self._make_comment( - login="random-user", - body=f"What is {triage_module.GRACE_COMMENT_MARKER}?", - ) - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_grace_warning("o/r", 1) is None - - def test_should_pick_latest_grace_marker(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body="old warning " + triage_module.GRACE_COMMENT_MARKER, - created_at="2026-05-18T03:00:00Z", - ), - self._make_comment( - login="github-actions[bot]", - body="newer warning " + triage_module.GRACE_COMMENT_MARKER, - created_at="2026-05-18T04:55:00Z", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - - import datetime as real_dt - - class FrozenDateTime(real_dt.datetime): - @classmethod - def now(cls, tz=None): - return real_dt.datetime(2026, 5, 18, 5, 0, 0, tzinfo=tz) - - frozen_module = type(triage_module.dt)("datetime") - frozen_module.datetime = FrozenDateTime - frozen_module.timezone = real_dt.timezone - monkeypatch.setattr(triage_module, "dt", frozen_module) - - age = triage_module.seconds_since_last_grace_warning("o/r", 1) - # Newer warning is 5 minutes (300s) before "now". - assert age == 300.0 - - -class TestTriageAllowlist: - """The dogfood allowlist gates `triage`: while non-empty it is the sole - author filter (only the named accounts are acted on) and it bypasses the - internal-author exemption for them, so a maintainer can dogfood on their - own org account. Emptying it restores the internal-author skip.""" - - def _make_pr(self, **overrides): - base = { - "number": 1, - "title": "PR title", - "body": "Body with no linked issue and no QA proof.", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - base.update(overrides) - return base - - def test_should_skip_author_not_on_allowlist(self, triage_module, monkeypatch): - pr = self._make_pr(user={"login": "random-oss-dev"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run for non-allowlisted author"), - ) - assert result["action"] == "skip-not-allowlisted" - - def test_should_act_on_allowlisted_internal_author( - self, triage_module, monkeypatch - ): - pr = self._make_pr(author_association="MEMBER", user={"login": "mateo-berri"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok"} - ), - ) - assert result["action"] == "pass-llm" - - def test_empty_allowlist_restores_internal_skip(self, triage_module, monkeypatch): - pr = self._make_pr( - author_association="MEMBER", user={"login": "krrishdholakia"} - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run for internal author"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_allowlist_constant_is_the_two_dogfood_accounts(self, triage_module): - assert triage_module.ALLOWLIST_LOGINS == frozenset( - {"mateo-berri", "swiftwinds"} - ) - for login in triage_module.ALLOWLIST_LOGINS: - assert login == login.lower(), login diff --git a/tests/test_litellm/test_github_triage_workflows.py b/tests/test_litellm/test_github_triage_workflows.py deleted file mode 100644 index f96c9b7e974..00000000000 --- a/tests/test_litellm/test_github_triage_workflows.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Static guardrails for the Agent Shin + Greptile workflow YAML files. - -These workflows can post comments and close PRs/issues on -BerriAI/litellm, so the gating logic that decides "is this a real -close-on-fail run?" must fail-safe on any unexpected input. The risk -is mostly maintenance: someone edits the bash gate, drops a quote, -inverts a comparison, or uses `!= "false"` (which treats "True", -"yes", "1", and typos as enabling closure) and the regression isn't -caught until a real OSS contributor's PR gets auto-closed. - -The tests below pin a set of invariants. The first two apply to every -workflow that gates a destructive `--close`: - - 1. The gate uses the fail-safe `= "true"` comparison — not `!= "false"`, - not `!= ""`. Only the literal string "true" should ever enable - closure. - 2. The gate also requires `AGENT_SHIN_ENABLED = "true"` (or the - scheduled-job equivalent) — disabling the variable must always - force dry-run. - -A third invariant covers every workflow that installs the OpenAI client. -These run with a write-scoped `GITHUB_TOKEN`, so a compromised package -release would execute in that context; the install must therefore come -from the hash-pinned `.github/scripts/triage-requirements.txt` via -`pip --require-hashes`, never a floating `pip install openai>=...`. - -Static parsing of the YAML + bash text is the right level of test here: -the gating logic lives in a `run:` block, not in a Python module we can -import, and end-to-end testing a GitHub Actions workflow from CI is -infeasible. A YAML-level guardrail is exactly what would have caught -the original `!= "false"` regression at PR time. -""" - -from __future__ import annotations - -from pathlib import Path - -import pytest -import yaml - -REPO_ROOT = Path(__file__).resolve().parents[2] -WORKFLOWS_DIR = REPO_ROOT / ".github" / "workflows" - -# Map of workflow file -> the env var name that drives the destructive -# gate inside that workflow's `run:` block. Keeping this table explicit -# (rather than scraping every workflow file) means a new workflow file -# that bypasses the dry-run gating doesn't silently slip past this test. -DESTRUCTIVE_GATE_ENV: dict[str, str] = { - "close_low_quality_prs.yml": "CLOSE_FLAG", - # The reconsider workflow has no per-run "really do it?" knob — its - # only kill switch is `AGENT_SHIN_ENABLED`, which already serves as - # both the destructive gate and the global enablement gate. - "triage_reconsider.yml": "AGENT_SHIN_ENABLED", -} - - -# Privileged workflows that install the OpenAI client. They run with a -# write-scoped GITHUB_TOKEN, so the install must be hash-pinned: a poisoned -# release would otherwise execute in that context. A new workflow that -# installs the client must be added here and use the same pinned file. -LLM_CLIENT_INSTALLER_WORKFLOWS = ( - "triage_reconsider.yml", -) - -PINNED_INSTALL = "--require-hashes -r .github/scripts/triage-requirements.txt" -REQUIREMENTS_FILE = REPO_ROOT / ".github" / "scripts" / "triage-requirements.txt" - - -def _load_workflow(name: str) -> dict: - return yaml.safe_load((WORKFLOWS_DIR / name).read_text()) - - -def _all_run_blocks(workflow: dict) -> list[str]: - """Return every `run:` step's command text, joined.""" - commands: list[str] = [] - jobs = workflow.get("jobs") or {} - for job in jobs.values(): - for step in job.get("steps", []) or []: - if not isinstance(step, dict): - continue - run = step.get("run") - if isinstance(run, str): - commands.append(run) - return commands - - -@pytest.mark.parametrize("workflow_file,env_var", sorted(DESTRUCTIVE_GATE_ENV.items())) -def test_should_use_failsafe_equals_true_comparison(workflow_file: str, env_var: str) -> None: - """The destructive `--close` gate must use `= "true"` (fail-safe), not - `!= "false"` (which would treat "True", "yes", "1", or any typo as - enabling closure). - - Both bare `${ENV_VAR}` and `${ENV_VAR:-false}` (with a default) are - accepted forms — what matters is the comparison operator. The - Greptile closer relies on an outer `AGENT_SHIN_ENABLED` gate so it - can use the bare form; the Agent Shin workflows include `:-false` - for defense in depth. Either is fine. - """ - workflow = _load_workflow(workflow_file) - text = "\n".join(_all_run_blocks(workflow)) - assert env_var in text, ( - f"{workflow_file} no longer references {env_var}; was the gating env var renamed without updating this test?" - ) - accepted_patterns = ( - f'"${{{env_var}}}" = "true"', - f'"${{{env_var}:-false}}" = "true"', - ) - assert any(p in text for p in accepted_patterns), ( - f"{workflow_file} must gate the destructive --close flag on the " - f'EXACT string "true" (one of: {accepted_patterns!r}). Mirror ' - 'the Greptile closer pattern; do NOT use `!= "false"` which ' - 'fail-opens on unknown values like "True", "yes", "1", or typos.' - ) - forbidden_patterns = ( - f'"${{{env_var}}}" != "false"', - f'"${{{env_var}:-false}}" != "false"', - f'"${{{env_var}:-true}}" != "false"', - ) - for forbidden in forbidden_patterns: - assert forbidden not in text, ( - f"{workflow_file} uses the fail-open pattern {forbidden!r}. " - 'Switch to `= "true"` so unknown values stay dry-run.' - ) - - -@pytest.mark.parametrize("workflow_file", sorted(DESTRUCTIVE_GATE_ENV)) -def test_should_require_agent_shin_enabled_for_close(workflow_file: str) -> None: - """Every destructive gate must also gate on the global enablement - variable, so flipping `AGENT_SHIN_ENABLED` off is a kill switch - regardless of any per-run input. - - Two patterns are equally fine: - - Positive: `[ "${AGENT_SHIN_ENABLED:-false}" = "true" ]` to enter - the close branch (Agent Shin workflows). - - Negative: `[ "${AGENT_SHIN_ENABLED:-false}" != "true" ]` then - bail out / force dry-run (Greptile closer). - - What matters is that the comparison value is the literal "true"; - `!= "false"` or `= "1"` etc. would not be a true kill switch. - """ - workflow = _load_workflow(workflow_file) - text = "\n".join(_all_run_blocks(workflow)) - accepted_patterns = ( - '"${AGENT_SHIN_ENABLED:-false}" = "true"', - '"${AGENT_SHIN_ENABLED:-false}" != "true"', - ) - assert any(p in text for p in accepted_patterns), ( - f"{workflow_file} must gate destructive actions on " - '`AGENT_SHIN_ENABLED = "true"` (or the inverted `!= "true"` ' - "guard that forces dry-run). Without this, an unset repo " - "variable would not be treated as a kill switch." - ) - - -@pytest.mark.parametrize("workflow_file", LLM_CLIENT_INSTALLER_WORKFLOWS) -def test_llm_client_install_is_hash_pinned(workflow_file: str) -> None: - """Every privileged workflow installs the OpenAI client from the - hash-pinned requirements file, never by floating version. - - A bare `pip install "openai>=1.40.0"` resolves to whatever PyPI serves - at run time and executes during install/import while a write-scoped - `GITHUB_TOKEN` is in scope, so a compromised release runs in a - privileged context. This test fails if that floating form comes back or - if the `--require-hashes` install is loosened. - """ - blocks = _all_run_blocks(_load_workflow(workflow_file)) - assert PINNED_INSTALL in "\n".join(blocks), ( - f"{workflow_file} must install the client via `pip install " - f"{PINNED_INSTALL}`; a floating install runs unverified code with a " - "write-scoped token." - ) - offenders = [b for b in blocks if "pip install" in b and "openai" in b] - assert not offenders, ( - f"{workflow_file} installs openai by name ({offenders!r}); pin it " - "through the hash-locked requirements file so the version and " - "checksum are fixed." - ) - - -def test_triage_requirements_are_fully_hash_pinned() -> None: - """The shared requirements file pins every package to an exact version - with a sha256 hash, which is what `pip --require-hashes` enforces at - install time. A loosened pin or a missing hash here would silently widen - the supply-chain surface for all the installer workflows. - """ - assert REQUIREMENTS_FILE.exists(), ( - f"the hash-pinned requirements file the triage workflows install from is missing at {REQUIREMENTS_FILE}" - ) - joined = REQUIREMENTS_FILE.read_text().replace("\\\n", " ") - entries = [line.strip() for line in joined.splitlines() if line.strip() and not line.strip().startswith("#")] - assert any(e.split()[0].startswith("openai==") for e in entries), ( - "openai must be pinned to an exact version in the triage requirements" - ) - for entry in entries: - spec = entry.split()[0] - assert "==" in spec, ( - f"requirement {spec!r} is not pinned to an exact version; " - "--require-hashes needs every package pinned with ==" - ) - assert "--hash=sha256:" in entry, ( - f"requirement {spec!r} has no sha256 hash; every pin must carry " - "checksums so --require-hashes can verify the download" - ) - - -def _reconsider_steps() -> list[dict]: - workflow = _load_workflow("triage_reconsider.yml") - return workflow["jobs"]["reconsider"]["steps"] - - -def _index_of_run_step(steps: list[dict], needle: str) -> int: - for i, step in enumerate(steps): - run = step.get("run") - if isinstance(run, str) and needle in run: - return i - raise AssertionError(f"no run step contains {needle!r}") - - -def _reaction_steps(steps: list[dict], content: str) -> list[tuple[int, dict]]: - return [ - (i, s) - for i, s in enumerate(steps) - if isinstance(s.get("run"), str) and f"content={content}" in s["run"] and "/reactions" in s["run"] - ] - - -class TestReconsiderReactions: - """The reconsider workflow acknowledges the triggering comment with a 👀 - reaction the moment it accepts the trigger, and a 👍 once the run finishes, - so the contributor gets feedback immediately instead of waiting on a cron. - - Both reactions are gated on `AGENT_SHIN_ENABLED == 'true'` so a dry-run - leaves no visible trace, and both target the comment that fired the event - (`github.event.comment.id`). The ordering (👀 before the triage run, 👍 - after) is the whole point — these tests fail if a refactor reorders the - steps, drops a reaction, or stops gating them. - """ - - def test_eyes_reaction_is_posted_before_the_triage_run(self) -> None: - steps = _reconsider_steps() - run_idx = _index_of_run_step(steps, "triage_with_llm.py") - eyes = _reaction_steps(steps, "eyes") - assert len(eyes) == 1, "expected exactly one 👀 (eyes) reaction step" - idx, step = eyes[0] - assert idx < run_idx, "👀 must be posted BEFORE the slow triage run, not after" - assert "github.event.comment.id" in (step.get("env") or {}).get("COMMENT_ID", ""), ( - "👀 must react to the comment that triggered the workflow" - ) - assert "${COMMENT_ID}" in step["run"], "👀 must react to the triggering comment, not a hardcoded id" - assert "vars.AGENT_SHIN_ENABLED == 'true'" in step["if"], ( - "👀 must be gated on AGENT_SHIN_ENABLED so dry-run stays inert" - ) - - def test_thumbs_up_reaction_is_posted_after_a_successful_run(self) -> None: - steps = _reconsider_steps() - run_idx = _index_of_run_step(steps, "triage_with_llm.py") - thumbs = _reaction_steps(steps, "+1") - assert len(thumbs) == 1, "expected exactly one 👍 (+1) reaction step" - idx, step = thumbs[0] - assert idx > run_idx, "👍 must come AFTER the triage run" - assert "success()" in step["if"], "👍 must only fire when the reconsider run succeeded" - assert "vars.AGENT_SHIN_ENABLED == 'true'" in step["if"], ( - "👍 must be gated on AGENT_SHIN_ENABLED so dry-run stays inert" - ) From a4624b6c6c4bb1ddf753bd57bbfd37b711598c90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:03:45 -0700 Subject: [PATCH 33/33] fix(gemini): derive the finish reason key set from the Candidates type Candidates.finishReason listed eleven values while the mapping key set carried twenty-one, so typed fixtures could not spell the reasons this PR handles. GeminiFinishReason is now the one list, the key set derives from it, and a test checks every documented reason has an explicit mapping instead of falling through to "stop" --- .../vertex_and_google_ai_studio_gemini.py | 29 ++------------ litellm/types/llms/vertex_ai.py | 39 ++++++++++++------- ...test_vertex_and_google_ai_studio_gemini.py | 9 ++++- 3 files changed, 36 insertions(+), 41 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a95b845718a..46f1b948026 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -6,7 +6,7 @@ import time from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args import httpx @@ -57,6 +57,7 @@ from litellm.types.llms.vertex_ai import ( ContentType, FunctionCallingConfig, FunctionDeclaration, + GeminiFinishReason, GeminiThinkingConfig, GenerateContentResponseBody, HttpxPartType, @@ -1330,31 +1331,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } - _GEMINI_FINISH_REASON_KEYS = frozenset( - { - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "FINISH_REASON_UNSPECIFIED", - "MALFORMED_FUNCTION_CALL", - "LANGUAGE", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "IMAGE_SAFETY", - "IMAGE_PROHIBITED_CONTENT", - "TOO_MANY_TOOL_CALLS", - "MALFORMED_RESPONSE", - "NO_IMAGE", - "IMAGE_RECITATION", - "IMAGE_OTHER", - "ESCALATION", - "UNEXPECTED_TOOL_CALL", - "MISSING_THOUGHT_SIGNATURE", - } - ) + _GEMINI_FINISH_REASON_KEYS: Final[frozenset[str]] = frozenset(get_args(GeminiFinishReason)) @staticmethod def get_finish_reason_mapping() -> dict[str, OpenAIChatCompletionFinishReason]: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 3b95b786631..ce51e46ef15 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -425,22 +425,35 @@ class UrlContextMetadata(TypedDict, total=False): urlMetadata: list[UrlMetadata] +GeminiFinishReason = Literal[ + "FINISH_REASON_UNSPECIFIED", + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "MALFORMED_FUNCTION_CALL", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "TOO_MANY_TOOL_CALLS", + "MALFORMED_RESPONSE", + "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + "ESCALATION", + "UNEXPECTED_TOOL_CALL", + "MISSING_THOUGHT_SIGNATURE", +] + + class Candidates(TypedDict, total=False): index: int content: HttpxContentType - finishReason: Literal[ - "FINISH_REASON_UNSPECIFIED", - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "MALFORMED_FUNCTION_CALL", - "IMAGE_SAFETY", - ] + finishReason: GeminiFinishReason safetyRatings: list[SafetyRatings] citationMetadata: CitationMetadata groundingMetadata: GroundingMetadata diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 4ee199bd9af..6c818016c87 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio import json import re from copy import deepcopy -from typing import Final, List, cast +from typing import Final, List, cast, get_args from unittest.mock import MagicMock, patch import httpx @@ -18,7 +18,7 @@ from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.types.llms.vertex_ai import UsageMetadata +from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import CustomStreamWrapper @@ -940,6 +940,11 @@ def test_check_finish_reason(): ) +def test_every_documented_gemini_finish_reason_has_an_explicit_mapping(): + documented: Final = frozenset(get_args(GeminiFinishReason)) + assert set(VertexGeminiConfig.get_finish_reason_mapping()) == documented + + def test_finish_reason_unspecified_and_malformed_function_call(): """ Test that FINISH_REASON_UNSPECIFIED and MALFORMED_FUNCTION_CALL