diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index 8d0957ff2be..4830bfcafb0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -27,6 +27,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, @@ -372,16 +373,16 @@ class Agent365Guardrail(CustomGuardrail): ) if isinstance(session_id, str) and session_id: return session_id - logging_obj: Final = data.get("litellm_logging_obj") - call_details: Final = getattr(logging_obj, "model_call_details", None) - if isinstance(call_details, Mapping): - tool_call_metadata: Final = call_details.get("mcp_tool_call_metadata") + raw_logging_obj: Final = data.get("litellm_logging_obj") + logging_obj: Final = raw_logging_obj if isinstance(raw_logging_obj, LiteLLMLoggingObj) else None + if logging_obj is not None: + tool_call_metadata: Final = logging_obj.model_call_details.get("mcp_tool_call_metadata") session_from_logging: Final = ( tool_call_metadata.get("mcp_session_id") if isinstance(tool_call_metadata, Mapping) else None ) if isinstance(session_from_logging, str) and session_from_logging: return session_from_logging - call_id: Final = data.get("litellm_call_id") or getattr(logging_obj, "litellm_call_id", None) + call_id: Final = data.get("litellm_call_id") or (logging_obj.litellm_call_id if logging_obj else None) if isinstance(call_id, str) and call_id: return call_id return str(uuid.uuid4()) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py index 8245aba6d73..04460b150f5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -7,6 +7,7 @@ import pytest from fastapi import HTTPException from litellm.exceptions import Timeout as LitellmTimeout +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.agent_365 import ( @@ -295,10 +296,16 @@ class TestConversationId: async def test_falls_back_to_logging_obj_session_id(self): handler: Final = FakeHandler([_token_response(), _allow_response()]) guardrail: Final = _make_guardrail(handler) - logging_obj: Final = SimpleNamespace( - model_call_details={"mcp_tool_call_metadata": {"mcp_session_id": "sess-from-logging"}}, + logging_obj: Final = LiteLLMLoggingObj( + model="mcp", + messages=[], + stream=False, + call_type="call_mcp_tool", + start_time=None, litellm_call_id="call-id-1", + function_id="fn-1", ) + logging_obj.model_call_details["mcp_tool_call_metadata"] = {"mcp_session_id": "sess-from-logging"} data: Final = _mcp_data(metadata={"headers": {}}, litellm_logging_obj=logging_obj) await _run(guardrail, data) assert handler.calls[1].json["conversationId"] == "sess-from-logging"