fix: route usage AI chat through proxy router

litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py

@@ -509,15 +509,56 @@ async def _process_tool_call(
     )
 
 
+def _get_proxy_router():
+    """Return the global proxy router when available."""
+    try:
+        from litellm.proxy.proxy_server import llm_router
+    except Exception:
+        return None
+    return llm_router
+
+
+async def _usage_ai_acompletion(
+    *,
+    model: str,
+    messages: List[Dict[str, Any]],
+    team_id: Optional[str] = None,
+    **kwargs: Any,
+):
+    """
+    Route Usage AI requests through the proxy router when available so proxy
+    aliases and model groups resolve consistently with normal proxy traffic.
+    """
+    router = _get_proxy_router()
+    request_kwargs: Dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        **kwargs,
+    }
+
+    if team_id:
+        metadata = dict(cast(Dict[str, Any], request_kwargs.get("metadata") or {}))
+        metadata["user_api_key_team_id"] = team_id
+        request_kwargs["metadata"] = metadata
+
+    if router is not None:
+        return await router.acompletion(**request_kwargs)
+
+    return await litellm.acompletion(**request_kwargs)
+
+
 async def _stream_final_response(
-    model: str, chat_messages: List[Dict[str, Any]]
+    model: str,
+    chat_messages: List[Dict[str, Any]],
+    team_id: Optional[str] = None,
 ) -> AsyncIterator[str]:
     """Stream the final LLM response after tool results are appended."""
     yield _sse({"type": "status", "message": "Analyzing results..."})
 
-    response = await litellm.acompletion(
+    response = await _usage_ai_acompletion(
         model=model,
         messages=chat_messages,
+        team_id=team_id,
         stream=True,
         temperature=USAGE_AI_TEMPERATURE,
     )
@@ -532,6 +573,7 @@ async def stream_usage_ai_chat(
     model: Optional[str] = None,
     user_id: Optional[str] = None,
     is_admin: bool = False,
+    team_id: Optional[str] = None,
 ) -> AsyncIterator[str]:
     """Stream SSE events: status → tool_call → chunk → done."""
     resolved_model = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL
@@ -546,9 +588,10 @@ async def stream_usage_ai_chat(
     try:
         yield _sse({"type": "status", "message": "Thinking..."})
         tools = get_tools_for_role(is_admin)
-        response = await litellm.acompletion(
+        response = await _usage_ai_acompletion(
             model=resolved_model,
             messages=chat_messages,
+            team_id=team_id,
             tools=tools,
             temperature=USAGE_AI_TEMPERATURE,
         )
@@ -564,7 +607,9 @@ async def stream_usage_ai_chat(
         for tc in choice.message.tool_calls:
             async for event in _process_tool_call(tc, chat_messages, user_id, is_admin):
                 yield event
-        async for event in _stream_final_response(resolved_model, chat_messages):
+        async for event in _stream_final_response(
+            resolved_model, chat_messages, team_id
+        ):
             yield event
         yield _sse({"type": "done"})
 

litellm/proxy/management_endpoints/usage_endpoints/endpoints.py

@@ -59,6 +59,7 @@ async def usage_ai_chat(
             model=data.model,
             user_id=user_id,
             is_admin=is_admin,
+            team_id=user_api_key_dict.team_id,
         ),
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},

tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py

@@ -177,11 +177,13 @@ class TestSummariseEntityData:
 
 
 class TestStreamUsageAiChat:
-    @pytest.mark.asyncio
-    async def test_stream_emits_status_events(self):
+    @staticmethod
+    def _make_tool_call_response(
+        tool_name: str = "get_usage_data",
+    ) -> tuple[MagicMock, MagicMock]:
         mock_tool_call = MagicMock()
         mock_tool_call.id = "call_123"
-        mock_tool_call.function.name = "get_usage_data"
+        mock_tool_call.function.name = tool_name
         mock_tool_call.function.arguments = json.dumps(
             {
                 "start_date": "2025-01-01",
@@ -200,18 +202,24 @@ class TestStreamUsageAiChat:
                     "id": "call_123",
                     "type": "function",
                     "function": {
-                        "name": "get_usage_data",
+                        "name": tool_name,
                         "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
                     },
                 }
             ],
         }
+        return mock_tool_call, mock_first_response
 
-        async def mock_stream():
-            chunk = MagicMock()
-            chunk.choices = [MagicMock()]
-            chunk.choices[0].delta.content = "Total spend is $50.25"
-            yield chunk
+    @staticmethod
+    async def _mock_stream(content: str = "Total spend is $50.25"):
+        chunk = MagicMock()
+        chunk.choices = [MagicMock()]
+        chunk.choices[0].delta.content = content
+        yield chunk
+
+    @pytest.mark.asyncio
+    async def test_stream_emits_status_events(self):
+        _, mock_first_response = self._make_tool_call_response()
 
         with patch(
             "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
@@ -222,7 +230,7 @@ class TestStreamUsageAiChat:
             mock_litellm.acompletion = AsyncMock(
                 side_effect=[
                     mock_first_response,
-                    mock_stream(),
+                    self._mock_stream(),
                 ]
             )
             mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE
@@ -251,39 +259,13 @@ class TestStreamUsageAiChat:
 
     @pytest.mark.asyncio
     async def test_stream_handles_team_tool(self):
-        mock_tool_call = MagicMock()
-        mock_tool_call.id = "call_team"
-        mock_tool_call.function.name = "get_team_usage_data"
-        mock_tool_call.function.arguments = json.dumps(
-            {
-                "start_date": "2025-01-01",
-                "end_date": "2025-01-31",
-            }
+        _, mock_first_response = self._make_tool_call_response(
+            tool_name="get_team_usage_data"
         )
-
-        mock_first_response = MagicMock()
-        mock_first_response.choices = [MagicMock()]
-        mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
-        mock_first_response.choices[0].message.model_dump.return_value = {
-            "role": "assistant",
-            "content": None,
-            "tool_calls": [
-                {
-                    "id": "call_team",
-                    "type": "function",
-                    "function": {
-                        "name": "get_team_usage_data",
-                        "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
-                    },
-                }
-            ],
-        }
-
-        async def mock_stream():
-            chunk = MagicMock()
-            chunk.choices = [MagicMock()]
-            chunk.choices[0].delta.content = "Engineering is the top team."
-            yield chunk
+        mock_first_response.choices[0].message.tool_calls[0].id = "call_team"
+        mock_first_response.choices[0].message.model_dump.return_value["tool_calls"][0][
+            "id"
+        ] = "call_team"
 
         with patch(
             "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
@@ -294,7 +276,7 @@ class TestStreamUsageAiChat:
             mock_litellm.acompletion = AsyncMock(
                 side_effect=[
                     mock_first_response,
-                    mock_stream(),
+                    self._mock_stream(content="Engineering is the top team."),
                 ]
             )
             mock_fetch.return_value = SAMPLE_TEAM_RESPONSE
@@ -330,41 +312,26 @@ class TestStreamUsageAiChat:
 
     @pytest.mark.asyncio
     async def test_non_admin_enforces_user_id(self):
-        mock_tool_call = MagicMock()
-        mock_tool_call.id = "call_456"
-        mock_tool_call.function.name = "get_usage_data"
-        mock_tool_call.function.arguments = json.dumps(
+        _, mock_first_response = self._make_tool_call_response()
+        mock_first_response.choices[0].message.tool_calls[0].id = "call_456"
+        mock_first_response.choices[0].message.tool_calls[
+            0
+        ].function.arguments = json.dumps(
             {
                 "start_date": "2025-01-01",
                 "end_date": "2025-01-31",
                 "user_id": "other-user",
             }
         )
-
-        mock_first_response = MagicMock()
-        mock_first_response.choices = [MagicMock()]
-        mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
-        mock_first_response.choices[0].message.model_dump.return_value = {
-            "role": "assistant",
-            "content": None,
-            "tool_calls": [
-                {
-                    "id": "call_456",
-                    "type": "function",
-                    "function": {
-                        "name": "get_usage_data",
-                        "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}',
-                    },
-                }
-            ],
+        mock_first_response.choices[0].message.model_dump.return_value["tool_calls"][0] = {
+            "id": "call_456",
+            "type": "function",
+            "function": {
+                "name": "get_usage_data",
+                "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}',
+            },
         }
 
-        async def mock_stream():
-            chunk = MagicMock()
-            chunk.choices = [MagicMock()]
-            chunk.choices[0].delta.content = "Data."
-            yield chunk
-
         mock_fetch = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE)
 
         with patch(
@@ -382,7 +349,7 @@ class TestStreamUsageAiChat:
             mock_litellm.acompletion = AsyncMock(
                 side_effect=[
                     mock_first_response,
-                    mock_stream(),
+                    self._mock_stream(content="Data."),
                 ]
             )
 
@@ -400,3 +367,79 @@ class TestStreamUsageAiChat:
                 end_date="2025-01-31",
                 user_id="my-user-id",
             )
+
+    @pytest.mark.asyncio
+    async def test_stream_uses_router_acompletion_for_proxy_alias_models(self):
+        _, mock_first_response = self._make_tool_call_response()
+        mock_router = MagicMock()
+        mock_router.acompletion = AsyncMock(
+            side_effect=[
+                mock_first_response,
+                self._mock_stream(),
+            ]
+        )
+
+        with patch(
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._get_proxy_router",
+            return_value=mock_router,
+        ), patch(
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
+        ) as mock_litellm, patch(
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data",
+            new_callable=AsyncMock,
+        ) as mock_fetch:
+            mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE
+
+            events = []
+            async for event in stream_usage_ai_chat(
+                messages=[{"role": "user", "content": "What is my total spend?"}],
+                model="mylitellmmodel",
+                team_id="team-123",
+                is_admin=True,
+            ):
+                events.append(json.loads(event.replace("data: ", "").strip()))
+
+            assert any(event["type"] == "done" for event in events)
+            assert mock_router.acompletion.await_count == 2
+            mock_litellm.acompletion.assert_not_called()
+
+            first_call = mock_router.acompletion.await_args_list[0].kwargs
+            second_call = mock_router.acompletion.await_args_list[1].kwargs
+
+            assert first_call["model"] == "mylitellmmodel"
+            assert first_call["metadata"]["user_api_key_team_id"] == "team-123"
+            assert second_call["metadata"]["user_api_key_team_id"] == "team-123"
+            assert second_call["stream"] is True
+
+    @pytest.mark.asyncio
+    async def test_stream_falls_back_to_litellm_when_router_missing(self):
+        _, mock_first_response = self._make_tool_call_response()
+
+        with patch(
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._get_proxy_router",
+            return_value=None,
+        ), patch(
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
+        ) as mock_litellm, patch(
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data",
+            new_callable=AsyncMock,
+        ) as mock_fetch:
+            mock_litellm.acompletion = AsyncMock(
+                side_effect=[
+                    mock_first_response,
+                    self._mock_stream(),
+                ]
+            )
+            mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE
+
+            events = []
+            async for event in stream_usage_ai_chat(
+                messages=[{"role": "user", "content": "What is my total spend?"}],
+                model="gpt-4o-mini",
+                team_id="team-123",
+                is_admin=True,
+            ):
+                events.append(json.loads(event.replace("data: ", "").strip()))
+
+            assert any(event["type"] == "done" for event in events)
+            assert mock_litellm.acompletion.await_count == 2

ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx

@@ -1,6 +1,7 @@
-import { screen } from "@testing-library/react";
-import { beforeAll, describe, expect, it, vi } from "vitest";
+import { act, fireEvent, screen, waitFor } from "@testing-library/react";
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { renderWithProviders } from "../../../../tests/test-utils";
+import { modelHubCall, usageAiChatStream } from "../../networking";
 import UsageAIChatPanel from "./UsageAIChatPanel";
 
 beforeAll(() => {
@@ -20,7 +21,7 @@ vi.mock("../../networking", () => ({
       { model_group: "claude-3-opus" },
     ],
   }),
-  usageAiChatStream: vi.fn(),
+  usageAiChatStream: vi.fn().mockResolvedValue(undefined),
 }));
 
 const defaultProps = {
@@ -30,6 +31,10 @@ const defaultProps = {
 };
 
 describe("UsageAIChatPanel", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
   it("should render the panel when open", () => {
     renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
 
@@ -82,4 +87,53 @@ describe("UsageAIChatPanel", () => {
     expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full");
     expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0");
   });
+
+  it("should submit the selected model value unchanged", async () => {
+    renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(modelHubCall).toHaveBeenCalledWith("test-token");
+    });
+
+    const modelSelect = screen.getByRole("combobox");
+    await act(async () => {
+      fireEvent.mouseDown(modelSelect);
+    });
+
+    await waitFor(() => {
+      expect(
+        screen.getByRole("option", { name: "claude-3-opus" })
+      ).toBeInTheDocument();
+    });
+
+    await act(async () => {
+      fireEvent.click(screen.getByRole("option", { name: "claude-3-opus" }));
+    });
+
+    await act(async () => {
+      fireEvent.change(screen.getByPlaceholderText("Ask about your usage..."), {
+        target: { value: "hello" },
+      });
+    });
+
+    await act(async () => {
+      fireEvent.click(screen.getByRole("button", { name: "Send" }));
+    });
+
+    await waitFor(() => {
+      expect(usageAiChatStream).toHaveBeenCalled();
+    });
+
+    expect(usageAiChatStream).toHaveBeenCalledWith(
+      "test-token",
+      [{ role: "user", content: "hello" }],
+      "claude-3-opus",
+      expect.any(Function),
+      expect.any(Function),
+      expect.any(Function),
+      expect.any(Function),
+      expect.any(Function),
+      expect.any(AbortSignal),
+    );
+  });
 });