mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(responses): don't require litellm_enterprise for background=true responses
This commit is contained in:
Devin AI
parent
bf02a4a47f
commit
170d348fb0
2 changed files with 192 additions and 40 deletions
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import Any, AsyncIterator, Dict, Optional, cast
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Optional, cast
|
||||
from uuid import uuid4
|
||||
|
||||
import fastapi
|
||||
|
|
@ -20,9 +20,66 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin
|
|||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
from litellm.types.responses.main import DeleteResponseResult
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
async def _store_background_response_in_managed_objects(
|
||||
response: ResponsesAPIResponse,
|
||||
proxy_logging_obj: "ProxyLogging",
|
||||
llm_router: Optional["Router"],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""
|
||||
Persist a queued/in_progress background Responses API result in the managed
|
||||
objects table so it can be polled later. Managed-object storage lives in the
|
||||
enterprise package, so this is a no-op on installs where it isn't available
|
||||
"""
|
||||
try:
|
||||
from litellm_enterprise.proxy.hooks.managed_files import (
|
||||
_PROXY_LiteLLMManagedFiles,
|
||||
)
|
||||
except ImportError:
|
||||
verbose_proxy_logger.debug(
|
||||
"litellm_enterprise not installed; skipping managed-object storage for background response %s",
|
||||
response.id,
|
||||
)
|
||||
return
|
||||
|
||||
managed_files_obj = cast(
|
||||
Optional[_PROXY_LiteLLMManagedFiles],
|
||||
proxy_logging_obj.get_proxy_hook("managed_files"),
|
||||
)
|
||||
if managed_files_obj is None or llm_router is None:
|
||||
return
|
||||
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
||||
model_id = hidden_params.get("model_id", None)
|
||||
if not model_id:
|
||||
verbose_proxy_logger.warning(
|
||||
f"No model_id found in response hidden params for response {response.id}, skipping managed object storage"
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
await managed_files_obj.store_unified_object_id(
|
||||
unified_object_id=response.id,
|
||||
file_object=response,
|
||||
litellm_parent_otel_span=None,
|
||||
model_object_id=response.id,
|
||||
file_purpose="response",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"Stored background response {response.id} in managed objects table with unified_id={response.id}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Failed to store background response in managed objects table: {str(e)}")
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/responses",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -213,45 +270,17 @@ async def responses_api(
|
|||
)
|
||||
|
||||
# Store in managed objects table if background mode is enabled
|
||||
if data.get("background") and isinstance(response, ResponsesAPIResponse):
|
||||
if response.status in ["queued", "in_progress"]:
|
||||
from litellm_enterprise.proxy.hooks.managed_files import ( # type: ignore
|
||||
_PROXY_LiteLLMManagedFiles,
|
||||
)
|
||||
|
||||
managed_files_obj = cast(
|
||||
Optional[_PROXY_LiteLLMManagedFiles],
|
||||
proxy_logging_obj.get_proxy_hook("managed_files"),
|
||||
)
|
||||
|
||||
if managed_files_obj and llm_router:
|
||||
try:
|
||||
# Get the actual deployment model_id from hidden params
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
||||
model_id = hidden_params.get("model_id", None)
|
||||
|
||||
if not model_id:
|
||||
verbose_proxy_logger.warning(
|
||||
f"No model_id found in response hidden params for response {response.id}, skipping managed object storage"
|
||||
)
|
||||
raise Exception("No model_id found in response hidden params")
|
||||
# Store in managed objects table
|
||||
await managed_files_obj.store_unified_object_id(
|
||||
unified_object_id=response.id,
|
||||
file_object=response,
|
||||
litellm_parent_otel_span=None,
|
||||
model_object_id=response.id,
|
||||
file_purpose="response",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Stored background response {response.id} in managed objects table with unified_id={response.id}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"Failed to store background response in managed objects table: {str(e)}"
|
||||
)
|
||||
if (
|
||||
data.get("background")
|
||||
and isinstance(response, ResponsesAPIResponse)
|
||||
and response.status in ["queued", "in_progress"]
|
||||
):
|
||||
await _store_background_response_in_managed_objects(
|
||||
response=response,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
return response
|
||||
except ModifyResponseException as e:
|
||||
|
|
|
|||
|
|
@ -711,3 +711,126 @@ class TestManagedResponsesSameProvider:
|
|||
call_kwargs: dict = {}
|
||||
handler._inject_credentials(call_kwargs, model="vertex_ai/gemini-2.0-flash")
|
||||
assert "custom_llm_provider" not in call_kwargs
|
||||
|
||||
|
||||
class TestStoreBackgroundResponseInManagedObjects:
|
||||
"""
|
||||
Regression for #32782: background=true Responses API requests must not fail
|
||||
with ModuleNotFoundError when the litellm_enterprise package is not installed.
|
||||
"""
|
||||
|
||||
def _response(self):
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
response = ResponsesAPIResponse(
|
||||
id="resp_bg123",
|
||||
created_at=1234567890,
|
||||
model="gpt-4o",
|
||||
object="response",
|
||||
status="queued",
|
||||
output=[],
|
||||
)
|
||||
response._hidden_params = {"model_id": "deployment-123"}
|
||||
return response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_enterprise_package_is_noop(self):
|
||||
import sys
|
||||
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_store_background_response_in_managed_objects,
|
||||
)
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.get_proxy_hook = MagicMock()
|
||||
|
||||
# Force the enterprise import to raise ImportError regardless of whether
|
||||
# the package happens to be installed in the test environment.
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"litellm_enterprise.proxy.hooks.managed_files": None},
|
||||
):
|
||||
await _store_background_response_in_managed_objects(
|
||||
response=self._response(),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
# Bailed out before ever looking up the managed_files hook.
|
||||
proxy_logging_obj.get_proxy_hook.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stores_object_when_enterprise_hook_present(self):
|
||||
import sys
|
||||
import types
|
||||
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_store_background_response_in_managed_objects,
|
||||
)
|
||||
|
||||
fake_module = types.ModuleType(
|
||||
"litellm_enterprise.proxy.hooks.managed_files"
|
||||
)
|
||||
fake_module._PROXY_LiteLLMManagedFiles = type("_FakeManagedFiles", (), {})
|
||||
|
||||
managed_files_obj = MagicMock()
|
||||
managed_files_obj.store_unified_object_id = AsyncMock()
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.get_proxy_hook = MagicMock(return_value=managed_files_obj)
|
||||
|
||||
response = self._response()
|
||||
user_api_key_dict = MagicMock()
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"litellm_enterprise.proxy.hooks.managed_files": fake_module},
|
||||
):
|
||||
await _store_background_response_in_managed_objects(
|
||||
response=response,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=MagicMock(),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
managed_files_obj.store_unified_object_id.assert_awaited_once()
|
||||
call_kwargs = managed_files_obj.store_unified_object_id.await_args.kwargs
|
||||
assert call_kwargs["unified_object_id"] == response.id
|
||||
assert call_kwargs["file_purpose"] == "response"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_when_no_model_id(self):
|
||||
import sys
|
||||
import types
|
||||
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_store_background_response_in_managed_objects,
|
||||
)
|
||||
|
||||
fake_module = types.ModuleType(
|
||||
"litellm_enterprise.proxy.hooks.managed_files"
|
||||
)
|
||||
fake_module._PROXY_LiteLLMManagedFiles = type("_FakeManagedFiles", (), {})
|
||||
|
||||
managed_files_obj = MagicMock()
|
||||
managed_files_obj.store_unified_object_id = AsyncMock()
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.get_proxy_hook = MagicMock(return_value=managed_files_obj)
|
||||
|
||||
response = self._response()
|
||||
response._hidden_params = {}
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"litellm_enterprise.proxy.hooks.managed_files": fake_module},
|
||||
):
|
||||
await _store_background_response_in_managed_objects(
|
||||
response=response,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
managed_files_obj.store_unified_object_id.assert_not_awaited()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue