diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index bca14ac22b1..4287ca3fc77 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -88,72 +88,81 @@ _DATE_PARAMS = { "end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"}, } -ALL_TOOLS = [ - { - "type": "function", - "function": { - "name": "get_usage_data", - "description": ( - "Fetch aggregated global usage/spend data. Returns daily spend, " - "token counts, request counts, and breakdowns by model, provider, " - "and API key. Use for overall spend, top models, top providers." - ), - "parameters": { - "type": "object", - "properties": { - **_DATE_PARAMS, - "user_id": { - "type": "string", - "description": "Optional user ID filter. Omit for global view.", - }, +_TOOL_USAGE = { + "type": "function", + "function": { + "name": "get_usage_data", + "description": ( + "Fetch aggregated global usage/spend data. Returns daily spend, " + "token counts, request counts, and breakdowns by model, provider, " + "and API key. Use for overall spend, top models, top providers." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "user_id": { + "type": "string", + "description": "Optional user ID filter. Omit for global view.", }, - "required": ["start_date", "end_date"], }, + "required": ["start_date", "end_date"], }, }, - { - "type": "function", - "function": { - "name": "get_team_usage_data", - "description": ( - "Fetch usage/spend data broken down by team. Use for questions " - "like 'which team spends the most' or 'show me team X usage'." - ), - "parameters": { - "type": "object", - "properties": { - **_DATE_PARAMS, - "team_ids": { - "type": "string", - "description": "Optional comma-separated team IDs. Omit for all teams.", - }, +} + +_TOOL_TEAM = { + "type": "function", + "function": { + "name": "get_team_usage_data", + "description": ( + "Fetch usage/spend data broken down by team. Use for questions " + "like 'which team spends the most' or 'show me team X usage'." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "team_ids": { + "type": "string", + "description": "Optional comma-separated team IDs. Omit for all teams.", }, - "required": ["start_date", "end_date"], }, + "required": ["start_date", "end_date"], }, }, - { - "type": "function", - "function": { - "name": "get_tag_usage_data", - "description": ( - "Fetch usage/spend data broken down by tag. Tags are labels " - "attached to requests (features, environments, credentials)." - ), - "parameters": { - "type": "object", - "properties": { - **_DATE_PARAMS, - "tags": { - "type": "string", - "description": "Optional comma-separated tag names. Omit for all tags.", - }, +} + +_TOOL_TAG = { + "type": "function", + "function": { + "name": "get_tag_usage_data", + "description": ( + "Fetch usage/spend data broken down by tag. Tags are labels " + "attached to requests (features, environments, credentials)." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "tags": { + "type": "string", + "description": "Optional comma-separated tag names. Omit for all tags.", }, - "required": ["start_date", "end_date"], }, + "required": ["start_date", "end_date"], }, }, -] +} + +TOOLS_BASE = [_TOOL_USAGE] +TOOLS_ADMIN = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] + + +def get_tools_for_role(is_admin: bool) -> List[Dict[str, Any]]: + """Return the tool list appropriate for the user's role.""" + return TOOLS_ADMIN if is_admin else TOOLS_BASE + SYSTEM_PROMPT = ( "You are an AI assistant embedded in the LiteLLM Usage dashboard. " @@ -508,10 +517,11 @@ async def stream_usage_ai_chat( try: yield _sse({"type": "status", "message": "Thinking..."}) + tools = get_tools_for_role(is_admin) response = await litellm.acompletion( model=resolved_model, messages=chat_messages, - tools=ALL_TOOLS, + tools=tools, temperature=USAGE_AI_TEMPERATURE, ) choice = response.choices[0] # type: ignore diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py index 7aa07110c10..0dbe518afb7 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -4,7 +4,7 @@ USAGE AI CHAT ENDPOINTS /usage/ai/chat - Stream AI chat responses about usage data """ -from typing import List, Optional +from typing import List, Literal, Optional from fastapi import APIRouter, Depends, Request from fastapi.responses import StreamingResponse @@ -17,7 +17,7 @@ router = APIRouter() class ChatMessage(BaseModel): - role: str + role: Literal["user", "assistant"] content: str diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index af95b7a0edd..f08f7ee58ec 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -8,9 +8,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( - ALL_TOOLS, SYSTEM_PROMPT, TOOL_HANDLERS, + TOOLS_ADMIN, + TOOLS_BASE, _summarise_entity_data, _summarise_usage_data, stream_usage_ai_chat, @@ -111,13 +112,17 @@ SAMPLE_TEAM_RESPONSE = { class TestToolSchemas: - def test_all_tools_defined(self): - assert len(ALL_TOOLS) == 3 - names = {t["function"]["name"] for t in ALL_TOOLS} + def test_admin_tools_include_all(self): + assert len(TOOLS_ADMIN) == 3 + names = {t["function"]["name"] for t in TOOLS_ADMIN} assert "get_usage_data" in names assert "get_team_usage_data" in names assert "get_tag_usage_data" in names + def test_base_tools_restricted_to_usage_only(self): + assert len(TOOLS_BASE) == 1 + assert TOOLS_BASE[0]["function"]["name"] == "get_usage_data" + def test_system_prompt_mentions_all_tools(self): assert "get_usage_data" in SYSTEM_PROMPT assert "get_team_usage_data" in SYSTEM_PROMPT