Address greptile review: security fixes and input validation

- Restrict team/tag tools to admin-only users (non-admins only get get_usage_data)
- Constrain ChatMessage.role to Literal['user', 'assistant'] to prevent system prompt injection
- Add test for base tools restriction (non-admin gets 1 tool, admin gets 3)
- Issues 3 (unused imports) and 4 (inline datetime) were already fixed in prior commit

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

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

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
 
 

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