feat(jql): add direct trace scope

Author: JJwilkin

README.md

@@ -88,13 +88,14 @@ from judgeval import Judgeval
 from judgeval.jql import spans
 
 client = Judgeval(project_name="my-project")
-result = client.query(spans().rows(), session_ids=["session-123"])
+result = client.query(spans().rows(), trace_ids=["trace-123"])
 ```
 
-`session_ids` is outside the JQL query object. Judgment resolves the sessions
+`trace_ids` and `session_ids` are mutually exclusive options outside the JQL
+query object. Trace IDs narrow the query directly. Judgment resolves session IDs
 within the authenticated organization and project, then narrows every part of
-the query to their traces. If none resolve, the request fails instead of
-falling back to the whole project. The same option works with `present()` and
+the query to their traces. If no session resolves, the request fails instead of
+falling back to the whole project. Both options work with `present()` and
 `discover()`.
 
 ## Integrations

scripts/jql_contract/public-openapi.json

@@ -278,6 +278,18 @@
             "maximum": 10000,
             "type": "integer"
           },
+          "trace_ids": {
+            "minItems": 1,
+            "maxItems": 1000,
+            "description": "Traces to bind directly to the query within the authenticated organization and project.",
+            "type": "array",
+            "items": {
+              "minLength": 1,
+              "maxLength": 512,
+              "pattern": ".*\\S.*",
+              "type": "string"
+            }
+          },
           "session_ids": {
             "minItems": 1,
             "maxItems": 1000,

src/judgeval/judgeval.py

@@ -196,42 +196,51 @@ def query(
         query: "QueryInput",
         *,
         limit: Optional[int] = None,
+        trace_ids: Optional[Sequence[str]] = None,
         session_ids: Optional[Sequence[str]] = None,
     ) -> "JqlQueryResponse":
-        """Run JQL for this project, optionally narrowed by session IDs."""
+        """Run JQL for this project, optionally narrowed by trace or session IDs."""
         from judgeval.jql import to_json
 
         return cast(
             "JqlQueryResponse",
-            self._run_jql("query", to_json(query), limit, session_ids),
+            self._run_jql("query", to_json(query), limit, trace_ids, session_ids),
         )
 
     def present(
         self,
         query: "QueryInput",
         *,
         limit: Optional[int] = None,
+        trace_ids: Optional[Sequence[str]] = None,
         session_ids: Optional[Sequence[str]] = None,
     ) -> "JqlPresentationResponse":
-        """Run a chart or table JQL query, optionally narrowed by session."""
+        """Run a chart or table JQL query, optionally narrowed by trace or session IDs."""
         from judgeval.jql import to_json
 
         return cast(
             "JqlPresentationResponse",
-            self._run_jql("query/presentation", to_json(query), limit, session_ids),
+            self._run_jql(
+                "query/presentation", to_json(query), limit, trace_ids, session_ids
+            ),
         )
 
     def _run_jql(
         self,
         path: str,
         query: Dict[str, Any],
         limit: Optional[int],
+        trace_ids: Optional[Sequence[str]],
         session_ids: Optional[Sequence[str]],
     ) -> Any:
+        if trace_ids is not None and session_ids is not None:
+            raise ValueError("trace_ids and session_ids are mutually exclusive")
         project_id = self._require_jql_project_id()
         payload: Dict[str, Any] = {"query": query}
         if limit is not None:
             payload["limit"] = limit
+        if trace_ids is not None:
+            payload["trace_ids"] = list(trace_ids)
         if session_ids is not None:
             payload["session_ids"] = list(session_ids)
         try:
@@ -251,14 +260,18 @@ def discover(
         kind: "DiscoveryKind",
         *,
         limit: Optional[int] = None,
+        trace_ids: Optional[Sequence[str]] = None,
         session_ids: Optional[Sequence[str]] = None,
         **options: Any,
     ) -> "JqlQueryResponse":
         """Discover project-scoped judges, fields, models, and related values."""
         from judgeval.jql import discovery
 
         return self.query(
-            discovery(kind, **options), limit=limit, session_ids=session_ids
+            discovery(kind, **options),
+            limit=limit,
+            trace_ids=trace_ids,
+            session_ids=session_ids,
         )
 
     def _require_jql_project_id(self) -> str:

src/tests/jql/test_session_scope.py

@@ -7,8 +7,17 @@
 from judgeval.jql import spans
 
 
-def test_judgeval_query_sends_session_scope_outside_jql(
+@pytest.mark.parametrize(
+    ("scope", "expected_scope"),
+    [
+        ({"trace_ids": ["trace-1"]}, {"trace_ids": ["trace-1"]}),
+        ({"session_ids": ["session-1"]}, {"session_ids": ["session-1"]}),
+    ],
+)
+def test_judgeval_query_sends_scope_outside_jql(
     monkeypatch: pytest.MonkeyPatch,
+    scope: dict[str, list[str]],
+    expected_scope: dict[str, list[str]],
 ) -> None:
     monkeypatch.setattr("judgeval.judgeval.resolve_project_id", lambda *_: "project-1")
     calls = []
@@ -30,7 +39,7 @@ def fake_request(self, method, url, payload, params=None):  # type: ignore[no-un
         api_url="https://api.example.com/",
     )
 
-    response = client.query(spans().rows(), session_ids=["session-1"])
+    response = client.query(spans().rows(), **scope)  # type: ignore[arg-type]
 
     assert {"calls": calls, "response": response} == {
         "calls": [
@@ -43,7 +52,7 @@ def fake_request(self, method, url, payload, params=None):  # type: ignore[no-un
                         "source": "spans",
                         "select": {"op": "rows"},
                     },
-                    "session_ids": ["session-1"],
+                    **expected_scope,
                 },
                 None,
             )
@@ -55,3 +64,43 @@ def fake_request(self, method, url, payload, params=None):  # type: ignore[no-un
             "elapsed_ms": 4,
         },
     }
+
+
+def test_judgeval_query_rejects_trace_and_session_scope_together(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr("judgeval.judgeval.resolve_project_id", lambda *_: "project-1")
+    calls = []
+
+    def fake_request(self, method, url, payload, params=None):  # type: ignore[no-untyped-def]
+        calls.append((method, url, payload, params))
+        raise AssertionError("must not run")
+
+    monkeypatch.setattr(JudgmentSyncClient, "_request", fake_request)
+    client = Judgeval(
+        project_name="demo",
+        api_key="api-key",
+        organization_id="org-1",
+        api_url="https://api.example.com/",
+    )
+
+    with pytest.raises(ValueError) as caught:
+        client.query(
+            spans().rows(),
+            trace_ids=["trace-1"],
+            session_ids=["session-1"],
+        )
+
+    assert {
+        "error": {
+            "type": type(caught.value).__name__,
+            "message": str(caught.value),
+        },
+        "calls": calls,
+    } == {
+        "error": {
+            "type": "ValueError",
+            "message": "trace_ids and session_ids are mutually exclusive",
+        },
+        "calls": [],
+    }