Merge 8e509bf6dc into 570f41daad

litellm/proxy/tracing_endpoints.py

@@ -185,6 +185,7 @@ async def list_agent_traces(
             start_ms=request.start_ms if request.start_ms is not None else now_ms - MS_PER_DAY,
             end_ms=request.end_ms if request.end_ms is not None else now_ms,
             cursor=request.cursor,
+            agent=request.agent,
         )
     except (TraceChanged, ValueError, OverflowError, RuntimeError) as error:
         raise read_failure(error) from error

litellm/tracing/generated/requests.py

@@ -35,6 +35,7 @@ class TraceListRequest(LiteLLMBaseModel):
 
     cursor: str | None = Field(None, max_length=512)
     end_ms: int | None = Field(None, description="Window end, unix ms. Default: now")
+    agent: str | None = Field(None, max_length=256)
     start_ms: int | None = Field(None, description="Window start, unix ms. Default: 24h ago")
 
 

litellm/tracing/receiver.py

@@ -12,8 +12,10 @@ class TraceReceiver:
     def __init__(self, storage: LensTraceStorage) -> None:
         self.storage: Final = storage
 
-    async def list_traces(self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None = None) -> TracePage:
-        return await self.storage.list_traces(scope, start_ms, end_ms, cursor, AGENT_TRACING_LIST_PAGE_SIZE)
+    async def list_traces(
+        self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None = None, agent: str | None = None
+    ) -> TracePage:
+        return await self.storage.list_traces(scope, start_ms, end_ms, cursor, AGENT_TRACING_LIST_PAGE_SIZE, agent)
 
     async def list_agents(self, scope: TraceScope, start_ms: int, end_ms: int) -> TraceAgentList:
         rows: Final = await self.storage.trace_agents(

litellm/tracing/remote.py

@@ -130,7 +130,7 @@ class RemoteTraceStore:
         response.raise_for_status()
 
     async def list_traces(
-        self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None, limit: int
+        self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None, limit: int, agent: str | None = None
     ) -> JsonValue:
         return await self._read(
             {
@@ -140,6 +140,7 @@ class RemoteTraceStore:
                 "end_ms": end_ms,
                 "cursor": cursor,
                 "limit": limit,
+                **({"agent": agent} if agent else {}),
             }
         )
 

litellm/tracing/storage.py

@@ -44,8 +44,9 @@ class LensTraceStorage:
         end_ms: int,
         cursor: str | None = None,
         limit: int = AGENT_TRACING_LIST_PAGE_SIZE,
+        agent: str | None = None,
     ) -> TracePage:
-        result: Final = await self._remote.list_traces(scope, start_ms, end_ms, cursor, limit)
+        result: Final = await self._remote.list_traces(scope, start_ms, end_ms, cursor, limit, agent)
         return _validate_query_response(_TRACE_PAGE, result)
 
     async def get_trace(

scripts/lens_assets/schemas/requests/TraceListRequest.json

@@ -18,6 +18,14 @@
         "null"
       ]
     },
+    "agent": {
+      "default": null,
+      "maxLength": 256,
+      "type": [
+        "string",
+        "null"
+      ]
+    },
     "start_ms": {
       "default": null,
       "description": "Window start, unix ms. Default: 24h ago",

scripts/lens_assets/source.json

@@ -14,7 +14,7 @@
     "requests": {
       "schemas/requests/TraceDetailRequest.json": "4666e0cae8439c8f1726ea5418db8d4dbd613836759c7d0f1578f92077f9ff70",
       "schemas/requests/TraceErrorPageRequest.json": "f5d8695d0a2b004fb3baca8d2566e35e8b33e3b3e81fd5ea8ce6c03f2274ea4f",
-      "schemas/requests/TraceListRequest.json": "5cdb08384256257ed481ce27a4989b13e5dabc56bf063457f6fa97c0ce04f007",
+      "schemas/requests/TraceListRequest.json": "fe2c44f27e9821426bd3f2b0245c1764c20b436f1f6144c7d5f62b9b27b9fa11",
       "schemas/requests/TraceQueryRequest.json": "dd2e2d25d4d54a2021b0585dbd8391a1e91334a959328b68b73ed1fea3303930",
       "schemas/requests/TraceSpanRequest.json": "7009bc131f190c0863ea6eb808966f1613d6a1611e1dd88a884953b7b19d3224"
     },

tests/unit/proxy/test_tracing_endpoints.py

@@ -179,7 +179,9 @@ def test_trace_read_permissions_with_retired_uploads(
     if scope is None:
         receiver.list_traces.assert_not_awaited()
     else:
-        receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, cursor=None)
+        receiver.list_traces.assert_awaited_once_with(
+            scope=scope, start_ms=1, end_ms=2, cursor=None, agent=None
+        )
 
     write: Final = client.post("/v1/traces", json={})
     assert write.status_code == 410
@@ -229,7 +231,7 @@ def test_gateway_uploads_return_setup_guidance_without_reading_the_body(
 
 
 def test_list_traces_passes_scope_window_and_cursor(client, receiver):
-    response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, "cursor": "abc"})
+    response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, "cursor": "abc", "agent": "agent-a"})
     assert response.status_code == 200
     assert response.json() == {"data": [], "next_cursor": None}
     receiver.list_traces.assert_awaited_once_with(
@@ -237,6 +239,7 @@ def test_list_traces_passes_scope_window_and_cursor(client, receiver):
         start_ms=1,
         end_ms=2,
         cursor="abc",
+        agent="agent-a",
     )
 
 
@@ -245,6 +248,7 @@ def test_list_traces_defaults_to_last_24h(client, receiver):
     kwargs = receiver.list_traces.call_args.kwargs
     assert kwargs["end_ms"] - kwargs["start_ms"] == tracing_endpoints.MS_PER_DAY
     assert kwargs["cursor"] is None
+    assert kwargs["agent"] is None
 
 
 @pytest.mark.parametrize(
@@ -270,6 +274,7 @@ def test_list_traces_resolves_default_bounds_from_injected_clock(
         start_ms=expected_start_ms,
         end_ms=expected_end_ms,
         cursor=None,
+        agent=None,
     )
 
 
@@ -453,6 +458,7 @@ def test_list_traces_forwards_large_and_negative_bounds_unchanged(client: TestCl
         start_ms=2**63,
         end_ms=-1,
         cursor="next",
+        agent=None,
     )
 
 
@@ -525,6 +531,7 @@ def test_trace_read_routes_accept_and_forward_512_character_cursors(client: Test
         start_ms=NOW_MS - tracing_endpoints.MS_PER_DAY,
         end_ms=NOW_MS,
         cursor=cursor,
+        agent=None,
     )
     receiver.get_trace.assert_awaited_once_with(
         "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor, None
@@ -570,8 +577,20 @@ def test_trace_read_routes_ignore_unknown_query_parameters(client: TestClient, r
     assert error_unknown_response.status_code == 200, error_unknown_response.text
     receiver.list_traces.assert_has_awaits(
         (
-            call(scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, cursor="list-cursor"),
-            call(scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, cursor="list-cursor"),
+            call(
+                scope={"all_teams": 0, "user_id": "user", "team_ids": ()},
+                start_ms=1,
+                end_ms=2,
+                cursor="list-cursor",
+                agent=None,
+            ),
+            call(
+                scope={"all_teams": 0, "user_id": "user", "team_ids": ()},
+                start_ms=1,
+                end_ms=2,
+                cursor="list-cursor",
+                agent=None,
+            ),
         )
     )
     receiver.get_trace.assert_has_awaits(

tests/unit/tracing/test_receiver.py

@@ -1,13 +1,17 @@
+import asyncio
+import json
 from datetime import datetime, timezone
 from typing import Final, cast
 
+import httpx
 import pytest
 
-from litellm.constants import AGENT_TRACING_AGENT_LIST_LIMIT
+from litellm.constants import AGENT_TRACING_AGENT_LIST_LIMIT, AGENT_TRACING_LIST_PAGE_SIZE
+from litellm.tracing import TraceReceiver
 from litellm.tracing.generated.models import TraceAgentRow, TraceAgentsParams
 from litellm.tracing.generated.types import TraceScope
+from litellm.tracing.remote import RemoteTraceStore
 from litellm.tracing.storage import LensTraceStorage
-from litellm.tracing import TraceReceiver
 from litellm.tracing.types import TraceAgent
 
 
@@ -69,3 +73,30 @@ async def test_list_agents_queries_the_reader_scope_and_shapes_rows(
             last_seen=datetime(1970, 1, 1, tzinfo=timezone.utc),
         ),
     )
+
+
+@pytest.mark.asyncio
+async def test_list_traces_forwards_agent_filter_to_lens() -> None:
+    requests: Final = asyncio.Queue[httpx.Request]()
+
+    def accept(request: httpx.Request) -> httpx.Response:
+        requests.put_nowait(request)
+        return httpx.Response(200, json={"data": [], "next_cursor": None})
+
+    scope: Final = TraceScope(all_teams=0, user_id="owner", team_ids=("team",))
+    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
+        receiver: Final = TraceReceiver(LensTraceStorage(RemoteTraceStore(client)))
+        page: Final = await receiver.list_traces(scope, start_ms=10, end_ms=20, agent="research-agent")
+
+    request: Final = requests.get_nowait()
+    assert request.url.path == "/internal/read"
+    assert json.loads(request.content) == {
+        "operation": "list",
+        "scope": {"all_teams": 0, "user_id": "owner", "team_ids": ["team"]},
+        "start_ms": 10,
+        "end_ms": 20,
+        "cursor": None,
+        "limit": AGENT_TRACING_LIST_PAGE_SIZE,
+        "agent": "research-agent",
+    }
+    assert page == {"data": (), "next_cursor": None}

tests/unit/tracing/test_remote.py

@@ -119,6 +119,32 @@ async def test_remote_reads_preserve_scope_and_pagination(operation: str) -> Non
         assert b"unused-local-secret" not in request.content
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("agent", "expected"), (("research-agent", "research-agent"), ("", None), (None, None)))
+async def test_remote_list_includes_only_a_non_empty_agent_filter(agent: str | None, expected: str | None) -> None:
+    requests: Final = asyncio.Queue[httpx.Request]()
+
+    def accept(request: httpx.Request) -> httpx.Response:
+        requests.put_nowait(request)
+        return httpx.Response(200, json={"data": [], "next_cursor": None})
+
+    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
+        await RemoteTraceStore(client).list_traces(
+            TraceScope(all_teams=1, user_id="", team_ids=()), 10, 20, None, 17, agent
+        )
+    request: Final = requests.get_nowait()
+    body: Final = json.loads(request.content)
+    assert body == {
+        "operation": "list",
+        "scope": {"all_teams": 1, "user_id": "", "team_ids": []},
+        "start_ms": 10,
+        "end_ms": 20,
+        "cursor": None,
+        "limit": 17,
+        **({"agent": expected} if expected else {}),
+    }
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "status,error",

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -84063,6 +84063,7 @@ export interface operations {
                 cursor?: string | null;
                 /** @description Window end, unix ms. Default: now */
                 end_ms?: number | null;
+                agent?: string | null;
                 /** @description Window start, unix ms. Default: 24h ago */
                 start_ms?: number | null;
             };