scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Add agent_id column to Tool Policies table

Browse source 
- 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 <krrishdholakia@gmail.com>
This commit is contained in:
Cursor Agent 2026-03-05 03:48:33 +00:00
parent 76f73de3f0
commit 67b6ca164e
7 changed files with 22 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -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):

2
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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,
)
)

4
litellm/proxy/db/tool_registry_writer.py
View file

@ -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": {

1
litellm/proxy/schema.prisma
View file

@ -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?

1
litellm/types/tool_management.py
View file

@ -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

15
ui/litellm-dashboard/src/components/ToolPolicies.tsx
View file

@ -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<ToolPoliciesProps> = ({ 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<ToolPoliciesProps> = ({ accessToken, onSelec
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Key Name" field="key_alias" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Agent ID" field="agent_id" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">User Agent</TableHeaderCell>
</TableRow>
</TableHead>
<TableBody>
{loading ? (
<TableRow>
<TableCell colSpan={9} className="h-8 text-center text-gray-500">
<TableCell colSpan={10} className="h-8 text-center text-gray-500">
Loading tools…
</TableCell>
</TableRow>
) : paginated.length === 0 ? (
<TableRow>
<TableCell colSpan={9} className="h-8 text-center text-gray-500">
<TableCell colSpan={10} className="h-8 text-center text-gray-500">
No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery.
</TableCell>
</TableRow>
@ -520,6 +524,11 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken, onSelec
<span className="max-w-[15ch] truncate block">{tool.key_alias ?? "-"}</span>
</Tooltip>
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<Tooltip title={tool.agent_id ?? "-"}>
<span className="font-mono max-w-[15ch] truncate block text-xs text-gray-500">{tool.agent_id ?? "-"}</span>
</Tooltip>
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<Tooltip title={tool.user_agent ?? "-"}>
<span className="font-mono max-w-[20ch] truncate block text-xs text-gray-500">{tool.user_agent ?? "-"}</span>

1
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -10157,6 +10157,7 @@ export interface ToolRow {
created_by?: string;
updated_by?: string;
user_agent?: string;
agent_id?: string;
last_used_at?: string;
}