From 67b6ca164ea88d7fffbadfcacbcd9123c038200c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 5 Mar 2026 03:48:33 +0000 Subject: [PATCH] Add agent_id column to Tool Policies table - Add agent_id field to LiteLLM_ToolTable in Prisma schema - Add agent_id to ToolDiscoveryQueueItem TypedDict - Add agent_id to LiteLLM_ToolTableRow Pydantic model - Update tool_registry_writer to persist agent_id on create - Extract agent_id from kwargs metadata in _enqueue_tool_registry_upsert - Add agent_id to frontend ToolRow interface - Add sortable Agent ID column to ToolPolicies.tsx table - Include agent_id in search filter Co-authored-by: Krish Dholakia --- litellm/proxy/_types.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 2 ++ litellm/proxy/db/tool_registry_writer.py | 4 ++++ litellm/proxy/schema.prisma | 1 + litellm/types/tool_management.py | 1 + .../src/components/ToolPolicies.tsx | 15 ++++++++++++--- .../src/components/networking.tsx | 1 + 7 files changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9acc80b0ddd..4aaea1b0200 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4215,6 +4215,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): team_id: Optional[str] # team that triggered discovery key_alias: Optional[str] # human-readable key alias user_agent: Optional[str] # HTTP User-Agent of the caller + agent_id: Optional[str] # agent that triggered this tool call class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4c96e079c9e..23c4871e615 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -215,6 +215,7 @@ class DBSpendUpdateWriter: _metadata = _litellm_params.get("metadata") or {} key_alias = _metadata.get("user_api_key_alias") or None user_agent = _metadata.get("user_agent") or None + agent_id: Optional[str] = kwargs.get("agent_id") or _metadata.get("agent_id") or None def _enqueue(tool_name: str, origin: str = "user_defined") -> None: self.tool_discovery_queue.add_update( @@ -225,6 +226,7 @@ class DBSpendUpdateWriter: team_id=team_id, key_alias=key_alias, user_agent=user_agent, + agent_id=agent_id, ) ) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 0eda012d515..4b1ecdea903 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -40,6 +40,7 @@ def _row_to_model(row: Union[dict, Any]) -> LiteLLM_ToolTableRow: "team_id", "key_alias", "user_agent", + "agent_id", "last_used_at", "created_at", "updated_at", @@ -59,6 +60,7 @@ def _row_to_model(row: Union[dict, Any]) -> LiteLLM_ToolTableRow: team_id=row.get("team_id"), key_alias=row.get("key_alias"), user_agent=row.get("user_agent"), + agent_id=row.get("agent_id"), last_used_at=row.get("last_used_at"), created_at=row.get("created_at"), updated_at=row.get("updated_at"), @@ -93,6 +95,7 @@ async def batch_upsert_tools( team_id = item.get("team_id") key_alias = item.get("key_alias") user_agent = item.get("user_agent") + agent_id = item.get("agent_id") await table.upsert( where={"tool_name": tool_name}, data={ @@ -109,6 +112,7 @@ async def batch_upsert_tools( "team_id": team_id, "key_alias": key_alias, "user_agent": user_agent, + "agent_id": agent_id, "last_used_at": now, }, "update": { diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9e847993efc..53d3fbfa00b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1109,6 +1109,7 @@ model LiteLLM_ToolTable { team_id String? // team that first called this tool key_alias String? // human-readable alias of the virtual key user_agent String? // user-agent of the first request that discovered this tool + agent_id String? // agent that first triggered this tool call last_used_at DateTime? // timestamp of the most recent call created_at DateTime @default(now()) created_by String? diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 1c5e1df9e9a..0ba8e54e745 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -25,6 +25,7 @@ class LiteLLM_ToolTableRow(BaseModel): team_id: Optional[str] = None key_alias: Optional[str] = None user_agent: Optional[str] = None + agent_id: Optional[str] = None last_used_at: Optional[datetime] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index cccbf79ddfa..928f21ba09b 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -40,7 +40,7 @@ function getTrendSubtitle(newToday: number, newYesterday: number): string | unde return `${diff} since yesterday`; } -type SortField = "tool_name" | "input_policy" | "output_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; +type SortField = "tool_name" | "input_policy" | "output_policy" | "team_id" | "key_alias" | "created_at" | "call_count" | "agent_id"; interface FilterValues { [key: string]: string; @@ -226,6 +226,7 @@ export const ToolPolicies: React.FC = ({ accessToken, onSelec (t.team_id ?? "").toLowerCase().includes(q) || (t.key_alias ?? "").toLowerCase().includes(q) || (t.key_hash ?? "").toLowerCase().includes(q) || + (t.agent_id ?? "").toLowerCase().includes(q) || t.input_policy.toLowerCase().includes(q) || t.output_policy.toLowerCase().includes(q); if (!matchesSearch) return false; @@ -447,19 +448,22 @@ export const ToolPolicies: React.FC = ({ accessToken, onSelec + + + User Agent {loading ? ( - + Loading tools… ) : paginated.length === 0 ? ( - + No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery. @@ -520,6 +524,11 @@ export const ToolPolicies: React.FC = ({ accessToken, onSelec {tool.key_alias ?? "-"} + + + {tool.agent_id ?? "-"} + + {tool.user_agent ?? "-"} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 9df2a3401e0..df3d01f86b5 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -10157,6 +10157,7 @@ export interface ToolRow { created_by?: string; updated_by?: string; user_agent?: string; + agent_id?: string; last_used_at?: string; }