diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index da8149eab9b..362581937d7 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -9,9 +9,11 @@ captured stdout back through the typed agentic loop plan. import json import time import uuid -from typing import Any, cast +from typing import Any, Literal, TypedDict, cast import litellm +from pydantic import ValidationError + from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.integrations.code_interpreter_interception import ( @@ -20,15 +22,93 @@ from litellm.types.integrations.code_interpreter_interception import ( from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + is_interception_internal_key, +) +from litellm.types.llms.openai import ( + ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, + ChatCompletionToolMessage, +) +from litellm.types.utils import ( + CallTypes, + ChatCompletionMessageToolCall, + ModelResponse, ) -from litellm.types.utils import CallTypes LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" _INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +_CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream" +_LITELLM_METADATA_KEY = "litellm_metadata" _CACHE_TTL_SECONDS = 15 * 60 +class CodeExecutionToolCall(TypedDict, total=False): + id: str | None + call_id: str | None + type: Literal["function"] + name: str + arguments: str + + +class CodeInterpreterLogOutput(TypedDict): + type: Literal["logs"] + logs: str + + +class CodeInterpreterCall(TypedDict): + id: str + type: Literal["code_interpreter_call"] + status: Literal["completed"] + code: str + container_id: str | None + outputs: list[CodeInterpreterLogOutput] + + +class CodeExecutionFunctionParameters(TypedDict): + type: Literal["object"] + properties: dict[str, dict[str, str]] + required: list[str] + + +class ResponsesFunctionTool(TypedDict): + type: Literal["function"] + name: str + description: str + parameters: CodeExecutionFunctionParameters + + +class ChatCompletionFunctionDefinition(TypedDict): + name: str + description: str + parameters: CodeExecutionFunctionParameters + + +class ChatCompletionFunctionTool(TypedDict): + type: Literal["function"] + function: ChatCompletionFunctionDefinition + + +CodeExecutionFunctionTool = ResponsesFunctionTool | ChatCompletionFunctionTool + + +class ResponsesFunctionToolChoice(TypedDict): + type: Literal["function"] + name: str + + +class ChatCompletionFunctionToolChoice(TypedDict): + type: Literal["function"] + function: dict[str, str] + + +CodeExecutionFunctionToolChoice = ( + ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice +) + + def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None: try: from litellm.sandbox.sandbox_tools import resolve_sandbox_tool @@ -97,9 +177,15 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if not kwargs.get("_agentic_loop_depth"): kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None) kwargs.pop(_SANDBOX_KEY, None) + self._strip_interception_metadata(kwargs) if not self.enabled: return None - if call_type not in (CallTypes.responses, CallTypes.aresponses): + if call_type not in ( + CallTypes.responses, + CallTypes.aresponses, + CallTypes.completion, + CallTypes.acompletion, + ): return None if ( self.enabled_providers is not None @@ -120,18 +206,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): kwargs[_SANDBOX_KEY] = uuid.uuid4().hex if kwargs.get("stream"): kwargs["stream"] = False - kwargs["_code_interpreter_interception_converted_stream"] = True + kwargs[_CONVERTED_STREAM_KEY] = True + self._write_interception_metadata(kwargs) - function_tool = { - "type": "function", - "name": LITELLM_CODE_EXECUTION_TOOL_NAME, - "description": "Execute python code in a sandbox and return stdout.", - "parameters": { - "type": "object", - "properties": {"code": {"type": "string"}}, - "required": ["code"], - }, - } + function_tool = self._get_function_tool(call_type=call_type) kwargs["tools"] = [ ( function_tool @@ -141,19 +219,90 @@ class CodeInterpreterInterceptionLogger(CustomLogger): for tool in tools ] if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): - kwargs["tool_choice"] = { - "type": "function", - "name": LITELLM_CODE_EXECUTION_TOOL_NAME, - } + kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type) return kwargs + @staticmethod + def _strip_interception_metadata(kwargs: dict[str, Any]) -> None: + metadata = kwargs.get(_LITELLM_METADATA_KEY) + if not isinstance(metadata, dict): + return + filtered_metadata = { + key: value + for key, value in metadata.items() + if not is_interception_internal_key(key) + and not key.startswith("_agentic_loop") + and key != "max_agentic_loops" + } + if filtered_metadata: + kwargs[_LITELLM_METADATA_KEY] = filtered_metadata + else: + kwargs.pop(_LITELLM_METADATA_KEY, None) + + @staticmethod + def _write_interception_metadata(kwargs: dict[str, Any]) -> None: + metadata = kwargs.get(_LITELLM_METADATA_KEY) + metadata = dict(metadata) if isinstance(metadata, dict) else {} + for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _CONVERTED_STREAM_KEY): + if key in kwargs: + metadata[key] = kwargs[key] + kwargs[_LITELLM_METADATA_KEY] = metadata + + @staticmethod + def _get_function_parameters() -> CodeExecutionFunctionParameters: + return { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + } + + def _get_function_tool( + self, call_type: CallTypes | None + ) -> CodeExecutionFunctionTool: + description = "Execute python code in a sandbox and return stdout." + if call_type in (CallTypes.completion, CallTypes.acompletion): + return { + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": description, + "parameters": self._get_function_parameters(), + }, + } + return { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": description, + "parameters": self._get_function_parameters(), + } + + @staticmethod + def _get_function_tool_choice( + call_type: CallTypes | None, + ) -> CodeExecutionFunctionToolChoice: + if call_type in (CallTypes.completion, CallTypes.acompletion): + return { + "type": "function", + "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}, + } + return { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + } + @staticmethod def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool: if not isinstance(tool_choice, dict): return False + function = tool_choice.get("function") return ( tool_choice.get("type") == "code_interpreter" or tool_choice.get("name") == "code_interpreter" + or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + or ( + isinstance(function, dict) + and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + ) ) def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None: @@ -188,7 +337,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ): return False, {} - tool_calls = self._extract_code_execution_tool_calls(response=response) + tool_calls = ( + self._extract_chat_completion_code_execution_tool_calls(response=response) + if kwargs.get("_agentic_loop_api_surface") + == CHAT_COMPLETION_AGENTIC_SURFACE + else self._extract_code_execution_tool_calls(response=response) + ) if not tool_calls: return False, {} @@ -206,15 +360,24 @@ class CodeInterpreterInterceptionLogger(CustomLogger): stream: bool, kwargs: dict, ) -> AgenticLoopPlan: + if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE: + return await self._build_chat_completion_agentic_loop_plan( + tools=tools, + model=model, + messages=messages, + optional_params=anthropic_messages_optional_request_params, + kwargs=kwargs, + ) + await self._prune_expired_cache() - tool_calls = cast(list[dict[str, Any]], tools.get("tool_calls", [])) + tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) sandbox_key = kwargs.get(_SANDBOX_KEY) container, params = await self._get_or_create_container(cache_key=sandbox_key) try: - container_id = getattr(container, "id", None) + container_id = cast(str | None, getattr(container, "id", None)) input_list = self._normalize_messages(messages) - code_interpreter_calls = [] + code_interpreter_calls: list[CodeInterpreterCall] = [] for tool_call in tool_calls: arguments = tool_call.get("arguments", "") code = self._parse_code(arguments) @@ -256,9 +419,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger): request_patch = AgenticLoopRequestPatch( model=model, messages=input_list, - tools=optional_params.get("tools"), - optional_params={k: v for k, v in optional_params.items() if k != "tools"}, - kwargs={k: v for k, v in kwargs.items() if k != "litellm_logging_obj"}, + tools=self._get_followup_tools( + tools=optional_params.get("tools"), + call_type=CallTypes.responses, + ), + optional_params=self._get_followup_optional_params(optional_params), + kwargs=self._filter_agentic_loop_kwargs(kwargs), ) return AgenticLoopPlan( @@ -271,12 +437,134 @@ class CodeInterpreterInterceptionLogger(CustomLogger): }, ) + async def _build_chat_completion_agentic_loop_plan( + self, + tools: dict[str, object], + model: str, + messages: list[dict], + optional_params: dict[str, object], + kwargs: dict[str, object], + ) -> AgenticLoopPlan: + await self._prune_expired_cache() + tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) + sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY)) + container, params = await self._get_or_create_container(cache_key=sandbox_key) + + try: + container_id = cast(str | None, getattr(container, "id", None)) + tool_results = [ + await self._build_chat_completion_tool_result( + container=container, + params=params, + tool_call=tool_call, + container_id=container_id, + ) + for tool_call in tool_calls + ] + except Exception: + await self._delete_container_for_cache_key(sandbox_key) + raise + tool_messages = [result[0] for result in tool_results] + code_interpreter_calls = [result[1] for result in tool_results] + + request_patch = AgenticLoopRequestPatch( + model=model, + messages=list(messages) + + [self._build_chat_completion_assistant_message(tool_calls)] + + tool_messages, + tools=self._get_followup_tools( + tools=optional_params.get("tools"), + call_type=CallTypes.completion, + ), + optional_params=self._get_followup_optional_params(optional_params), + kwargs=self._filter_agentic_loop_kwargs(kwargs), + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={ + "tool_type": "code_interpreter", + "sandbox_key": sandbox_key or "", + "code_interpreter_calls": code_interpreter_calls, + "response_format": "openai", + }, + ) + + async def _build_chat_completion_tool_result( + self, + container: object, + params: dict[str, Any] | None, + tool_call: CodeExecutionToolCall, + container_id: str | None, + ) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]: + arguments = tool_call.get("arguments", "") + code = self._parse_code(arguments) + stdout = await self._run_tool_call( + container=container, params=params, arguments=arguments + ) + tool_call_id = ( + tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex + ) + return ( + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": stdout, + }, + { + "id": f"ci_{uuid.uuid4().hex}", + "type": "code_interpreter_call", + "status": "completed", + "code": code, + "container_id": container_id, + "outputs": [{"type": "logs", "logs": stdout}] if stdout else [], + }, + ) + async def async_agentic_loop_cleanup_hook( self, plan: AgenticLoopPlan, kwargs: dict ) -> None: metadata = plan.metadata or {} if plan else {} await self._delete_container_for_cache_key(metadata.get("sandbox_key")) + @staticmethod + def _filter_agentic_loop_kwargs(kwargs: dict[str, object]) -> dict[str, object]: + return { + k: v + for k, v in kwargs.items() + if k not in {"litellm_logging_obj", "acompletion"} + and not is_interception_internal_key( + k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES + ) + } + + def _get_followup_tools( + self, tools: object, call_type: CallTypes | None + ) -> list[dict[str, Any]] | None: + if not isinstance(tools, list): + return None + return [ + ( + self._get_function_tool(call_type=call_type) + if isinstance(tool, dict) and tool.get("type") == "code_interpreter" + else tool + ) + for tool in tools + ] + + def _get_followup_optional_params( + self, optional_params: dict[str, object] + ) -> dict[str, object]: + drop_tool_choice = self._tool_choice_targets_code_interpreter( + optional_params.get("tool_choice") + ) + return { + k: v + for k, v in optional_params.items() + if k != "tools" and not (k == "tool_choice" and drop_tool_choice) + } + async def async_post_agentic_loop_response_hook( self, response: Any, plan: AgenticLoopPlan, kwargs: dict ) -> Any: @@ -420,7 +708,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return list(messages) return [] - def _extract_code_execution_tool_calls(self, response: Any) -> list[dict[str, Any]]: + def _extract_code_execution_tool_calls( + self, response: object + ) -> list[CodeExecutionToolCall]: if isinstance(response, dict): output = response.get("output", []) else: @@ -446,6 +736,82 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if self._is_code_execution_call(item) ] + def _extract_chat_completion_code_execution_tool_calls( + self, response: ModelResponse | dict[str, Any] + ) -> list[CodeExecutionToolCall]: + model_response = self._to_model_response(response) + if model_response is None: + return [] + choices = model_response.choices or [] + if not choices: + return [] + message = choices[0].message + tool_calls = message.tool_calls or [] + + return [ + normalized + for tool_call in tool_calls + if (normalized := self._normalize_chat_completion_tool_call(tool_call)) + is not None + ] + + @staticmethod + def _normalize_chat_completion_tool_call( + tool_call: ChatCompletionMessageToolCall, + ) -> CodeExecutionToolCall | None: + if ( + tool_call.type != "function" + or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME + ): + return None + + arguments = tool_call.function.arguments + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + elif not isinstance(arguments, str): + arguments = "" if arguments is None else str(arguments) + + return { + "id": tool_call.id, + "call_id": tool_call.id, + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": arguments, + } + + @staticmethod + def _build_chat_completion_assistant_message( + tool_calls: list[CodeExecutionToolCall], + ) -> ChatCompletionAssistantMessage: + return { + "role": "assistant", + "tool_calls": [ + cast( + ChatCompletionAssistantToolCall, + { + "id": tool_call.get("id"), + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": tool_call.get("arguments", ""), + }, + }, + ) + for tool_call in tool_calls + ], + } + + @staticmethod + def _to_model_response( + response: ModelResponse | dict[str, Any], + ) -> ModelResponse | None: + if isinstance(response, ModelResponse): + return response + try: + return ModelResponse(**response) + except (TypeError, ValidationError): + return None + def _is_code_execution_call(self, item: Any) -> bool: if isinstance(item, dict): return ( diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py new file mode 100644 index 00000000000..938e892bd50 --- /dev/null +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -0,0 +1,332 @@ +# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api + +import json +from typing import cast + +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + AgenticLoopPlan, + AgenticLoopRequestPatch, + is_interception_internal_key, +) +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper + +_FOLLOWUP_INTERNAL_PARAMS = frozenset( + ( + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + "_agentic_loop_api_surface", + ) +) + + +def _gate_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_should_run_agentic_loop + func = type(callback).async_should_run_agentic_loop + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _build_plan_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_build_agentic_loop_plan + func = type(callback).async_build_agentic_loop_plan + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _post_hook_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_post_agentic_loop_response_hook + func = type(callback).async_post_agentic_loop_response_hook + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _coerce_int(value: object, default: int) -> int: + return int(value) if isinstance(value, (int, str)) else default + + +def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]: + depth = _coerce_int(kwargs.get("_agentic_loop_depth"), 0) + max_loops = max(_coerce_int(kwargs.get("max_agentic_loops"), 3), 1) + raw_fingerprints = kwargs.get("_agentic_loop_fingerprints") + fingerprints = ( + [str(fp) for fp in raw_fingerprints] + if isinstance(raw_fingerprints, list) + else [] + ) + return depth, max_loops, fingerprints + + +def _fingerprint_tools(tool_calls: object) -> str: + try: + return json.dumps(tool_calls, sort_keys=True, default=str) + except Exception: + return str(tool_calls) + + +def _check_agentic_loop_safety( + tool_calls: object, + fingerprints: list[str], + depth: int, + max_loops: int, + model: str, +) -> str: + fingerprint = _fingerprint_tools(tool_calls) + if fingerprint in fingerprints: + raise ValueError( + "Agentic loop detected repeated tool-call fingerprint; aborting rerun" + ) + if depth >= max_loops: + raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}") + return fingerprint + + +def _wrap_response_as_fake_stream(response: object) -> object: + if getattr(response, "object", None) == "chat.completion.chunk": + return response + if not hasattr(response, "choices"): + return response + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) + + return convert_model_response_to_streaming(cast(ModelResponse, response)) + + +def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None: + metadata = kwargs_for_followup.get("litellm_metadata") + metadata = dict(metadata) if isinstance(metadata, dict) else {} + for key, value in kwargs_for_followup.items(): + if ( + key.startswith("_agentic_loop") + or key == "max_agentic_loops" + or is_interception_internal_key(key) + ): + metadata[key] = value + kwargs_for_followup["litellm_metadata"] = metadata + + +def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: + return { + k: v + for k, v in source.items() + if not is_interception_internal_key( + k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES + ) + and k not in _FOLLOWUP_INTERNAL_PARAMS + } + + +async def _execute_chat_completion_agentic_plan( + *, + plan: AgenticLoopPlan, + callback: CustomLogger, + model: str, + optional_params: dict[str, object], + kwargs: dict[str, object], + logging_obj: object, + custom_llm_provider: str, + depth: int, + max_loops: int, + fingerprints: list[str], + fingerprint: str, +) -> object: + import litellm + + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = patch.model or model + if "/" not in full_model_name: + full_model_name = f"{custom_llm_provider}/{full_model_name}" + + optional_params_for_followup = {**optional_params, **patch.optional_params} + if patch.tools is not None: + optional_params_for_followup["tools"] = patch.tools + if "tool_choice" not in patch.optional_params: + optional_params_for_followup.pop("tool_choice", None) + + kwargs_for_followup = _filter_followup_kwargs(kwargs) + kwargs_for_followup.update( + { + k: v + for k, v in _filter_followup_kwargs(patch.kwargs).items() + if k not in optional_params_for_followup + } + ) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + _add_agentic_loop_metadata(kwargs_for_followup) + + try: + response_followup = await litellm.acompletion( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) + if _post_hook_overridden(callback): + try: + response_followup = ( + await callback.async_post_agentic_loop_response_hook( + response=response_followup, plan=plan, kwargs=kwargs + ) + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + if kwargs.get("_code_interpreter_interception_converted_stream") and not depth: + return _wrap_response_as_fake_stream(response_followup) + return response_followup + finally: + try: + await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + +async def maybe_run_chat_completion_agentic_loop( + *, + response: ModelResponse, + model: str, + messages: list, + optional_params: dict, + kwargs: dict, + logging_obj: object, + custom_llm_provider: str, + stream: bool, +) -> ModelResponse | CustomStreamWrapper | None: + import litellm + + callbacks = litellm.callbacks + ( + getattr(logging_obj, "dynamic_success_callbacks", None) or [] + ) + depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) + tools = optional_params.get("tools", []) + + for callback in callbacks: + if not isinstance(callback, CustomLogger): + continue + if not _gate_overridden(callback): + continue + + gate_kwargs = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + try: + should_run, tool_calls = await callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=gate_kwargs, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in chat completion agentic gate: %s", + str(e), + ) + continue + + if not should_run: + continue + + fingerprint = _check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + plan_kwargs = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + if not _build_plan_overridden(callback): + return await callback.async_run_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=plan_kwargs, + ) + + plan = await callback.async_build_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=plan_kwargs, + ) + + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + return response + if not plan.run_agentic_loop: + continue + + return await _execute_chat_completion_agentic_plan( + plan=plan, + callback=callback, + model=model, + optional_params=optional_params, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s", + str(e), + ) + + if ( + kwargs.get("_code_interpreter_interception_converted_stream") + and not depth + and hasattr(response, "choices") + ): + return cast( + "ModelResponse | CustomStreamWrapper", + _wrap_response_as_fake_stream(response), + ) + return None diff --git a/litellm/main.py b/litellm/main.py index a0fc662c6f6..a80109f4bff 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -81,6 +81,9 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, +) from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, @@ -654,6 +657,39 @@ async def acompletion( response_object=response, model_response_object=litellm.ModelResponse(), ) + # Provider-agnostic dispatch point for the chat-completions agentic loop + # (code-interpreter interception, etc). Chat routing forks per provider + # before this (OpenAI goes through the OpenAI SDK in openai.py, others + # through the shared httpx handler), so a dispatch inside any single + # provider handler would miss the others. Here is where every fork + # reconverges, so the loop runs once for all providers. Responses needs + # no equivalent: every provider already funnels through one shared + # handler where the loop is dispatched. + if isinstance(response, litellm.ModelResponse): + looped = await maybe_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + optional_params={ + k: v + for k, v in completion_kwargs.items() + if v is not None + and k + not in ( + "model", + "messages", + "stream", + "acompletion", + "deployment_id", + ) + }, + kwargs=kwargs, + logging_obj=kwargs.get("litellm_logging_obj"), + custom_llm_provider=custom_llm_provider, + stream=bool(stream), + ) + if looped is not None: + response = looped if isinstance(response, CustomStreamWrapper): response.set_logging_event_loop( loop=loop diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index b5726a11ca0..26a0be36ef4 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -2,6 +2,25 @@ from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field +CHAT_COMPLETION_AGENTIC_SURFACE = "chat_completions" +CODE_INTERPRETER_INTERCEPTION_PREFIX = "_code_interpreter_interception" +NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES = frozenset( + ("_websearch_interception", "_compression_interception") +) +INTERCEPTION_INTERNAL_PREFIXES = frozenset( + ( + *NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + CODE_INTERPRETER_INTERCEPTION_PREFIX, + ) +) + + +def is_interception_internal_key( + key: str, + prefixes: frozenset[str] = INTERCEPTION_INTERNAL_PREFIXES, +) -> bool: + return any(key.startswith(prefix) for prefix in prefixes) + class StandardCustomLoggerInitParams(BaseModel): """ diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 00b095f33ca..24d6e84fba7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3163,8 +3163,26 @@ class CustomPricingLiteLLMParams(BaseModel): regional_processing_uplift_multiplier_us: Optional[float] = None +# Server-controlled fields that bound or drive an interceptor's agentic loop +# (depth, cycle fingerprints, ceiling, code-interpreter sandbox state). Listed +# in all_litellm_params so they are treated as LiteLLM-level and excluded from +# get_non_default_completion_params; otherwise the OpenAI param builder sweeps +# any unrecognized top-level key into extra_body and leaks them to the provider. +# This is what lets the loop carry state across rerun calls without a provider +# scrubber. +agentic_loop_internal_litellm_params = [ + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_agentic_loop_api_surface", + "max_agentic_loops", + "_code_interpreter_interception_active", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_converted_stream", +] + all_litellm_params = ( - [ + agentic_loop_internal_litellm_params + + [ "metadata", "litellm_metadata", "litellm_trace_id", diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py index f33814b86df..7ff58ba6324 100644 --- a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py +++ b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py @@ -1,8 +1,8 @@ """ Unit tests for CodeInterpreterInterceptionLogger. -All sandbox dependencies are injected (dependency injection, no monkeypatch): -a FakeSandbox stands in for the real e2b config and records how it is called. +All sandbox dependencies are injected: a FakeSandbox stands in for the real e2b +config and records how it is called. """ import time @@ -12,13 +12,17 @@ import pytest from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, LITELLM_CODE_EXECUTION_TOOL_NAME, + _INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY, + _SANDBOX_KEY, +) +from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + is_interception_internal_key, ) from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult from litellm.types.utils import CallTypes -_ACTIVE_KEY = "_code_interpreter_interception_active" -_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" - class FakeHandle: def __init__(self, sandbox_id="sbx_fake"): @@ -51,6 +55,13 @@ class FakeLogging: def __init__(self, litellm_call_id="k1"): self.litellm_call_id = litellm_call_id self.model_call_details = {} + self.dynamic_success_callbacks = [] + + def pre_call(self, *args, **kwargs): + return None + + def post_call(self, *args, **kwargs): + return None def _function_call_item(call_id="c1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): @@ -62,6 +73,17 @@ def _function_call_item(call_id="c1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): } +def _chat_function_call_item(call_id="call_1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): + return { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": '{"code":"print(40 + 2)"}', + }, + } + + class FakeResponse: def __init__(self, output): self.output = output @@ -74,6 +96,18 @@ def _iter_messages(plan): return patch.messages +def test_interception_internal_key_prefix_sets_preserve_code_interpreter_state(): + assert is_interception_internal_key("_code_interpreter_interception_active") + assert not is_interception_internal_key( + "_code_interpreter_interception_active", + prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + ) + assert is_interception_internal_key( + "_websearch_interception_converted_stream", + prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + ) + + @pytest.mark.asyncio async def test_build_plan_runs_code_and_feeds_output_back(): sandbox = FakeSandbox(stdout="42") @@ -133,6 +167,30 @@ async def test_pre_call_converts_code_interpreter_tool(): assert LITELLM_CODE_EXECUTION_TOOL_NAME in names +@pytest.mark.asyncio +async def test_pre_call_converts_code_interpreter_tool_for_chat_completions(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": {"type": "code_interpreter"}, + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + tool = result["tools"][0] + assert tool["type"] == "function" + assert tool["function"]["name"] == LITELLM_CODE_EXECUTION_TOOL_NAME + assert tool["function"]["parameters"]["required"] == ["code"] + assert result["tool_choice"] == { + "type": "function", + "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}, + } + assert result["litellm_metadata"][_ACTIVE_KEY] is True + assert result["litellm_metadata"][_SANDBOX_KEY] == result[_SANDBOX_KEY] + + @pytest.mark.asyncio @pytest.mark.parametrize( "tool_choice", @@ -184,9 +242,24 @@ async def test_pre_call_noop_on_non_responses(): "custom_llm_provider": "openai", } + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aembedding) + + assert result is None + + +@pytest.mark.asyncio +async def test_pre_call_noop_on_chat_completion_without_code_interpreter_tool(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "web_search"}], + "custom_llm_provider": "openai", + } + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) assert result is None + assert _ACTIVE_KEY not in kwargs + assert _SANDBOX_KEY not in kwargs @pytest.mark.asyncio @@ -524,14 +597,142 @@ async def test_gate_rechecks_provider_scope(): assert should_run is False +@pytest.mark.asyncio +async def test_chat_completion_gate_detects_code_execution_tool_call(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = { + "choices": [ + {"message": {"tool_calls": [_chat_function_call_item(call_id="call_123")]}} + ] + } + + should_run, payload = await logger.async_should_run_agentic_loop( + response=response, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={ + _ACTIVE_KEY: True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert should_run is True + assert payload["tool_calls"][0]["id"] == "call_123" + assert payload["tool_calls"][0]["arguments"] == '{"code":"print(40 + 2)"}' + + +@pytest.mark.asyncio +async def test_chat_completion_gate_refuses_without_server_active_marker(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = {"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]} + + should_run, payload = await logger.async_should_run_agentic_loop( + response=response, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE}, + ) + + assert should_run is False + assert payload == {} + + +@pytest.mark.asyncio +async def test_chat_completion_build_plan_runs_code_and_appends_tool_message(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + native_chat_tool = {"type": "code_interpreter", "container": {"type": "auto"}} + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response={ + "choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}] + }, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "tools": [native_chat_tool], + "tool_choice": {"type": "code_interpreter", "container": {"type": "auto"}}, + "temperature": 0, + }, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={ + "acompletion": True, + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert sandbox.run_calls[0]["code"] == "print(40 + 2)" + patch = plan.request_patch + assert patch is not None + assert patch.tools == [ + { + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": "Execute python code in a sandbox and return stdout.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + assert patch.optional_params == {"temperature": 0} + assert patch.kwargs == { + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + } + assert patch.messages is not None + assert patch.messages[-2]["role"] == "assistant" + assert patch.messages[-2]["tool_calls"][0]["id"] == "call_1" + assert patch.messages[-1] == { + "role": "tool", + "tool_call_id": "call_1", + "content": "42", + } + assert plan.metadata["code_interpreter_calls"][0]["code"] == "print(40 + 2)" + + @pytest.mark.asyncio async def test_pre_call_strips_client_forged_marker_on_initial_request(): - """A client cannot pre-set the active marker on the original request.""" + """A client cannot pre-set the active marker on the original request: with no + native code_interpreter tool, any client-supplied interception markers in + litellm_metadata are scrubbed and the active flag in kwargs is cleared.""" logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) kwargs = { "tools": [{"type": "web_search"}], "custom_llm_provider": "openai", _ACTIVE_KEY: True, + "litellm_metadata": { + _ACTIVE_KEY: True, + _SANDBOX_KEY: "client-forged", + "safe_user_value": "kept", + }, } await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) @@ -540,6 +741,42 @@ async def test_pre_call_strips_client_forged_marker_on_initial_request(): "no native code_interpreter tool was present, so a client-supplied " "active marker must be cleared" ) + assert kwargs["litellm_metadata"] == {"safe_user_value": "kept"} + + +@pytest.mark.asyncio +async def test_pre_call_strips_forged_loop_controls_then_mints_own_markers(): + """On an INITIAL request (no server-set _agentic_loop_depth) a client cannot + smuggle loop-control state: forged _agentic_loop_depth / max_agentic_loops and + interception markers in litellm_metadata are stripped before the interceptor + activates, so the only interception markers that survive are the ones the + server mints for the converted code_interpreter tool.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "litellm_metadata": { + _ACTIVE_KEY: True, + _SANDBOX_KEY: "client-forged", + "_agentic_loop_depth": 99, + "max_agentic_loops": 999, + "safe_user_value": "kept", + }, + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + metadata = result["litellm_metadata"] + assert metadata["safe_user_value"] == "kept" + assert "_agentic_loop_depth" not in metadata, "forged loop depth must be stripped" + assert "max_agentic_loops" not in metadata, "forged loop cap must be stripped" + assert metadata[_ACTIVE_KEY] is True + assert metadata[_SANDBOX_KEY] == result[_SANDBOX_KEY] + assert metadata[_SANDBOX_KEY] != "client-forged", ( + "the surviving sandbox key must be the server-minted one, not the forged " + "value the client supplied" + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py new file mode 100644 index 00000000000..f1196ab4692 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -0,0 +1,422 @@ +""" +Tests for the provider-agnostic chat completion agentic loop dispatcher +(`litellm/litellm_core_utils/chat_completion_agentic_loop.py`) and the +code-interpreter interception integration that drives it. + +The load-bearing regression here protects a reviewer requirement: the internal +agentic/interception control fields must NEVER reach the outbound provider HTTP +request body. The relevant fields are: + + _agentic_loop_depth + _agentic_loop_fingerprints + _agentic_loop_api_surface + max_agentic_loops + _code_interpreter_interception_active + _code_interpreter_interception_sandbox_key + _code_interpreter_interception_converted_stream + +A scrubber in gpt_transformation.py used to strip these. That scrubber was +removed, so `test_internal_control_fields_never_leak_into_provider_body` proves +they stay out of the body even without it. +""" + +import os +import sys +from typing import Any, Dict, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, +) +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import ( + Choices, + Function, + ChatCompletionMessageToolCall, + Message, + ModelResponse, +) + +# The internal control fields that must never reach a provider request body. +_INTERNAL_CONTROL_FIELDS = ( + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_agentic_loop_api_surface", + "max_agentic_loops", + "_code_interpreter_interception_active", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_converted_stream", + "litellm_metadata", +) + + +@pytest.fixture +def restore_callbacks(): + """Save/restore litellm.callbacks so a registered fake logger never pollutes + other tests in the suite.""" + saved = list(litellm.callbacks) + try: + yield + finally: + litellm.callbacks = saved + + +class _SandboxResult: + def __init__(self, stdout: str) -> None: + self.stdout = stdout + self.error = None + + +class FakeSandboxConfig: + """Injected sandbox so the interception loop runs no real network / E2B.""" + + def __init__(self) -> None: + self.created = 0 + self.deleted = 0 + self.run_codes: List[str] = [] + + async def acreate_sandbox(self) -> Any: + self.created += 1 + return MagicMock(id="sandbox-123") + + async def arun_code(self, container: Any, code: str) -> _SandboxResult: + self.run_codes.append(code) + return _SandboxResult(stdout="42\n") + + async def adelete_sandbox(self, container: Any) -> None: + self.deleted += 1 + + +def _tool_call_model_response() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_abc", + type="function", + function=Function( + name="litellm_code_execution", + arguments='{"code": "print(6*7)"}', + ), + ) + ], + ), + ) + ] + ) + + +def _plain_model_response(content: str = "The answer is 42") -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="stop", + message=Message(role="assistant", content=content), + ) + ] + ) + + +def _raw_response_for(model_response: ModelResponse) -> MagicMock: + """Wrap a ModelResponse as the OpenAI `with_raw_response.create` return value + (an object exposing `.headers` and `.parse()` -> something with model_dump).""" + parsed = MagicMock() + parsed.model_dump.return_value = model_response.model_dump() + raw = MagicMock() + raw.headers = {} + raw.parse.return_value = parsed + return raw + + +# --------------------------------------------------------------------------- +# A) PROVIDER-PAYLOAD REGRESSION +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_internal_control_fields_never_leak_into_provider_body(restore_callbacks): + """Drive a real acompletion with a native code_interpreter tool through the + interception logger + agentic loop, capturing every outbound OpenAI request + body. None of the internal control fields may appear at top-level or inside + extra_body on ANY of the captured calls.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandboxConfig()) + litellm.callbacks = [logger] + + # First create -> model emits a code_execution tool call (triggers the loop). + # Second create -> model returns a plain answer (loop terminates). + create = AsyncMock( + side_effect=[ + _raw_response_for(_tool_call_model_response()), + _raw_response_for(_plain_model_response()), + ] + ) + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create = create + + response = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + tools=[{"type": "code_interpreter"}], + tool_choice={"type": "code_interpreter"}, + api_key="sk-test", + client=mock_client, + ) + + # The loop must have actually fired (sanity: two provider calls). + assert create.await_count == 2, ( + "expected the agentic loop to issue a follow-up provider call; " + f"got {create.await_count} call(s)" + ) + + for idx, call in enumerate(create.await_args_list): + body = call.kwargs + extra_body = body.get("extra_body") or {} + for field in _INTERNAL_CONTROL_FIELDS: + assert field not in body, ( + f"provider call #{idx}: internal field {field!r} leaked into " + f"top-level request body: {sorted(body.keys())}" + ) + assert field not in extra_body, ( + f"provider call #{idx}: internal field {field!r} leaked into " + f"extra_body: {sorted(extra_body.keys())}" + ) + # The native code_interpreter tool must have been swapped for the + # function tool, never sent raw to OpenAI as a chat-completions request. + for tool in body.get("tools") or []: + assert tool.get("type") != "code_interpreter" + + # The final response is the post-loop answer, not the tool-call turn. + assert response.choices[0].message.content == "The answer is 42" + + +# --------------------------------------------------------------------------- +# B) DISPATCHER UNIT TESTS +# --------------------------------------------------------------------------- + + +class _LoggingStub: + """Minimal logging_obj: dispatcher only reads dynamic_success_callbacks and + litellm_call_id off it.""" + + litellm_call_id = "call-test" + dynamic_success_callbacks: List[Any] = [] + + +class _GateOnlyLogger(CustomLogger): + """Overrides the gate to fire, but builds a plan from request_patch.""" + + def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None: + super().__init__() + self._plan = plan + self._tool_calls = tool_calls + self.cleanup_calls = 0 + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict[str, Any], + ) -> Tuple[bool, Dict[str, Any]]: + return True, self._tool_calls + + async def async_build_agentic_loop_plan( + self, + tools: Dict[str, Any], + model: str, + messages: List[Dict[str, Any]], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict[str, Any], + logging_obj: Any, + stream: bool, + kwargs: Dict[str, Any], + ) -> AgenticLoopPlan: + return self._plan + + async def async_agentic_loop_cleanup_hook( + self, plan: AgenticLoopPlan, kwargs: Dict[str, Any] + ) -> None: + self.cleanup_calls += 1 + + +def _patched_messages() -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "what is 6*7?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "litellm_code_execution", + "arguments": '{"code": "print(6*7)"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc", "content": "42\n"}, + ] + + +@pytest.mark.asyncio +async def test_dispatcher_returns_none_when_no_callback_gates(restore_callbacks): + """No callback overrides the gate -> dispatcher returns None so the caller + keeps the original response untouched.""" + litellm.callbacks = [] + + result = await maybe_run_chat_completion_agentic_loop( + response=_plain_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_dispatcher_runs_followup_with_incremented_depth_and_patched_messages( + restore_callbacks, +): + """A gating logger with a request_patch -> the dispatcher calls + litellm.acompletion exactly once with _agentic_loop_depth == 1 and the + patched messages. Loop-control state rides as litellm-level kwargs and is + mirrored into litellm_metadata; the provider-surface transient + _agentic_loop_api_surface is never forwarded. (Provider-body stripping of + these litellm-level kwargs is asserted separately in test A.)""" + followup = _plain_model_response("done") + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages()), + ) + logger = _GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]}) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock(return_value=followup) + with patch.object(litellm, "acompletion", acompletion_mock): + result = await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={"temperature": 0.1}, + kwargs={"_code_interpreter_interception_active": True}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + assert result is followup + acompletion_mock.assert_awaited_once() + call_kwargs = acompletion_mock.await_args.kwargs + + assert call_kwargs["_agentic_loop_depth"] == 1 + assert call_kwargs["messages"] == _patched_messages() + # Preserved non-internal optional param survives the rerun. + assert call_kwargs["temperature"] == 0.1 + # Loop-control state is carried at the litellm level for the follow-up. + assert call_kwargs["max_agentic_loops"] >= 1 + assert "_agentic_loop_fingerprints" in call_kwargs + # Interception markers are mirrored into litellm_metadata for the follow-up. + assert ( + call_kwargs["litellm_metadata"]["_code_interpreter_interception_active"] is True + ) + # The transient surface marker is NOT forwarded to the follow-up call. + assert "_agentic_loop_api_surface" not in call_kwargs + # Cleanup hook always runs. + assert logger.cleanup_calls == 1 + + +@pytest.mark.asyncio +async def test_dispatcher_raises_when_depth_reaches_max_agentic_loops( + restore_callbacks, +): + """depth >= max_agentic_loops -> ValueError mentioning max_agentic_loops, + before any follow-up call is attempted.""" + logger = _GateOnlyLogger( + plan=AgenticLoopPlan(run_agentic_loop=True), + tool_calls={"tool_calls": [{"id": "call_abc"}]}, + ) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + with pytest.raises(ValueError, match="max_agentic_loops"): + await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + acompletion_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callbacks): + """A tool_calls fingerprint already present in _agentic_loop_fingerprints -> + ValueError about the repeated fingerprint (cycle guard), with no follow-up + call.""" + import json + + # The dispatcher fingerprints the whole value the gate returns as its second + # tuple element, so the seeded fingerprint must mirror that dict exactly. + gate_tool_calls = { + "tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}] + } + fingerprint = json.dumps(gate_tool_calls, sort_keys=True, default=str) + + logger = _GateOnlyLogger( + plan=AgenticLoopPlan(run_agentic_loop=True), + tool_calls=gate_tool_calls, + ) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + with pytest.raises(ValueError, match="fingerprint"): + await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={ + "_agentic_loop_depth": 0, + "max_agentic_loops": 3, + "_agentic_loop_fingerprints": [fingerprint], + }, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + acompletion_mock.assert_not_awaited()