From db5b1f86b381badac1d5e7502b3e9dadcb6b18c5 Mon Sep 17 00:00:00 2001 From: Radu Raicea Date: Tue, 20 Oct 2026 15:48:14 -0400 Subject: [PATCH 1/3] feat(llma): add datasets to offline evaluations --- .sampo/changesets/ai-evaluation-datasets.md | 5 + .../example-offline-evaluations/README.md | 163 +++- examples/example-offline-evaluations/main.py | 93 ++- posthog/ai/evaluations/__init__.py | 24 +- posthog/ai/evaluations/_async_datasets.py | 460 +++++++++++ posthog/ai/evaluations/_client.py | 9 +- .../ai/evaluations/_dataset_serialization.py | 209 +++++ posthog/ai/evaluations/_dataset_types.py | 111 +++ posthog/ai/evaluations/_datasets.py | 453 ++++++++++ posthog/ai/evaluations/_types.py | 40 + posthog/test/ai/test_evaluations_datasets.py | 771 ++++++++++++++++++ references/public_api_snapshot.txt | 100 +++ 12 files changed, 2414 insertions(+), 24 deletions(-) create mode 100644 .sampo/changesets/ai-evaluation-datasets.md create mode 100644 posthog/ai/evaluations/_async_datasets.py create mode 100644 posthog/ai/evaluations/_dataset_serialization.py create mode 100644 posthog/ai/evaluations/_dataset_types.py create mode 100644 posthog/ai/evaluations/_datasets.py create mode 100644 posthog/test/ai/test_evaluations_datasets.py diff --git a/.sampo/changesets/ai-evaluation-datasets.md b/.sampo/changesets/ai-evaluation-datasets.md new file mode 100644 index 000000000..22d689b31 --- /dev/null +++ b/.sampo/changesets/ai-evaluation-datasets.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +Manage AI observability datasets and items through `OfflineEvaluations.datasets` and `AsyncOfflineEvaluations.datasets`, including history, tags, and archive/restore operations. Pin a dataset snapshot for lazy, consistent reads and use `EvaluationItem.from_dataset_item()` to preserve case provenance across offline experiments and scorers. Item creates with a caller-supplied `client_item_id` support safe retries; updates use explicit version checks. diff --git a/examples/example-offline-evaluations/README.md b/examples/example-offline-evaluations/README.md index 279fe4e70..7da7dd1c8 100644 --- a/examples/example-offline-evaluations/README.md +++ b/examples/example-offline-evaluations/README.md @@ -1,6 +1,6 @@ # Offline evaluations -Upload evaluations you have already run, manage reusable scorers, and explicitly finish an experiment. The SDK returns server acknowledgments for every operation. It does not execute your application or judges, create traces, or upload in the background. +Upload evaluations you have already run, manage reusable scorers and datasets, and explicitly finish an experiment. The SDK returns server acknowledgments for every operation. It does not execute your application or judges, create traces, or upload in the background. ## Credentials and setup @@ -11,11 +11,15 @@ Offline evaluations must be enabled for your PostHog project. Use the project's | Create experiments, upload results, complete/fail | Personal API key or project secret key with `offline_evaluation_ingestion:write` | | Fetch scorers and versions | Personal API key with `llm_analytics:read` | | Create/update scorers and create versions | Personal API key with `llm_analytics:write` | +| Fetch datasets, items, and history | Personal API key with `dataset:read` | +| Create/update/archive/restore datasets and items | Personal API key with `dataset:write` | | Read experiment tags | Personal API key with `evaluation:read` | | Replace experiment tags | Personal API key with `evaluation:write` | A personal key can cover both scorer management and uploads. Project secret keys can upload using a known scorer-version UUID without scorer read access. All operations use the explicit project/environment; scorer names are not unique. +Dataset operations also require the datasets feature to be enabled. Dataset reads, scorer reads, and experiment uploads use their own scopes; an upload-only credential cannot fetch a dataset. Keep personal keys in server-side scripts or CI secrets. + From this repository's root: ```bash @@ -85,6 +89,163 @@ EvaluationItem(input=None) # Payload contains a null input. The client context manager releases connections. Call `complete()` explicitly after uploads succeed, or `fail()` when you deliberately want to close the experiment as failed. Neither exceptions nor async cancellation automatically close an experiment. Expected counts include every outcome status; the server checks them on completion. +## Datasets and reproducible evaluations + +Use `evaluations.datasets` to manage PostHog-hosted datasets with the same client configuration as scorers and experiments. Dataset names are unique within a project; keep the returned UUID for subsequent operations. + +Provision a dataset once with a personal key that has `dataset:write`: + +```python +dataset = evaluations.datasets.create( + name="Capital city regression cases", + description="Cases for the offline evaluation example", + tags=["regression"], +) +case = evaluations.datasets.create_item( + dataset.id, + client_item_id="capital-france", + input="Capital of France?", + expected_output="Paris", + metadata={"category": "geography"}, +) +print(dataset.id) +``` + +An identical item create with the same `client_item_id` returns the existing active item. Reusing that key with different current content, or for an archived item, raises `EvaluationAPIError` with `code="client_item_id_conflict"`. Choose a stable key per case and retain it across process restarts. Creating the dataset itself is not idempotent; keep its ID instead of provisioning it on each evaluation run. + +Pin both the dataset and scorer before running your application: + +```python +from posthog.ai.evaluations import EvaluationItem, OfflineEvaluations + +def capital_answer(question): + return {"Capital of France?": "Paris", "Capital of Italy?": "Rome"}.get( + question, "I don't know" + ) + +with OfflineEvaluations(project_id=123, secret_key=personal_api_key) as evaluations: + snapshot = evaluations.datasets.snapshot(dataset_id) # Resolve latest once. + scorer = evaluations.scorers.get(scorer_id) # A boolean exact-match scorer. + version_id = scorer.current_version_id + experiment = evaluations.create_experiment( + name="Capital city regression", + dataset_revision_id=snapshot.revision_id, + application_version="your-commit-sha", + run_source="ci", + ) + + for case in snapshot.iter_items(): + output = capital_answer(case.input) + item = EvaluationItem.from_dataset_item(case, output=output) + if case.expected_output is None: + experiment.upload_result( + item=item, scorer_version_id=version_id, status="not_applicable", + ) + else: + experiment.upload_result( + item=item, scorer_version_id=version_id, + value=output == case.expected_output, + ) + + experiment.complete() +``` + +This example uses string inputs. Replace `capital_answer` and the equality check with your own application and judge. The SDK does not run scorer definitions. For multiple scorers, reuse the same `EvaluationItem` and upload one result per scorer-version UUID. For independent trials, call `from_dataset_item` again and pass a distinct `trial`, such as `trial="2"`; each call generates a new execution UUID. The helper defaults `case_key` to `client_item_id`, falling back to the stable dataset item UUID, and accepts an explicit override. + +`from_dataset_item` copies input, expected output, metadata, and immutable item provenance. It never uses `source_output` as a new execution's output or expected answer, and never uses the original source trace as the new application's trace. Pass `application_trace_id` explicitly for the new execution. Omit `output` if execution failed before producing one, then upload an error outcome. Omitted output differs from `output=None`. + +### Snapshot and pagination behavior + +- `snapshot(dataset_id, revision=12)` selects a historical **number**. The returned `revision_id` is the **UUID** to pass to `create_experiment`. The backend derives the remaining hosted-dataset identifiers from that UUID. +- Every item request uses the pinned revision, even if another process edits or archives items during iteration. `DatasetItem.dataset_revision_id` identifies the revision that introduced that item version; it can differ from the snapshot's revision ID. Do not use it as the experiment's revision. +- `snapshot()` fetches metadata only. `snapshot.iter_items(page_size=25)` lazily fetches active items, one page at a time. Keep the client open. Repeating iteration starts again at the first page without caching; read errors propagate rather than silently returning a partial dataset. +- A new dataset has no revision until its first item is added, so `snapshot()` raises `ValueError`. An existing revision with zero active items yields an empty iterator. Archived datasets remain readable. +- `snapshot.dataset` contains current descriptive metadata observed when the snapshot was resolved. Names, descriptions, tags, and dataset metadata are not historical item content. + +For explicit paging, each `list*` method returns typed `results`, `count`, and `next_offset`. Continue with the same filters until `next_offset is None`. Dataset/revision pages allow at most 100 entries, item pages 25, and item-history pages 100. `list_items` without a revision reads live state on each request; use a snapshot for an evaluation run. A historical snapshot may make multiple metadata requests to find its revision UUID. + +| Method | Purpose | +| --- | --- | +| `list(search=..., tags=..., tag_match="any", archived=False, order_by="-created_at")` | Find datasets; `archived=True` selects archived datasets only | +| `get(dataset_id)` | Read current dataset metadata and revision references | +| `create(name=..., description=..., metadata=..., tags=...)` | Create an empty dataset | +| `update(dataset_id, name=..., description=..., metadata=..., tags=...)` | Replace supplied metadata fields without creating an item revision | +| `archive(dataset_id)` / `restore(dataset_id)` | Change dataset archive state; restore leaves individual item states unchanged | +| `list_revisions(dataset_id)` | Read revision history, newest first | +| `list_items(dataset_id, revision=..., archived=False)` | Read one page of current or historical items | +| `get_item(item_id, revision=...)` | Read one stable item at the current or selected dataset revision | +| `create_item(dataset_id, input=..., ...)` | Add a case and its first version | +| `update_item(item_id, base_version=..., input=..., expected_output=..., metadata=...)` | Update editable content with optimistic concurrency | +| `archive_item(item_id, base_version=...)` | Archive an active item with a new version | +| `restore_item(item_id, base_version=..., source_version=...)` | Restore an archived item, optionally copying a historical item version | +| `list_item_versions(item_id)` | Read item history, newest first | + +Dataset and item metadata are JSON objects. Inputs accept any non-null JSON value; expected/source outputs can be null. Writes require finite numbers, JSON string keys, and at most 32 nested containers. The server enforces dataset, item, version, and payload size limits; failures retain their codes and details in `EvaluationAPIError`. JSON values returned by reads are local copies. Editing them does not write to PostHog. + +To save source provenance on creation, supply `source_trace_id` together with a timezone-aware `source_timestamp`; optionally add `source_event_id`. Source fields and `client_item_id` cannot be edited later. + +### Updates and recovery + +```python +case = evaluations.datasets.get_item(item_id) +updated = evaluations.datasets.update_item( + case.id, + base_version=case.version, + expected_output=None, # Clear the answer; omitting it preserves it. + metadata={}, # Clear metadata; omitting it preserves it. +) +``` + +Pass the **observed item version number** as `base_version` for updates, archive, and restore. A stale version raises a 409 `stale_version` conflict with details in `error.response`, including `current_version`. The SDK does not fetch a newer version and overwrite it automatically. Changed content creates a version and dataset revision; identical content is a no-op. Dataset descriptions clear with `""`, metadata with `{}`, and tags with `[]`; those fields do not accept `None`. + +Reads retry within the configured transport limits. Item creates retry only when you supply a `client_item_id`. Other mutations never retry automatically. After an unknown write outcome, inspect the current state before issuing another mutation. Dataset errors use the same `EvaluationAPIError` as scorer and experiment operations, including `status`, `code`, `response`, and `persistence`. + +Run the supplied example against the provisioned dataset: + +```bash +uv run python examples/example-offline-evaluations/main.py \ + --dataset-id '' --state dataset-run.json + +# Optional: select an earlier numeric revision for a new run. +uv run python examples/example-offline-evaluations/main.py \ + --dataset-id '' --dataset-revision 1 --state historical-run.json +``` + +The example prepares and saves the pinned revision, execution identities, scorer version, and results **before the first experiment write**. If the state file already exists, replay uses that file without fetching a newer dataset, selecting another scorer version, or running the application again. Dataset selection arguments apply only when preparing a new state file. Choose a new file for a new run. This small example materializes its results; for larger workloads, use caller-sized batches and persist each prepared batch before upload. + +### Async datasets + +Await management methods and snapshot creation. Use `async for` on `iter_items()`; iteration itself is an async iterator, not an awaitable list. + +```python +from posthog.ai.evaluations import AsyncOfflineEvaluations, EvaluationItem + +async with AsyncOfflineEvaluations( + project_id=123, secret_key=personal_api_key, +) as evaluations: + snapshot = await evaluations.datasets.snapshot(dataset_id) + scorer = await evaluations.scorers.get(scorer_id) + experiment = await evaluations.create_experiment( + name="Async regression", dataset_revision_id=snapshot.revision_id, + ) + async for case in snapshot.iter_items(): + output = capital_answer(case.input) + item = EvaluationItem.from_dataset_item(case, output=output) + if case.expected_output is None: + await experiment.upload_result( + item=item, scorer_version_id=scorer.current_version_id, + status="not_applicable", + ) + else: + await experiment.upload_result( + item=item, scorer_version_id=scorer.current_version_id, + value=output == case.expected_output, + ) + await experiment.complete() +``` + +Install `posthog[async]` for async support. Cancellation propagates and does not automatically complete or fail an experiment. + ## Typed scorer configuration and polarity Configurations are public `TypedDict` types: editors suggest fields and type checkers connect `create(kind=...)` to the correct configuration. They are ordinary dictionaries at runtime; server validation remains authoritative. diff --git a/examples/example-offline-evaluations/main.py b/examples/example-offline-evaluations/main.py index a3c34171d..6c2e7826f 100644 --- a/examples/example-offline-evaluations/main.py +++ b/examples/example-offline-evaluations/main.py @@ -14,6 +14,7 @@ EvaluationAPIError, EvaluationItem, EvaluationResult, + JSONValue, OfflineEvaluations, ) @@ -33,8 +34,39 @@ def resolve_version(evaluations: OfflineEvaluations) -> str: return scorer.current_version_id +def capital_answer(input: JSONValue) -> str: + """A deterministic application to replace with your own model or agent.""" + if not isinstance(input, str): + raise ValueError("The capital-city example expects string dataset inputs.") + return {"Capital of France?": "Paris", "Capital of Italy?": "Milan"}.get( + input, "I don't know" + ) + + +def exact_match_result(item: EvaluationItem, version_id: str) -> EvaluationResult: + payload = item.to_dict()["payload"] + if payload.get("expected_output") is None: + return EvaluationResult( + item=item, + scorer_version_id=version_id, + status="not_applicable", + reasoning="This case has no expected answer.", + ) + return EvaluationResult( + item=item, + scorer_version_id=version_id, + value=payload["output"] == payload["expected_output"], + reasoning="Exact string comparison", + ) + + def prepare_state( - evaluations: OfflineEvaluations, project_id: int, host: str + evaluations: OfflineEvaluations, + project_id: int, + host: str, + *, + dataset_id: str | None = None, + dataset_revision: int | None = None, ) -> dict[str, Any]: version_id = resolve_version(evaluations) cases = [ @@ -51,24 +83,24 @@ def prepare_state( "expected": "Rome", }, ] - items = [ - EvaluationItem( - case_key=case["key"], - input=case["input"], - output=case["output"], - expected_output=case["expected"], - ) - for case in cases - ] - results = [ - EvaluationResult( - item=item, - scorer_version_id=version_id, - value=case["output"] == case["expected"], - reasoning="Exact string comparison", - ) - for item, case in zip(items, cases) - ] + snapshot = None + if dataset_id is not None: + snapshot = evaluations.datasets.snapshot(dataset_id, revision=dataset_revision) + items = [ + EvaluationItem.from_dataset_item(case, output=capital_answer(case.input)) + for case in snapshot.iter_items() + ] + else: + items = [ + EvaluationItem( + case_key=case["key"], + input=case["input"], + output=case["output"], + expected_output=case["expected"], + ) + for case in cases + ] + results = [exact_match_result(item, version_id) for item in items] return { "project_id": project_id, "host": host, @@ -79,6 +111,11 @@ def prepare_state( "run_source": "local", "expected_item_count": len(items), "expected_result_count": len(results), + **( + {"dataset_revision_id": snapshot.revision_id} + if snapshot is not None + else {} + ), }, "items": [item.to_dict() for item in items], "results": [result.to_dict() for result in results], @@ -108,7 +145,17 @@ def main() -> None: default=Path("offline-evaluation-state.json"), help="Create a new prepared run here, or replay the existing run exactly.", ) + parser.add_argument( + "--dataset-id", help="Run the capital-city application on a PostHog dataset." + ) + parser.add_argument( + "--dataset-revision", + type=int, + help="Pin this numeric revision; defaults to latest once.", + ) args = parser.parse_args() + if args.dataset_revision is not None and args.dataset_id is None: + parser.error("--dataset-revision requires --dataset-id") project_id = int(os.environ["POSTHOG_PROJECT_ID"]) host = os.environ.get("POSTHOG_HOST", "https://us.posthog.com") with OfflineEvaluations( @@ -133,7 +180,13 @@ def main() -> None: "The saved run belongs to a different project or host." ) else: - state = prepare_state(evaluations, project_id, host) + state = prepare_state( + evaluations, + project_id, + host, + dataset_id=args.dataset_id, + dataset_revision=args.dataset_revision, + ) save_state(args.state, state) items = {item["id"]: EvaluationItem.from_dict(item) for item in state["items"]} diff --git a/posthog/ai/evaluations/__init__.py b/posthog/ai/evaluations/__init__.py index 7d3e89b0d..2dbd0cc35 100644 --- a/posthog/ai/evaluations/__init__.py +++ b/posthog/ai/evaluations/__init__.py @@ -1,4 +1,4 @@ -"""Typed offline evaluation uploads, experiment lifecycle, and scorer management. +"""Typed offline evaluation uploads, experiments, scorers, and dataset management. from posthog.ai.evaluations import OfflineEvaluations, EvaluationItem @@ -10,6 +10,17 @@ # experiment.complete() """ +from ._async_datasets import AsyncDatasets, AsyncDatasetSnapshot +from ._dataset_types import ( + Dataset, + DatasetItem, + DatasetItemPage, + DatasetOrder, + DatasetPage, + DatasetRevision, + DatasetRevisionPage, +) +from ._datasets import Datasets, DatasetSnapshot from ._client import ( AsyncExperiment, AsyncOfflineEvaluations, @@ -51,6 +62,8 @@ ) __all__ = [ + "AsyncDatasets", + "AsyncDatasetSnapshot", "AsyncExperiment", "AsyncOfflineEvaluations", "AsyncScorers", @@ -60,6 +73,15 @@ "CategoricalPassingRule", "CategoricalScorerConfig", "CategoricalScorerOption", + "Dataset", + "DatasetItem", + "DatasetItemPage", + "DatasetOrder", + "DatasetPage", + "DatasetRevision", + "DatasetRevisionPage", + "Datasets", + "DatasetSnapshot", "EvaluationAPIError", "EvaluationItem", "EvaluationResult", diff --git a/posthog/ai/evaluations/_async_datasets.py b/posthog/ai/evaluations/_async_datasets.py new file mode 100644 index 000000000..46a9cba0a --- /dev/null +++ b/posthog/ai/evaluations/_async_datasets.py @@ -0,0 +1,460 @@ +"""Asynchronous dataset management and revision-pinned iteration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from typing import AsyncIterator, Literal, Sequence +from uuid import UUID + +from ._dataset_serialization import ( + _body, + _dataset, + _integer, + _item, + _page, + _page_params, + _parse, + _revision, + _tag_params, +) +from ._dataset_types import ( + Dataset, + DatasetItem, + DatasetItemPage, + DatasetOrder, + DatasetPage, + DatasetRevisionPage, +) +from ._serialization import ( + JSONValue, + _UNSET, + _Unset, + _boolean, + _identifier, + _timestamp, + _uuid, +) +from ._transport import AsyncTransport + + +class AsyncDatasets: + """Manage datasets through ``AsyncOfflineEvaluations.datasets``. + + Requires a personal key with ``dataset:read`` or ``dataset:write`` and the + datasets feature enabled. Mutations return after acknowledgment; only item + creates with a caller-supplied ``client_item_id`` are automatically retried. + Other lost write responses require inspecting state before retrying. + """ + + def __init__(self, transport: AsyncTransport, project_id: int) -> None: + self._transport = transport + self._project_id = project_id + self._base = f"/api/projects/{project_id}/datasets/" + self._items = f"/api/projects/{project_id}/dataset_items/" + + async def create( + self, + *, + name: str, + description: str | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + tags: Sequence[str] | _Unset = _UNSET, + ) -> Dataset: + """Create an empty dataset. Names must be unique within the project. + + No revision exists until the first item is created. This operation is + not automatically retried, including after an ambiguous acknowledgment. + """ + return _parse( + lambda data: _dataset(data, self._project_id), + await self._transport.request( + "POST", + self._base, + body=_body( + dict( + name=name, description=description, metadata=metadata, tags=tags + ) + ), + retry_safe=False, + ), + ) + + async def get(self, dataset_id: str | UUID) -> Dataset: + """Fetch current metadata, including both current revision identifiers.""" + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + await self._transport.request("GET", f"{self._base}{identity}/"), + ) + + async def list( + self, + *, + limit: int = 50, + offset: int = 0, + search: str | None = None, + archived: bool = False, + order_by: DatasetOrder = "-created_at", + tags: Sequence[str] | None = None, + tag_match: Literal["any", "all"] = "any", + ) -> DatasetPage: + """Fetch one page of active datasets, or archived datasets if requested. + + Search matches names, descriptions, and metadata. Keep filters and pass + ``next_offset`` to continue. Maximum limit is 100; tags accept 20 names. + """ + params = _page_params(limit, offset, 100) + params.update(_tag_params(tags, tag_match)) + params["archived"] = "true" if _boolean(archived) else "false" + if order_by not in ("created_at", "-created_at", "updated_at", "-updated_at"): + raise ValueError("Unsupported dataset order.") + params["order_by"] = order_by + if search is not None: + if not isinstance(search, str) or not search.strip(): + raise ValueError("search must be a nonempty string.") + params["search"] = search + return _parse( + lambda data: DatasetPage( + *_page(data, lambda row: _dataset(row, self._project_id), offset) + ), + await self._transport.request("GET", self._base, params=params), + ) + + async def update( + self, + dataset_id: str | UUID, + *, + name: str | _Unset = _UNSET, + description: str | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + tags: Sequence[str] | _Unset = _UNSET, + ) -> Dataset: + """Replace supplied metadata fields; omitted fields remain unchanged. + + Use description="", metadata={}, or tags=[] to clear them. Metadata + updates do not create revisions. Use archive/restore for archive state. + """ + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + await self._transport.request( + "PATCH", + f"{self._base}{identity}/", + body=_body( + dict( + name=name, description=description, metadata=metadata, tags=tags + ) + ), + retry_safe=False, + ), + ) + + async def archive(self, dataset_id: str | UUID) -> Dataset: + """Archive a dataset. It stays readable but rejects item mutations.""" + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + await self._transport.request( + "POST", f"{self._base}{identity}/archive/", body=b"{}", retry_safe=False + ), + ) + + async def restore(self, dataset_id: str | UUID) -> Dataset: + """Restore an archived dataset without changing individual item states.""" + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + await self._transport.request( + "POST", f"{self._base}{identity}/restore/", body=b"{}", retry_safe=False + ), + ) + + async def list_revisions( + self, dataset_id: str | UUID, *, limit: int = 50, offset: int = 0 + ) -> DatasetRevisionPage: + """List immutable revisions newest first. Maximum page limit is 100.""" + identity = _uuid(dataset_id) + params = _page_params(limit, offset, 100) + return _parse( + lambda data: DatasetRevisionPage( + *_page( + data, lambda row: _revision(row, self._project_id, identity), offset + ) + ), + await self._transport.request( + "GET", f"{self._base}{identity}/revisions/", params=params + ), + ) + + async def list_items( + self, + dataset_id: str | UUID, + *, + revision: int | None = None, + archived: bool = False, + limit: int = 25, + offset: int = 0, + ) -> DatasetItemPage: + """Fetch one item page, optionally at an exact numeric dataset revision. + + Omitted revision reads current state for this request. Use ``snapshot`` + for consistent multi-page evaluation reads. Archived items are excluded + by default; archived=True selects only archived items. Maximum limit is 25. + """ + identity = _uuid(dataset_id) + params = _page_params(limit, offset, 25) + params["dataset"] = identity + params["archived"] = "true" if _boolean(archived) else "false" + if revision is not None: + params["revision"] = _integer(revision) + return _parse( + lambda data: DatasetItemPage( + *_page( + data, + lambda row: _item( + row, + self._project_id, + dataset_id=identity, + revision=revision, + archived=archived, + ), + offset, + ) + ), + await self._transport.request("GET", self._items, params=params), + ) + + async def get_item( + self, item_id: str | UUID, *, revision: int | None = None + ) -> DatasetItem: + """Fetch an item by stable UUID, optionally at a dataset revision number.""" + identity = _uuid(item_id) + params: dict[str, str | int] = ( + {} if revision is None else {"revision": _integer(revision)} + ) + return _parse( + lambda data: _item(data, self._project_id, id=identity, revision=revision), + await self._transport.request( + "GET", f"{self._items}{identity}/", params=params + ), + ) + + async def create_item( + self, + dataset_id: str | UUID, + *, + input: JSONValue, + expected_output: JSONValue | _Unset = _UNSET, + source_output: JSONValue | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + client_item_id: str | None = None, + source_trace_id: str | None = None, + source_event_id: str | None = None, + source_timestamp: datetime | str | None = None, + ) -> DatasetItem: + """Add a case, creating its first version and a dataset revision. + + ``input`` accepts any non-null JSON value. Supply a stable, case-sensitive + ``client_item_id`` to make identical creates retryable across processes. + A matching key with different current content or an archived item raises + a conflict. Source trace ID and timestamp must be supplied together; + source event ID additionally identifies an event within that trace. + """ + identity = _uuid(dataset_id) + client_item_id = _identifier(client_item_id, "client_item_id") + source_trace_id = _identifier(source_trace_id, "source_trace_id") + source_event_id = _identifier(source_event_id, "source_event_id") + if (source_trace_id is None) != (source_timestamp is None) or ( + source_trace_id is None and source_event_id is not None + ): + raise ValueError( + "Source provenance requires both source_trace_id and source_timestamp." + ) + body = _body( + dict( + dataset=identity, + input=input, + expected_output=expected_output, + source_output=source_output, + metadata=metadata, + client_item_id=client_item_id, + source_trace_id=source_trace_id, + source_event_id=source_event_id, + source_timestamp=None + if source_timestamp is None + else _timestamp(source_timestamp), + ) + ) + return _parse( + lambda data: _item( + data, self._project_id, dataset_id=identity, archived=False + ), + await self._transport.request( + "POST", self._items, body=body, retry_safe=client_item_id is not None + ), + ) + + async def update_item( + self, + item_id: str | UUID, + *, + base_version: int, + input: JSONValue | _Unset = _UNSET, + expected_output: JSONValue | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + ) -> DatasetItem: + """Update editable content using the observed item ``version``. + + Changed content creates a version and dataset revision; identical content + is a no-op. Omit fields to preserve them, use expected_output=None to clear + the answer, or metadata={} to clear metadata. Source fields are immutable. + Stale versions raise a conflict without automatic refresh or retry. + """ + identity = _uuid(item_id) + body = _body( + dict( + base_version=_integer(base_version), + input=input, + expected_output=expected_output, + metadata=metadata, + ) + ) + return _parse( + lambda data: _item(data, self._project_id, id=identity, archived=False), + await self._transport.request( + "PATCH", f"{self._items}{identity}/", body=body, retry_safe=False + ), + ) + + async def archive_item( + self, item_id: str | UUID, *, base_version: int + ) -> DatasetItem: + """Archive an active item by creating a version; require its observed version.""" + identity = _uuid(item_id) + body = _body(dict(base_version=_integer(base_version))) + return _parse( + lambda data: _item(data, self._project_id, id=identity, archived=True), + await self._transport.request( + "POST", f"{self._items}{identity}/archive/", body=body, retry_safe=False + ), + ) + + async def restore_item( + self, + item_id: str | UUID, + *, + base_version: int, + source_version: int | None = None, + ) -> DatasetItem: + """Restore an archived item, optionally copying a historical item version. + + ``base_version`` is the observed current version, even when copying an + older ``source_version``. Restoration creates a new immutable version. + """ + identity = _uuid(item_id) + body = _body( + dict( + base_version=_integer(base_version), + source_version=None + if source_version is None + else _integer(source_version), + ) + ) + return _parse( + lambda data: _item(data, self._project_id, id=identity, archived=False), + await self._transport.request( + "POST", f"{self._items}{identity}/restore/", body=body, retry_safe=False + ), + ) + + async def list_item_versions( + self, item_id: str | UUID, *, limit: int = 25, offset: int = 0 + ) -> DatasetItemPage: + """List every item version, newest first. Maximum page limit is 100.""" + identity = _uuid(item_id) + params = _page_params(limit, offset, 100) + return _parse( + lambda data: DatasetItemPage( + *_page( + data, lambda row: _item(row, self._project_id, id=identity), offset + ) + ), + await self._transport.request( + "GET", f"{self._items}{identity}/versions/", params=params + ), + ) + + async def snapshot( + self, dataset_id: str | UUID, *, revision: int | None = None + ) -> AsyncDatasetSnapshot: + """Pin a revision now; item iteration subsequently fetches that exact revision. + + Omit revision to resolve latest once. Historical revisions are resolved + through paginated revision metadata. No items are fetched here. Raises + ValueError for a never-populated dataset or a nonexistent revision. + Keep the owning client open while iterating the returned snapshot. + """ + if revision is not None: + _integer(revision) + dataset = await self.get(dataset_id) + selected = dataset.current_revision if revision is None else revision + if selected is None: + raise ValueError( + "This dataset has no revision. Add an item before taking a snapshot." + ) + if dataset.current_revision is None or selected > dataset.current_revision: + raise ValueError(f"Dataset revision {selected} does not exist.") + if ( + selected == dataset.current_revision + and dataset.current_revision_id is not None + ): + return AsyncDatasetSnapshot( + dataset, selected, dataset.current_revision_id, self + ) + offset = 0 + while True: + page = await self.list_revisions(dataset.id, limit=100, offset=offset) + for candidate in page.results: + if candidate.revision == selected: + return AsyncDatasetSnapshot(dataset, selected, candidate.id, self) + if page.next_offset is None or ( + page.results and page.results[-1].revision < selected + ): + raise ValueError(f"Dataset revision {selected} does not exist.") + offset = page.next_offset + + +@dataclass(frozen=True) +class AsyncDatasetSnapshot: + """A revision-pinned read handle owned by an open AsyncOfflineEvaluations client. + + ``dataset`` is metadata observed when resolving the snapshot; descriptive + metadata is not historically versioned. Use ``revision_id`` when creating + an experiment, not the introducing revision on an individual DatasetItem. + Repeated iteration fetches the same active item versions without caching. + """ + + dataset: Dataset + revision: int + revision_id: str + _datasets: AsyncDatasets = field(repr=False, compare=False) + + async def iter_items(self, *, page_size: int = 25) -> AsyncIterator[DatasetItem]: + """Use async for to fetch active items at this revision, at most 25 per request. + + Each iteration starts at the first page. Memory is bounded by one page; + fetch failures propagate so callers cannot mistake a partial read for a + complete dataset. An existing revision with no active items yields none. + """ + _page_params(page_size, 0, 25) + offset = 0 + while True: + page = await self._datasets.list_items( + self.dataset.id, revision=self.revision, limit=page_size, offset=offset + ) + for item in page.results: + yield item + if page.next_offset is None: + return + offset = page.next_offset diff --git a/posthog/ai/evaluations/_client.py b/posthog/ai/evaluations/_client.py index 69fc622cb..db0a5d695 100644 --- a/posthog/ai/evaluations/_client.py +++ b/posthog/ai/evaluations/_client.py @@ -11,6 +11,8 @@ from typing_extensions import Unpack +from ._async_datasets import AsyncDatasets +from ._datasets import Datasets from ._errors import EvaluationAPIError from ._scorers import AsyncScorers, Scorers from ._serialization import ( @@ -402,11 +404,12 @@ async def fail(self) -> ExperimentReceipt: class OfflineEvaluations: - """Upload offline evaluations and manage scorers in one project/environment. + """Upload offline evaluations and manage scorers and datasets in one project. ``secret_key`` accepts personal or project secret keys with offline_evaluation_ingestion:write. Scorer reads/writes require a personal - key with llm_analytics:read/write. The offline feature must be enabled. + key with llm_analytics:read/write; datasets require dataset:read/write. + Each operation requires its corresponding product feature to be enabled. ``host`` defaults to the US app API; known ingestion hosts are mapped to their app hosts. Custom hosts are preserved. Requests use a 15 second timeout and at most three retries by default. Scorer creation/version bumps @@ -427,6 +430,7 @@ def __init__( ) self._project_id = project_id self.scorers = Scorers(self._transport, project_id) + self.datasets = Datasets(self._transport, project_id) self._path = f"/api/projects/{project_id}/ai_observability/offline_experiments/" def create_experiment( @@ -501,6 +505,7 @@ def __init__( ) self._project_id = project_id self.scorers = AsyncScorers(self._transport, project_id) + self.datasets = AsyncDatasets(self._transport, project_id) self._path = f"/api/projects/{project_id}/ai_observability/offline_experiments/" async def create_experiment( diff --git a/posthog/ai/evaluations/_dataset_serialization.py b/posthog/ai/evaluations/_dataset_serialization.py new file mode 100644 index 000000000..ae7ff7614 --- /dev/null +++ b/posthog/ai/evaluations/_dataset_serialization.py @@ -0,0 +1,209 @@ +"""Dataset wire validation shared by the synchronous and asynchronous clients.""" + +from __future__ import annotations + +import json +from typing import Any, Callable, Sequence, TypeVar, cast +from urllib.parse import parse_qs, urlsplit + +from ._dataset_types import Dataset, DatasetItem, DatasetRevision +from ._errors import EvaluationAPIError +from ._serialization import ( + JSONValue, + _Unset, + _boolean, + _date, + _encode, + _identifier, + _object, + _uuid, + _validate_json, +) +from ._tags import _response_tags, _tag_names + +_T = TypeVar("_T") + + +def _integer(value: Any, minimum: int = 1) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"Expected an integer greater than or equal to {minimum}.") + return value + + +def _text(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("Expected a string.") + return value + + +def _json(value: Any) -> JSONValue: + return cast(JSONValue, json.loads(_encode(value))) + + +def _metadata(value: Any) -> dict[str, JSONValue]: + return cast(dict[str, JSONValue], _json(_object(value))) + + +def _user(value: Any) -> dict[str, Any] | None: + return None if value is None else _metadata(value) + + +def _owner(data: dict[str, Any], project_id: int) -> int: + if _integer(data["team_id"]) != project_id: + raise ValueError("The dataset response belongs to another project.") + return project_id + + +def _dataset(data: dict[str, Any], project_id: int, id: str | None = None) -> Dataset: + identity = _uuid(data["id"]) + if id is not None and identity != id: + raise ValueError("Mismatched dataset ID.") + revision = data["current_revision"] + revision_id = data["current_revision_id"] + if (revision is None) != (revision_id is None): + raise ValueError("Incomplete dataset revision reference.") + return Dataset( + id=identity, + name=_text(data["name"]), + description=_text(data["description"]), + metadata=_metadata(data["metadata"]), + tags=_response_tags(data.get("tags", [])), + archived=_boolean(data["archived"]), + current_revision=None if revision is None else _integer(revision), + current_revision_id=None if revision_id is None else _uuid(revision_id), + created_at=_date(data["created_at"]), + updated_at=None if data["updated_at"] is None else _date(data["updated_at"]), + created_by=_user(data["created_by"]), + team_id=_owner(data, project_id), + ) + + +def _revision( + data: dict[str, Any], project_id: int, dataset_id: str +) -> DatasetRevision: + if _uuid(data["dataset_id"]) != dataset_id: + raise ValueError("Mismatched revision dataset ID.") + return DatasetRevision( + id=_uuid(data["id"]), + dataset_id=dataset_id, + revision=_integer(data["revision"]), + created_at=_date(data["created_at"]), + created_by=_user(data["created_by"]), + team_id=_owner(data, project_id), + ) + + +def _item( + data: dict[str, Any], + project_id: int, + *, + id: str | None = None, + dataset_id: str | None = None, + revision: int | None = None, + archived: bool | None = None, +) -> DatasetItem: + identity, owner = _uuid(data["id"]), _uuid(data["dataset"]) + introduced = _integer(data["dataset_revision"]) + is_archived = _boolean(data["archived"]) + if (id is not None and identity != id) or ( + dataset_id is not None and owner != dataset_id + ): + raise ValueError("Mismatched dataset item identity.") + if revision is not None and introduced > revision: + raise ValueError("The item version is newer than the requested revision.") + if archived is not None and is_archived != archived: + raise ValueError("The item does not match the requested archive filter.") + if data["input"] is None: + raise ValueError("Dataset inputs cannot be null.") + return DatasetItem( + id=identity, + dataset_id=owner, + client_item_id=_identifier(data["client_item_id"], "client_item_id"), + version=_integer(data["version"]), + version_id=_uuid(data["version_id"]), + dataset_revision=introduced, + dataset_revision_id=_uuid(data["dataset_revision_id"]), + archived=is_archived, + input=_json(data["input"]), + expected_output=_json(data["expected_output"]), + source_output=_json(data["source_output"]), + metadata=_metadata(data["metadata"]), + source_trace_id=_identifier(data["source_trace_id"], "source_trace_id"), + source_event_id=_identifier(data["source_event_id"], "source_event_id"), + source_timestamp=None + if data["source_timestamp"] is None + else _date(data["source_timestamp"]), + created_at=_date(data["created_at"]), + updated_at=None if data["updated_at"] is None else _date(data["updated_at"]), + created_by=_user(data["created_by"]), + version_created_at=_date(data["version_created_at"]), + version_created_by=_user(data["version_created_by"]), + team_id=_owner(data, project_id), + ) + + +def _parse(parser: Callable[[dict[str, Any]], _T], data: dict[str, Any]) -> _T: + try: + return parser(data) + except (KeyError, TypeError, ValueError, OverflowError): + raise EvaluationAPIError( + code="invalid_response", + detail="The dataset API returned an invalid response.", + response=data, + persistence="unknown", + ) from None + + +def _page( + data: dict[str, Any], parser: Callable[[dict[str, Any]], _T], offset: int +) -> tuple[int, tuple[_T, ...], int | None]: + rows = data["results"] + if not isinstance(rows, list): + raise ValueError("Expected a list of dataset results.") + results = tuple(parser(_object(row)) for row in rows) + count = _integer(data["count"], 0) + next_offset = None + if data["next"] is not None: + # Extract only the continuation. Never send credentials to a response URL. + offsets = parse_qs(urlsplit(_text(data["next"])).query).get("offset") + if offsets is None or len(offsets) != 1: + raise ValueError("Missing dataset page offset.") + next_offset = _integer(int(offsets[0]), offset + 1) + if not results: + raise ValueError("An empty page must not have a continuation.") + return count, results, next_offset + + +def _page_params(limit: int, offset: int, maximum: int) -> dict[str, str | int]: + if _integer(limit) > maximum: + raise ValueError(f"Page limit must be between 1 and {maximum}.") + return {"limit": limit, "offset": _integer(offset, 0)} + + +def _body(fields: dict[str, Any]) -> bytes: + data = { + key: value for key, value in fields.items() if not isinstance(value, _Unset) + } + for key in ("name", "description"): + if key in data: + value = _text(data[key]) + maximum = 400 if key == "name" else 10_000 + if len(value) > maximum or (key == "name" and not value.strip()): + raise ValueError(f"Invalid dataset {key}.") + if "input" in data and data["input"] is None: + raise ValueError("Dataset inputs cannot be null.") + if "metadata" in data: + _object(data["metadata"]) + if "tags" in data: + data["tags"] = _tag_names(data["tags"]) + _validate_json(data) + return _encode(data) + + +def _tag_params(tags: Sequence[str] | None, tag_match: str) -> dict[str, str | int]: + if tag_match not in ("any", "all"): + raise ValueError("tag_match must be any or all.") + params: dict[str, str | int] = {"tag_match": tag_match} + if tags is not None: + params["tags"] = json.dumps(_tag_names(tags, limit=20), ensure_ascii=False) + return params diff --git a/posthog/ai/evaluations/_dataset_types.py b/posthog/ai/evaluations/_dataset_types.py new file mode 100644 index 000000000..5f51ca253 --- /dev/null +++ b/posthog/ai/evaluations/_dataset_types.py @@ -0,0 +1,111 @@ +"""Typed dataset metadata, immutable version references, and paginated reads.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Literal + +from ._serialization import JSONValue + +DatasetOrder = Literal["created_at", "-created_at", "updated_at", "-updated_at"] + + +@dataclass(frozen=True) +class Dataset: + """A dataset's current metadata, separate from its immutable item revisions. + + Names are unique within a project. Use ``id`` for subsequent operations. + Both revision fields are None until the first item is created. Metadata + dictionaries are local copies, not live server state. + """ + + id: str + name: str + description: str + metadata: dict[str, JSONValue] = field(repr=False) + tags: tuple[str, ...] + archived: bool + current_revision: int | None + current_revision_id: str | None + created_at: datetime + updated_at: datetime | None + created_by: dict[str, Any] | None + team_id: int + + +@dataclass(frozen=True) +class DatasetRevision: + """An immutable dataset revision. Reads use ``revision``; experiments use ``id``.""" + + id: str + dataset_id: str + revision: int + created_at: datetime + created_by: dict[str, Any] | None + team_id: int + + +@dataclass(frozen=True) +class DatasetItem: + """One version of a dataset case, with stable ``id`` and immutable ``version_id``. + + ``dataset_revision`` and ``dataset_revision_id`` identify the revision that + introduced this version, which may precede the snapshot being read. Use + the snapshot's revision ID when creating an experiment. ``source_output`` + is the original trace's output, not a new execution or expected answer. + JSON values are local copies; changing them does not update the dataset. + """ + + id: str + dataset_id: str + client_item_id: str | None + version: int + version_id: str + dataset_revision: int + dataset_revision_id: str + archived: bool + input: JSONValue = field(repr=False) + expected_output: JSONValue = field(repr=False) + source_output: JSONValue = field(repr=False) + metadata: dict[str, JSONValue] = field(repr=False) + source_trace_id: str | None + source_event_id: str | None + source_timestamp: datetime | None + created_at: datetime + updated_at: datetime | None + created_by: dict[str, Any] | None + version_created_at: datetime + version_created_by: dict[str, Any] | None + team_id: int + + +@dataclass(frozen=True) +class DatasetPage: + """One dataset page. Keep filters and pass ``next_offset`` to continue.""" + + count: int + results: tuple[Dataset, ...] + next_offset: int | None + + +@dataclass(frozen=True) +class DatasetRevisionPage: + """One page of revisions, newest first. None ``next_offset`` ends the list.""" + + count: int + results: tuple[DatasetRevision, ...] + next_offset: int | None + + +@dataclass(frozen=True) +class DatasetItemPage: + """One page of items or item history. None ``next_offset`` ends the list. + + Keep the dataset, revision, and archive filter unchanged when continuing + an item listing. For a complete pinned read, prefer ``snapshot.iter_items``. + """ + + count: int + results: tuple[DatasetItem, ...] + next_offset: int | None diff --git a/posthog/ai/evaluations/_datasets.py b/posthog/ai/evaluations/_datasets.py new file mode 100644 index 000000000..571325cdf --- /dev/null +++ b/posthog/ai/evaluations/_datasets.py @@ -0,0 +1,453 @@ +"""Dataset management and revision-pinned iteration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from typing import Iterator, Literal, Sequence +from uuid import UUID + +from ._dataset_serialization import ( + _body, + _dataset, + _integer, + _item, + _page, + _page_params, + _parse, + _revision, + _tag_params, +) +from ._dataset_types import ( + Dataset, + DatasetItem, + DatasetItemPage, + DatasetOrder, + DatasetPage, + DatasetRevisionPage, +) +from ._serialization import ( + JSONValue, + _UNSET, + _Unset, + _boolean, + _identifier, + _timestamp, + _uuid, +) +from ._transport import SyncTransport + + +class Datasets: + """Manage datasets through ``OfflineEvaluations.datasets``. + + Requires a personal key with ``dataset:read`` or ``dataset:write`` and the + datasets feature enabled. Mutations return after acknowledgment; only item + creates with a caller-supplied ``client_item_id`` are automatically retried. + Other lost write responses require inspecting state before retrying. + """ + + def __init__(self, transport: SyncTransport, project_id: int) -> None: + self._transport = transport + self._project_id = project_id + self._base = f"/api/projects/{project_id}/datasets/" + self._items = f"/api/projects/{project_id}/dataset_items/" + + def create( + self, + *, + name: str, + description: str | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + tags: Sequence[str] | _Unset = _UNSET, + ) -> Dataset: + """Create an empty dataset. Names must be unique within the project. + + No revision exists until the first item is created. This operation is + not automatically retried, including after an ambiguous acknowledgment. + """ + return _parse( + lambda data: _dataset(data, self._project_id), + self._transport.request( + "POST", + self._base, + body=_body( + dict( + name=name, description=description, metadata=metadata, tags=tags + ) + ), + retry_safe=False, + ), + ) + + def get(self, dataset_id: str | UUID) -> Dataset: + """Fetch current metadata, including both current revision identifiers.""" + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + self._transport.request("GET", f"{self._base}{identity}/"), + ) + + def list( + self, + *, + limit: int = 50, + offset: int = 0, + search: str | None = None, + archived: bool = False, + order_by: DatasetOrder = "-created_at", + tags: Sequence[str] | None = None, + tag_match: Literal["any", "all"] = "any", + ) -> DatasetPage: + """Fetch one page of active datasets, or archived datasets if requested. + + Search matches names, descriptions, and metadata. Keep filters and pass + ``next_offset`` to continue. Maximum limit is 100; tags accept 20 names. + """ + params = _page_params(limit, offset, 100) + params.update(_tag_params(tags, tag_match)) + params["archived"] = "true" if _boolean(archived) else "false" + if order_by not in ("created_at", "-created_at", "updated_at", "-updated_at"): + raise ValueError("Unsupported dataset order.") + params["order_by"] = order_by + if search is not None: + if not isinstance(search, str) or not search.strip(): + raise ValueError("search must be a nonempty string.") + params["search"] = search + return _parse( + lambda data: DatasetPage( + *_page(data, lambda row: _dataset(row, self._project_id), offset) + ), + self._transport.request("GET", self._base, params=params), + ) + + def update( + self, + dataset_id: str | UUID, + *, + name: str | _Unset = _UNSET, + description: str | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + tags: Sequence[str] | _Unset = _UNSET, + ) -> Dataset: + """Replace supplied metadata fields; omitted fields remain unchanged. + + Use description="", metadata={}, or tags=[] to clear them. Metadata + updates do not create revisions. Use archive/restore for archive state. + """ + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + self._transport.request( + "PATCH", + f"{self._base}{identity}/", + body=_body( + dict( + name=name, description=description, metadata=metadata, tags=tags + ) + ), + retry_safe=False, + ), + ) + + def archive(self, dataset_id: str | UUID) -> Dataset: + """Archive a dataset. It stays readable but rejects item mutations.""" + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + self._transport.request( + "POST", f"{self._base}{identity}/archive/", body=b"{}", retry_safe=False + ), + ) + + def restore(self, dataset_id: str | UUID) -> Dataset: + """Restore an archived dataset without changing individual item states.""" + identity = _uuid(dataset_id) + return _parse( + lambda data: _dataset(data, self._project_id, identity), + self._transport.request( + "POST", f"{self._base}{identity}/restore/", body=b"{}", retry_safe=False + ), + ) + + def list_revisions( + self, dataset_id: str | UUID, *, limit: int = 50, offset: int = 0 + ) -> DatasetRevisionPage: + """List immutable revisions newest first. Maximum page limit is 100.""" + identity = _uuid(dataset_id) + params = _page_params(limit, offset, 100) + return _parse( + lambda data: DatasetRevisionPage( + *_page( + data, lambda row: _revision(row, self._project_id, identity), offset + ) + ), + self._transport.request( + "GET", f"{self._base}{identity}/revisions/", params=params + ), + ) + + def list_items( + self, + dataset_id: str | UUID, + *, + revision: int | None = None, + archived: bool = False, + limit: int = 25, + offset: int = 0, + ) -> DatasetItemPage: + """Fetch one item page, optionally at an exact numeric dataset revision. + + Omitted revision reads current state for this request. Use ``snapshot`` + for consistent multi-page evaluation reads. Archived items are excluded + by default; archived=True selects only archived items. Maximum limit is 25. + """ + identity = _uuid(dataset_id) + params = _page_params(limit, offset, 25) + params["dataset"] = identity + params["archived"] = "true" if _boolean(archived) else "false" + if revision is not None: + params["revision"] = _integer(revision) + return _parse( + lambda data: DatasetItemPage( + *_page( + data, + lambda row: _item( + row, + self._project_id, + dataset_id=identity, + revision=revision, + archived=archived, + ), + offset, + ) + ), + self._transport.request("GET", self._items, params=params), + ) + + def get_item( + self, item_id: str | UUID, *, revision: int | None = None + ) -> DatasetItem: + """Fetch an item by stable UUID, optionally at a dataset revision number.""" + identity = _uuid(item_id) + params: dict[str, str | int] = ( + {} if revision is None else {"revision": _integer(revision)} + ) + return _parse( + lambda data: _item(data, self._project_id, id=identity, revision=revision), + self._transport.request("GET", f"{self._items}{identity}/", params=params), + ) + + def create_item( + self, + dataset_id: str | UUID, + *, + input: JSONValue, + expected_output: JSONValue | _Unset = _UNSET, + source_output: JSONValue | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + client_item_id: str | None = None, + source_trace_id: str | None = None, + source_event_id: str | None = None, + source_timestamp: datetime | str | None = None, + ) -> DatasetItem: + """Add a case, creating its first version and a dataset revision. + + ``input`` accepts any non-null JSON value. Supply a stable, case-sensitive + ``client_item_id`` to make identical creates retryable across processes. + A matching key with different current content or an archived item raises + a conflict. Source trace ID and timestamp must be supplied together; + source event ID additionally identifies an event within that trace. + """ + identity = _uuid(dataset_id) + client_item_id = _identifier(client_item_id, "client_item_id") + source_trace_id = _identifier(source_trace_id, "source_trace_id") + source_event_id = _identifier(source_event_id, "source_event_id") + if (source_trace_id is None) != (source_timestamp is None) or ( + source_trace_id is None and source_event_id is not None + ): + raise ValueError( + "Source provenance requires both source_trace_id and source_timestamp." + ) + body = _body( + dict( + dataset=identity, + input=input, + expected_output=expected_output, + source_output=source_output, + metadata=metadata, + client_item_id=client_item_id, + source_trace_id=source_trace_id, + source_event_id=source_event_id, + source_timestamp=None + if source_timestamp is None + else _timestamp(source_timestamp), + ) + ) + return _parse( + lambda data: _item( + data, self._project_id, dataset_id=identity, archived=False + ), + self._transport.request( + "POST", self._items, body=body, retry_safe=client_item_id is not None + ), + ) + + def update_item( + self, + item_id: str | UUID, + *, + base_version: int, + input: JSONValue | _Unset = _UNSET, + expected_output: JSONValue | _Unset = _UNSET, + metadata: dict[str, JSONValue] | _Unset = _UNSET, + ) -> DatasetItem: + """Update editable content using the observed item ``version``. + + Changed content creates a version and dataset revision; identical content + is a no-op. Omit fields to preserve them, use expected_output=None to clear + the answer, or metadata={} to clear metadata. Source fields are immutable. + Stale versions raise a conflict without automatic refresh or retry. + """ + identity = _uuid(item_id) + body = _body( + dict( + base_version=_integer(base_version), + input=input, + expected_output=expected_output, + metadata=metadata, + ) + ) + return _parse( + lambda data: _item(data, self._project_id, id=identity, archived=False), + self._transport.request( + "PATCH", f"{self._items}{identity}/", body=body, retry_safe=False + ), + ) + + def archive_item(self, item_id: str | UUID, *, base_version: int) -> DatasetItem: + """Archive an active item by creating a version; require its observed version.""" + identity = _uuid(item_id) + body = _body(dict(base_version=_integer(base_version))) + return _parse( + lambda data: _item(data, self._project_id, id=identity, archived=True), + self._transport.request( + "POST", f"{self._items}{identity}/archive/", body=body, retry_safe=False + ), + ) + + def restore_item( + self, + item_id: str | UUID, + *, + base_version: int, + source_version: int | None = None, + ) -> DatasetItem: + """Restore an archived item, optionally copying a historical item version. + + ``base_version`` is the observed current version, even when copying an + older ``source_version``. Restoration creates a new immutable version. + """ + identity = _uuid(item_id) + body = _body( + dict( + base_version=_integer(base_version), + source_version=None + if source_version is None + else _integer(source_version), + ) + ) + return _parse( + lambda data: _item(data, self._project_id, id=identity, archived=False), + self._transport.request( + "POST", f"{self._items}{identity}/restore/", body=body, retry_safe=False + ), + ) + + def list_item_versions( + self, item_id: str | UUID, *, limit: int = 25, offset: int = 0 + ) -> DatasetItemPage: + """List every item version, newest first. Maximum page limit is 100.""" + identity = _uuid(item_id) + params = _page_params(limit, offset, 100) + return _parse( + lambda data: DatasetItemPage( + *_page( + data, lambda row: _item(row, self._project_id, id=identity), offset + ) + ), + self._transport.request( + "GET", f"{self._items}{identity}/versions/", params=params + ), + ) + + def snapshot( + self, dataset_id: str | UUID, *, revision: int | None = None + ) -> DatasetSnapshot: + """Pin a revision now; item iteration subsequently fetches that exact revision. + + Omit revision to resolve latest once. Historical revisions are resolved + through paginated revision metadata. No items are fetched here. Raises + ValueError for a never-populated dataset or a nonexistent revision. + Keep the owning client open while iterating the returned snapshot. + """ + if revision is not None: + _integer(revision) + dataset = self.get(dataset_id) + selected = dataset.current_revision if revision is None else revision + if selected is None: + raise ValueError( + "This dataset has no revision. Add an item before taking a snapshot." + ) + if dataset.current_revision is None or selected > dataset.current_revision: + raise ValueError(f"Dataset revision {selected} does not exist.") + if ( + selected == dataset.current_revision + and dataset.current_revision_id is not None + ): + return DatasetSnapshot(dataset, selected, dataset.current_revision_id, self) + offset = 0 + while True: + page = self.list_revisions(dataset.id, limit=100, offset=offset) + for candidate in page.results: + if candidate.revision == selected: + return DatasetSnapshot(dataset, selected, candidate.id, self) + if page.next_offset is None or ( + page.results and page.results[-1].revision < selected + ): + raise ValueError(f"Dataset revision {selected} does not exist.") + offset = page.next_offset + + +@dataclass(frozen=True) +class DatasetSnapshot: + """A revision-pinned read handle owned by an open OfflineEvaluations client. + + ``dataset`` is metadata observed when resolving the snapshot; descriptive + metadata is not historically versioned. Use ``revision_id`` when creating + an experiment, not the introducing revision on an individual DatasetItem. + Repeated iteration fetches the same active item versions without caching. + """ + + dataset: Dataset + revision: int + revision_id: str + _datasets: Datasets = field(repr=False, compare=False) + + def iter_items(self, *, page_size: int = 25) -> Iterator[DatasetItem]: + """Lazily fetch active items at this revision, at most 25 per request. + + Each iteration starts at the first page. Memory is bounded by one page; + fetch failures propagate so callers cannot mistake a partial read for a + complete dataset. An existing revision with no active items yields none. + """ + _page_params(page_size, 0, 25) + offset = 0 + while True: + page = self._datasets.list_items( + self.dataset.id, revision=self.revision, limit=page_size, offset=offset + ) + yield from page.results + if page.next_offset is None: + return + offset = page.next_offset diff --git a/posthog/ai/evaluations/_types.py b/posthog/ai/evaluations/_types.py index c6862b37d..7992ac1c9 100644 --- a/posthog/ai/evaluations/_types.py +++ b/posthog/ai/evaluations/_types.py @@ -8,6 +8,7 @@ from typing_extensions import TypedDict +from ._dataset_types import DatasetItem from ._errors import EvaluationAPIError from ._serialization import ( JSONValue, @@ -135,6 +136,45 @@ def to_dict(self) -> dict[str, Any]: """Return a defensive, JSON-serializable copy of the complete declaration.""" return json.loads(self._json) + @classmethod + def from_dataset_item( + cls, + item: DatasetItem, + *, + output: JSONValue | _Unset = _UNSET, + id: str | UUID | None = None, + case_key: str | None = None, + trial: str | None = None, + application_trace_id: str | None = None, + ) -> EvaluationItem: + """Declare one execution of a dataset case, preserving its exact version. + + Copies input, expected output, and metadata; supply the new execution's + output and trace ID. Source output/provenance are never substituted. + A fresh execution UUID is generated unless id is supplied. Reuse the + returned item across scorers; call again for each independent trial. + ``case_key`` defaults to the caller-owned client item ID, or stable item + UUID. Omit output for an execution that failed before producing one. + Use the owning snapshot's revision_id when creating the experiment. + """ + if not isinstance(item, DatasetItem): + raise TypeError("Expected a DatasetItem.") + if item.archived: + raise ValueError("Archived dataset items cannot be evaluated.") + return cls( + id=id, + input=item.input, + output=output, + expected_output=item.expected_output, + metadata=item.metadata, + case_key=(item.client_item_id or item.id) if case_key is None else case_key, + trial=trial, + dataset_item_identifier=item.id, + dataset_item_version_identifier=item.version_id, + dataset_item_version_id=item.version_id, + application_trace_id=application_trace_id, + ) + @classmethod def from_dict(cls, data: Mapping[str, Any]) -> EvaluationItem: """Restore a saved declaration, including its required original UUID.""" diff --git a/posthog/test/ai/test_evaluations_datasets.py b/posthog/test/ai/test_evaluations_datasets.py new file mode 100644 index 000000000..504cc491b --- /dev/null +++ b/posthog/test/ai/test_evaluations_datasets.py @@ -0,0 +1,771 @@ +import asyncio +import copy +import importlib.util +import inspect +import json +from dataclasses import FrozenInstanceError, replace +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import AsyncMock, Mock, patch +from uuid import UUID + +import httpx +import pytest +import requests + +from posthog.ai.evaluations import ( + AsyncOfflineEvaluations, + DatasetItemPage, + EvaluationAPIError, + EvaluationItem, + EvaluationResult, + OfflineEvaluations, +) +from posthog.ai.evaluations import _client +from posthog.ai.evaluations._transport import AsyncTransport, SyncTransport + +from .test_evaluations import respond as respond_to_experiment +from .test_evaluations_scorers import scorer_response + +DATASET_ID = "015e4ac4-3659-456b-998c-4345b4e4c846" +ITEM_ID = "7d1e9d28-9b64-4c95-b4ad-3e6fa6b975b1" +OTHER_ITEM_ID = "b492bf6b-bb45-4374-aefb-e40a8eecfc62" +VERSION_ID = "9e2bdd1a-ad9e-4a96-9bc0-77a7f4d8b35e" +REVISION_ID = "ba71a1bb-836f-4c19-9d06-1d36d44948a4" +OLD_REVISION_ID = "a4137aa9-4e27-4829-928b-568cc55a0082" +TIMESTAMP = "2026-10-20T10:00:00Z" +BASE = "/api/projects/123/datasets/" +ITEMS = "/api/projects/123/dataset_items/" + + +def dataset_response(**overrides): + return { + "id": DATASET_ID, + "name": "Capital cities", + "description": "Regression cases", + "metadata": {"owner": "support"}, + "tags": ["regression"], + "archived": False, + "current_revision": 3, + "current_revision_id": REVISION_ID, + "created_at": TIMESTAMP, + "updated_at": None, + "created_by": None, + "team_id": 123, + **overrides, + } + + +def revision_response(**overrides): + return { + "id": OLD_REVISION_ID, + "dataset_id": DATASET_ID, + "revision": 1, + "created_at": TIMESTAMP, + "created_by": None, + "team_id": 123, + **overrides, + } + + +def item_response(**overrides): + return { + "id": ITEM_ID, + "dataset": DATASET_ID, + "client_item_id": "capital-france", + "version": 2, + "version_id": VERSION_ID, + "dataset_revision": 2, + "dataset_revision_id": OLD_REVISION_ID, + "archived": False, + "input": {"question": "Capital of France?"}, + "expected_output": "Paris", + "source_output": "Lyon", + "metadata": {"category": ["geography"]}, + "source_trace_id": "source-trace", + "source_event_id": "source-event", + "source_timestamp": TIMESTAMP, + "created_at": TIMESTAMP, + "updated_at": TIMESTAMP, + "created_by": None, + "version_created_at": TIMESTAMP, + "version_created_by": None, + "team_id": 123, + **overrides, + } + + +def page_response(*rows, next_offset=None, count=None): + return { + "count": len(rows) if count is None else count, + "results": list(rows), + "next": None + if next_offset is None + else f"https://us.posthog.com{ITEMS}?offset={next_offset}", + "previous": None, + } + + +def invoke(method, *args, **kwargs): + result = method(*args, **kwargs) + return asyncio.run(result) if inspect.isawaitable(result) else result + + +def collect(iterator): + async def consume(): + return [item async for item in iterator] + + return asyncio.run(consume()) if hasattr(iterator, "__aiter__") else list(iterator) + + +@pytest.fixture(params=[False, True], ids=["sync", "async"]) +def client(request, monkeypatch): + transport = Mock(spec=SyncTransport) + if request.param: + transport = AsyncMock(spec=AsyncTransport) + monkeypatch.setattr(_client, "AsyncTransport", lambda *args: transport) + evaluations = AsyncOfflineEvaluations(project_id=123, secret_key="phx_example") + else: + monkeypatch.setattr(_client, "SyncTransport", lambda *args: transport) + evaluations = OfflineEvaluations(project_id=123, secret_key="phx_example") + return evaluations, transport + + +def test_dataset_metadata_omission_clearing_and_empty_revision(client): + evaluations, transport = client + response = dataset_response(current_revision=None, current_revision_id=None) + transport.request.return_value = response + dataset = invoke(evaluations.datasets.create, name="Capital cities") + assert dataset.current_revision is None + assert dataset.current_revision_id is None + assert dataset.created_at.tzinfo == timezone.utc + assert json.loads(transport.request.call_args.kwargs["body"]) == { + "name": "Capital cities" + } + assert transport.request.call_args.kwargs["retry_safe"] is False + + response["metadata"]["owner"] = "changed" + assert dataset.metadata == {"owner": "support"} + with pytest.raises(FrozenInstanceError): + dataset.id = OTHER_ITEM_ID + + invoke( + evaluations.datasets.update, + UUID(DATASET_ID), + description="", + metadata={}, + tags=[], + ) + assert json.loads(transport.request.call_args.kwargs["body"]) == { + "description": "", + "metadata": {}, + "tags": [], + } + invoke(evaluations.datasets.update, DATASET_ID, name="Renamed") + assert json.loads(transport.request.call_args.kwargs["body"]) == {"name": "Renamed"} + assert transport.request.call_args.kwargs["retry_safe"] is False + + +@pytest.mark.parametrize("input", [False, 0, "", [], {"messages": []}]) +def test_item_creation_preserves_json_values_and_source_provenance(client, input): + evaluations, transport = client + transport.request.return_value = item_response(input=input) + item = invoke( + evaluations.datasets.create_item, + DATASET_ID, + input=input, + expected_output=None, + source_output="original", + metadata={}, + client_item_id="capital-france", + source_trace_id="trace", + source_event_id="event", + source_timestamp=datetime(2026, 10, 20, tzinfo=timezone.utc), + ) + body = json.loads(transport.request.call_args.kwargs["body"]) + assert item.input == input + assert body == { + "dataset": DATASET_ID, + "input": input, + "expected_output": None, + "source_output": "original", + "metadata": {}, + "client_item_id": "capital-france", + "source_trace_id": "trace", + "source_event_id": "event", + "source_timestamp": "2026-10-20T00:00:00+00:00", + } + assert transport.request.call_args.kwargs["retry_safe"] is True + + +def test_item_updates_distinguish_omission_from_clear_and_protect_source_fields(client): + evaluations, transport = client + transport.request.return_value = item_response() + invoke( + evaluations.datasets.update_item, + ITEM_ID, + base_version=2, + expected_output=None, + metadata={}, + ) + assert json.loads(transport.request.call_args.kwargs["body"]) == { + "base_version": 2, + "expected_output": None, + "metadata": {}, + } + assert transport.request.call_args.kwargs["retry_safe"] is False + invoke(evaluations.datasets.update_item, ITEM_ID, base_version=2, input="new") + assert json.loads(transport.request.call_args.kwargs["body"]) == { + "base_version": 2, + "input": "new", + } + transport.request.reset_mock() + with pytest.raises(TypeError): + invoke( + evaluations.datasets.update_item, + ITEM_ID, + base_version=2, + source_output="overwrite", + ) + transport.request.assert_not_called() + + +@pytest.mark.parametrize("scope", ["dataset", "item"]) +def test_archive_and_restore_use_explicit_transitions(client, scope): + evaluations, transport = client + namespace = evaluations.datasets + if scope == "dataset": + transport.request.return_value = dataset_response(archived=True) + archived = invoke(namespace.archive, DATASET_ID) + assert transport.request.call_args.args == ( + "POST", + BASE + DATASET_ID + "/archive/", + ) + transport.request.return_value = dataset_response() + restored = invoke(namespace.restore, DATASET_ID) + assert transport.request.call_args.args == ( + "POST", + BASE + DATASET_ID + "/restore/", + ) + else: + transport.request.return_value = item_response(archived=True, version=3) + archived = invoke(namespace.archive_item, ITEM_ID, base_version=2) + assert json.loads(transport.request.call_args.kwargs["body"]) == { + "base_version": 2 + } + transport.request.return_value = item_response(version=4) + restored = invoke( + namespace.restore_item, ITEM_ID, base_version=3, source_version=1 + ) + assert json.loads(transport.request.call_args.kwargs["body"]) == { + "base_version": 3, + "source_version": 1, + } + assert transport.request.call_args.args == ( + "POST", + ITEMS + ITEM_ID + "/restore/", + ) + assert archived.archived is True + assert restored.archived is False + assert all( + call.kwargs["retry_safe"] is False for call in transport.request.call_args_list + ) + + +def test_filtered_dataset_pages_keep_continuation_and_normalized_tags(client): + evaluations, transport = client + transport.request.return_value = page_response( + dataset_response(tags=["quality"]), next_offset=50, count=51 + ) + page = invoke( + evaluations.datasets.list, + search="capital", + archived=True, + order_by="updated_at", + tags=[" Quality "], + tag_match="all", + ) + assert page.next_offset == 50 + assert page.count == 51 + assert page.results[0].tags == ("quality",) + assert transport.request.call_args.kwargs["params"] == { + "limit": 50, + "offset": 0, + "search": "capital", + "archived": "true", + "order_by": "updated_at", + "tags": '[" Quality "]', + "tag_match": "all", + } + + +def test_item_history_and_historical_get_keep_stable_and_version_ids_separate(client): + evaluations, transport = client + transport.request.return_value = page_response( + item_response(), item_response(version=1) + ) + page = invoke(evaluations.datasets.list_item_versions, UUID(ITEM_ID), limit=100) + assert isinstance(page, DatasetItemPage) + assert [row.version for row in page.results] == [2, 1] + assert transport.request.call_args.args == ("GET", ITEMS + ITEM_ID + "/versions/") + transport.request.return_value = item_response() + item = invoke(evaluations.datasets.get_item, UUID(ITEM_ID), revision=3) + assert item.id == ITEM_ID and item.version_id == VERSION_ID + assert item.dataset_revision == 2 + assert transport.request.call_args.kwargs["params"] == {"revision": 3} + + +def test_snapshot_pins_before_iteration_and_never_follows_remote_pagination_urls( + client, +): + evaluations, transport = client + latest = dataset_response() + transport.request.return_value = latest + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID) + assert transport.request.call_count == 1 + assert snapshot.revision == 3 and snapshot.revision_id == REVISION_ID + latest.update(current_revision=4, current_revision_id=OLD_REVISION_ID) + iterator = snapshot.iter_items(page_size=1) + assert transport.request.call_count == 1 + first = page_response(item_response(), next_offset=1, count=2) + first["next"] = "https://untrusted.example/api/?offset=1&revision=4" + transport.request.side_effect = [ + first, + page_response(item_response(id=OTHER_ITEM_ID), count=2), + ] + items = collect(iterator) + assert [item.id for item in items] == [ITEM_ID, OTHER_ITEM_ID] + for call in transport.request.call_args_list[1:]: + assert call.args == ("GET", ITEMS) + assert call.kwargs["params"]["revision"] == 3 + assert call.kwargs["params"]["archived"] == "false" + assert transport.request.call_args.kwargs["params"]["offset"] == 1 + + +def test_historical_snapshot_resolves_revision_uuid_across_pages(client): + evaluations, transport = client + transport.request.side_effect = [ + dataset_response(), + page_response( + revision_response(id=REVISION_ID, revision=3), next_offset=1, count=2 + ), + page_response(revision_response(), count=2), + page_response(item_response(dataset_revision=1)), + ] + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID, revision=1) + assert snapshot.revision_id == OLD_REVISION_ID + assert transport.request.call_args.kwargs["params"]["offset"] == 1 + assert snapshot.dataset.current_revision == 3 + assert len(collect(snapshot.iter_items())) == 1 + assert transport.request.call_args.kwargs["params"]["revision"] == 1 + + +def test_never_populated_and_missing_revisions_are_distinct_from_empty_snapshot(client): + evaluations, transport = client + transport.request.return_value = dataset_response( + current_revision=None, current_revision_id=None + ) + with pytest.raises(ValueError, match="no revision"): + invoke(evaluations.datasets.snapshot, DATASET_ID) + transport.request.side_effect = [dataset_response(), page_response()] + with pytest.raises(ValueError, match="does not exist"): + invoke(evaluations.datasets.snapshot, DATASET_ID, revision=9) + assert transport.request.call_args.args == ("GET", BASE + DATASET_ID + "/") + transport.request.side_effect = [dataset_response(archived=True), page_response()] + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID) + assert collect(snapshot.iter_items()) == [] + + +def test_missing_historical_revision_stops_once_history_passes_it(client): + evaluations, transport = client + transport.request.side_effect = [ + dataset_response(), + page_response(revision_response(), next_offset=1, count=2), + ] + with pytest.raises(ValueError, match="revision 2 does not exist"): + invoke(evaluations.datasets.snapshot, DATASET_ID, revision=2) + assert transport.request.call_count == 2 + + +def test_snapshot_propagates_partial_read_failure_and_can_be_iterated_again(client): + evaluations, transport = client + transport.request.return_value = dataset_response() + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID) + failure = EvaluationAPIError(status=503, persistence="unknown") + transport.request.side_effect = [ + page_response(item_response(), next_offset=1, count=2), + failure, + ] + with pytest.raises(EvaluationAPIError) as caught: + collect(snapshot.iter_items(page_size=1)) + assert caught.value is failure + transport.request.side_effect = [page_response(item_response())] + assert len(collect(snapshot.iter_items())) == 1 + assert transport.request.call_args.kwargs["params"]["offset"] == 0 + + +@pytest.mark.parametrize( + "next_url,rows", + [ + ("?offset=0", [item_response()]), + ("?offset=-1", [item_response()]), + ("?offset=one", [item_response()]), + ("?offset=1&offset=2", [item_response()]), + ("?cursor=next", [item_response()]), + ("?offset=1", []), + ], +) +def test_invalid_continuations_cannot_loop_or_silently_truncate(client, next_url, rows): + evaluations, transport = client + transport.request.return_value = {"count": 2, "results": rows, "next": next_url} + with pytest.raises(EvaluationAPIError, match="no acknowledgment") as caught: + invoke(evaluations.datasets.list_items, DATASET_ID) + assert caught.value.code == "invalid_response" + + +@pytest.mark.parametrize( + "overrides", + [ + {"dataset": OTHER_ITEM_ID}, + {"team_id": 999}, + {"dataset_revision": 4}, + {"archived": True}, + {"version_id": "invalid"}, + {"metadata": []}, + {"input": None}, + {"version": True}, + {"created_at": "2026-10-20"}, + ], +) +def test_invalid_item_responses_cannot_enter_a_pinned_run(client, overrides): + evaluations, transport = client + transport.request.return_value = page_response(item_response(**overrides)) + with pytest.raises(EvaluationAPIError) as caught: + invoke(evaluations.datasets.list_items, DATASET_ID, revision=3) + assert caught.value.code == "invalid_response" + assert caught.value.persistence == "unknown" + + +@pytest.mark.parametrize( + "overrides", + [ + {"id": OTHER_ITEM_ID}, + {"team_id": 999}, + {"current_revision": None}, + {"current_revision_id": None}, + {"current_revision": True}, + ], +) +def test_invalid_dataset_metadata_cannot_pin_a_snapshot(client, overrides): + evaluations, transport = client + transport.request.return_value = dataset_response(**overrides) + with pytest.raises(EvaluationAPIError) as caught: + invoke(evaluations.datasets.snapshot, DATASET_ID) + assert caught.value.code == "invalid_response" + + +@pytest.mark.parametrize( + "method,args,kwargs", + [ + ("create", (), {"name": " "}), + ("update", (DATASET_ID,), {"description": None}), + ("update", (DATASET_ID,), {"metadata": None}), + ("update", (DATASET_ID,), {"tags": "tag"}), + ("get", ("../other-project",), {}), + ("list", (), {"limit": 101}), + ("list", (), {"offset": -1}), + ("list", (), {"archived": "false"}), + ("list", (), {"tag_match": "invalid"}), + ("list_items", (DATASET_ID,), {"limit": 26}), + ("list_items", (DATASET_ID,), {"revision": 0}), + ("snapshot", (DATASET_ID,), {"revision": True}), + ("create_item", (DATASET_ID,), {"input": None}), + ("create_item", (DATASET_ID,), {"input": float("nan")}), + ("create_item", (DATASET_ID,), {"input": {1: "invalid key"}}), + ("create_item", (DATASET_ID,), {"input": "q", "client_item_id": ""}), + ("create_item", (DATASET_ID,), {"input": "q", "source_trace_id": "trace"}), + ("create_item", (DATASET_ID,), {"input": "q", "source_event_id": "event"}), + ("create_item", (DATASET_ID,), {"input": "q", "source_timestamp": TIMESTAMP}), + ( + "create_item", + (DATASET_ID,), + { + "input": "q", + "source_trace_id": "trace", + "source_timestamp": "2026-10-20", + }, + ), + ("update_item", (ITEM_ID,), {"base_version": 0}), + ("update_item", (ITEM_ID,), {"base_version": 2, "input": None}), + ("restore_item", (ITEM_ID,), {"base_version": 2, "source_version": False}), + ], +) +def test_invalid_requests_fail_before_network_io(client, method, args, kwargs): + evaluations, transport = client + with pytest.raises(ValueError): + invoke(getattr(evaluations.datasets, method), *args, **kwargs) + transport.request.assert_not_called() + + +def test_conflicts_and_scope_errors_preserve_actionable_server_details(client): + evaluations, transport = client + conflict = EvaluationAPIError( + status=409, + code="stale_version", + response={"current_version": 3, "current_item_id": ITEM_ID}, + persistence="rejected", + ) + transport.request.side_effect = conflict + with pytest.raises(EvaluationAPIError) as caught: + invoke( + evaluations.datasets.update_item, + ITEM_ID, + base_version=2, + expected_output="new", + ) + assert caught.value is conflict + assert caught.value.response["current_version"] == 3 + transport.request.assert_called_once() + transport.request.side_effect = EvaluationAPIError( + status=403, persistence="rejected" + ) + with pytest.raises(EvaluationAPIError) as caught: + invoke(evaluations.datasets.get, DATASET_ID) + assert caught.value.status == 403 + + +def test_dataset_to_execution_copies_payload_and_keeps_trials_independent(client): + evaluations, transport = client + response = item_response() + transport.request.return_value = response + case = invoke(evaluations.datasets.get_item, ITEM_ID) + original = copy.deepcopy(response) + response["metadata"]["category"].append("changed") + output = {"answer": "Paris"} + first = EvaluationItem.from_dataset_item( + case, output=output, trial="1", application_trace_id="new-trace" + ) + second = EvaluationItem.from_dataset_item(case, output="Paris", trial="2") + assert first.id != second.id and first.id != case.id + data = first.to_dict() + assert data["dataset_item_version_id"] == VERSION_ID + assert data["dataset_item_identifier"] == ITEM_ID + assert data["dataset_item_version_identifier"] == VERSION_ID + assert data["case_key"] == "capital-france" + assert data["application_trace_id"] == "new-trace" + assert data["payload"] == { + "input": original["input"], + "output": output, + "expected_output": "Paris", + "metadata": original["metadata"], + } + output["answer"] = "changed" + case.metadata["category"].append("changed locally") + assert first.to_dict() == data + assert "application_trace_id" not in second.to_dict() + assert EvaluationItem.from_dict(data).to_dict() == data + assert ( + EvaluationItem.from_dataset_item(replace(case, client_item_id=None)).to_dict()[ + "case_key" + ] + == ITEM_ID + ) + assert "output" not in EvaluationItem.from_dataset_item(case).to_dict()["payload"] + explicit = EvaluationItem.from_dataset_item( + case, id=first.id, case_key="override", output=None + ) + assert explicit.id == first.id + assert explicit.to_dict()["case_key"] == "override" + assert explicit.to_dict()["payload"]["output"] is None + with pytest.raises(ValueError, match="Archived"): + EvaluationItem.from_dataset_item(replace(case, archived=True)) + + +def test_dataset_scorer_experiment_flow_pins_both_versions_and_reuses_execution(client): + evaluations, transport = client + bodies = [] + + def respond(method, path, *, body=None, **kwargs): + if path == BASE + DATASET_ID + "/": + return dataset_response() + if path == ITEMS: + return page_response(item_response()) + if "score_definitions" in path: + return scorer_response() + if path.endswith("/upload/"): + bodies.append(json.loads(body)) + return respond_to_experiment(method, path, body=body, **kwargs) + + transport.request.side_effect = respond + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID) + scorer = invoke(evaluations.scorers.get, scorer_response()["id"]) + experiment = invoke( + evaluations.create_experiment, + name="Capital cities", + dataset_revision_id=snapshot.revision_id, + ) + case = collect(snapshot.iter_items())[0] + item = EvaluationItem.from_dataset_item(case, output="Paris") + results = [ + EvaluationResult(item=item, scorer_version_id=version, value=True) + for version in [scorer.current_version_id, OTHER_ITEM_ID] + ] + invoke(experiment.upload_results, results) + invoke(experiment.complete) + assert experiment.submission["dataset_revision_id"] == REVISION_ID + assert experiment.submission["dataset_revision_id"] != case.dataset_revision_id + assert len(bodies[0]["items"]) == 1 + assert {result["item_id"] for result in bodies[0]["results"]} == {item.id} + assert bodies[0]["items"][0]["dataset_item_version_id"] == VERSION_ID + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("client_item_id", [None, "stable-key"]) +def test_lost_item_create_acknowledgment_only_retries_with_explicit_identity( + asynchronous, client_item_id +): + attempts = [] + + def payload(attempt): + attempts.append(attempt) + if len(attempts) == 1: + raise ( + requests.ReadTimeout() + if not asynchronous + else httpx.ReadTimeout("lost acknowledgment") + ) + return item_response(client_item_id=client_item_id) + + async def run_async(): + async with AsyncOfflineEvaluations( + project_id=123, secret_key="phx_example" + ) as evaluations: + + def respond(request): + return httpx.Response(200, json=payload((request.url, request.content))) + + await evaluations._transport._client.aclose() + evaluations._transport._client = httpx.AsyncClient( + transport=httpx.MockTransport(respond) + ) + with patch( + "posthog.ai.evaluations._transport.asyncio.sleep", + new_callable=AsyncMock, + ): + return await evaluations.datasets.create_item( + DATASET_ID, input="q", client_item_id=client_item_id + ) + + def run_sync(): + with OfflineEvaluations( + project_id=123, secret_key="phx_example" + ) as evaluations: + + def respond(method, url, **kwargs): + data = payload((url, kwargs["data"])) + result = requests.Response() + result.status_code = 200 + result._content = json.dumps(data).encode() + return result + + with ( + patch.object( + evaluations._transport._session, "request", side_effect=respond + ), + patch("posthog.ai.evaluations._transport.time.sleep"), + ): + return evaluations.datasets.create_item( + DATASET_ID, input="q", client_item_id=client_item_id + ) + + def run(): + return asyncio.run(run_async()) if asynchronous else run_sync() + + if client_item_id is None: + with pytest.raises(EvaluationAPIError) as caught: + run() + assert caught.value.persistence == "unknown" + assert len(attempts) == 1 + else: + assert run().id == ITEM_ID + assert len(attempts) == 2 + assert attempts[0] == attempts[1] + + +def test_dataset_example_replays_persisted_cases_without_refetching_or_rerunning( + monkeypatch, tmp_path +): + example = Path(__file__).parents[3] / "examples/example-offline-evaluations/main.py" + spec = importlib.util.spec_from_file_location("dataset_evaluation_example", example) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + state_path = tmp_path / "dataset-run.json" + monkeypatch.setenv("POSTHOG_PROJECT_ID", "123") + monkeypatch.setenv("POSTHOG_SECRET_KEY", "never-save-this-secret") + monkeypatch.setenv("POSTHOG_HOST", "https://us.posthog.com") + monkeypatch.setenv("POSTHOG_SCORER_VERSION_ID", VERSION_ID) + monkeypatch.setattr( + "sys.argv", + [str(example), "--dataset-id", DATASET_ID, "--state", str(state_path)], + ) + transport = Mock(spec=SyncTransport) + monkeypatch.setattr(_client, "SyncTransport", lambda *args: transport) + + def respond(method, path, **kwargs): + if method == "GET": + assert not state_path.exists() + if path == BASE + DATASET_ID + "/": + return dataset_response() + return page_response(item_response(input="Capital of France?")) + assert state_path.exists() + assert "never-save-this-secret" not in state_path.read_text() + return respond_to_experiment(method, path, **kwargs) + + transport.request.side_effect = respond + module.main() + state = json.loads(state_path.read_text()) + assert state["experiment"]["dataset_revision_id"] == REVISION_ID + assert state["items"][0]["dataset_item_version_id"] == VERSION_ID + assert state["results"][0]["value"] is True + first_writes = [ + call for call in transport.request.call_args_list if call.args[0] == "POST" + ] + transport.request.reset_mock() + monkeypatch.setattr( + module, + "capital_answer", + Mock(side_effect=AssertionError("Execution must not replay")), + ) + module.main() + assert transport.request.call_args_list == first_writes + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +def test_snapshot_cannot_read_after_owning_client_closes(asynchronous): + async def run_async(): + async with AsyncOfflineEvaluations( + project_id=123, secret_key="phx_example" + ) as evaluations: + with patch.object( + evaluations._transport, + "request", + new_callable=AsyncMock, + return_value=dataset_response(), + ): + snapshot = await evaluations.datasets.snapshot(DATASET_ID) + return [item async for item in snapshot.iter_items()] + + def run_sync(): + with OfflineEvaluations( + project_id=123, secret_key="phx_example" + ) as evaluations: + with patch.object( + evaluations._transport, "request", return_value=dataset_response() + ): + snapshot = evaluations.datasets.snapshot(DATASET_ID) + return list(snapshot.iter_items()) + + with pytest.raises(EvaluationAPIError) as caught: + asyncio.run(run_async()) if asynchronous else run_sync() + assert caught.value.code == "client_closed" + assert caught.value.persistence == "not_sent" diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 5a093b650..537e028f2 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -392,9 +392,13 @@ attribute posthog.ai.anthropic.anthropic_providers.AsyncAnthropicBedrock.message attribute posthog.ai.anthropic.anthropic_providers.AsyncAnthropicVertex.messages = AsyncWrappedMessages(self) attribute posthog.ai.claude_agent_sdk.client.log = logging.getLogger('posthog') attribute posthog.ai.claude_agent_sdk.processor.log = logging.getLogger('posthog') +attribute posthog.ai.evaluations.AsyncDatasetSnapshot.dataset: Dataset +attribute posthog.ai.evaluations.AsyncDatasetSnapshot.revision: int +attribute posthog.ai.evaluations.AsyncDatasetSnapshot.revision_id: str attribute posthog.ai.evaluations.AsyncExperiment.id: str attribute posthog.ai.evaluations.AsyncExperiment.receipt = receipt attribute posthog.ai.evaluations.AsyncExperiment.submission: dict[str, Any] | None +attribute posthog.ai.evaluations.AsyncOfflineEvaluations.datasets = AsyncDatasets(self._transport, project_id) attribute posthog.ai.evaluations.AsyncOfflineEvaluations.scorers = AsyncScorers(self._transport, project_id) attribute posthog.ai.evaluations.BooleanScorerConfig.false_label: str attribute posthog.ai.evaluations.BooleanScorerConfig.true_is_failure: bool | None @@ -412,6 +416,58 @@ attribute posthog.ai.evaluations.CategoricalScorerConfig.passing_rule: Categoric attribute posthog.ai.evaluations.CategoricalScorerConfig.selection_mode: Literal['single', 'multiple'] attribute posthog.ai.evaluations.CategoricalScorerOption.key: str attribute posthog.ai.evaluations.CategoricalScorerOption.label: str +attribute posthog.ai.evaluations.Dataset.archived: bool +attribute posthog.ai.evaluations.Dataset.created_at: datetime +attribute posthog.ai.evaluations.Dataset.created_by: dict[str, Any] | None +attribute posthog.ai.evaluations.Dataset.current_revision: int | None +attribute posthog.ai.evaluations.Dataset.current_revision_id: str | None +attribute posthog.ai.evaluations.Dataset.description: str +attribute posthog.ai.evaluations.Dataset.id: str +attribute posthog.ai.evaluations.Dataset.metadata: dict[str, JSONValue] = field(repr=False) +attribute posthog.ai.evaluations.Dataset.name: str +attribute posthog.ai.evaluations.Dataset.tags: tuple[str, ...] +attribute posthog.ai.evaluations.Dataset.team_id: int +attribute posthog.ai.evaluations.Dataset.updated_at: datetime | None +attribute posthog.ai.evaluations.DatasetItem.archived: bool +attribute posthog.ai.evaluations.DatasetItem.client_item_id: str | None +attribute posthog.ai.evaluations.DatasetItem.created_at: datetime +attribute posthog.ai.evaluations.DatasetItem.created_by: dict[str, Any] | None +attribute posthog.ai.evaluations.DatasetItem.dataset_id: str +attribute posthog.ai.evaluations.DatasetItem.dataset_revision: int +attribute posthog.ai.evaluations.DatasetItem.dataset_revision_id: str +attribute posthog.ai.evaluations.DatasetItem.expected_output: JSONValue = field(repr=False) +attribute posthog.ai.evaluations.DatasetItem.id: str +attribute posthog.ai.evaluations.DatasetItem.input: JSONValue = field(repr=False) +attribute posthog.ai.evaluations.DatasetItem.metadata: dict[str, JSONValue] = field(repr=False) +attribute posthog.ai.evaluations.DatasetItem.source_event_id: str | None +attribute posthog.ai.evaluations.DatasetItem.source_output: JSONValue = field(repr=False) +attribute posthog.ai.evaluations.DatasetItem.source_timestamp: datetime | None +attribute posthog.ai.evaluations.DatasetItem.source_trace_id: str | None +attribute posthog.ai.evaluations.DatasetItem.team_id: int +attribute posthog.ai.evaluations.DatasetItem.updated_at: datetime | None +attribute posthog.ai.evaluations.DatasetItem.version: int +attribute posthog.ai.evaluations.DatasetItem.version_created_at: datetime +attribute posthog.ai.evaluations.DatasetItem.version_created_by: dict[str, Any] | None +attribute posthog.ai.evaluations.DatasetItem.version_id: str +attribute posthog.ai.evaluations.DatasetItemPage.count: int +attribute posthog.ai.evaluations.DatasetItemPage.next_offset: int | None +attribute posthog.ai.evaluations.DatasetItemPage.results: tuple[DatasetItem, ...] +attribute posthog.ai.evaluations.DatasetOrder = Literal['created_at', '-created_at', 'updated_at', '-updated_at'] +attribute posthog.ai.evaluations.DatasetPage.count: int +attribute posthog.ai.evaluations.DatasetPage.next_offset: int | None +attribute posthog.ai.evaluations.DatasetPage.results: tuple[Dataset, ...] +attribute posthog.ai.evaluations.DatasetRevision.created_at: datetime +attribute posthog.ai.evaluations.DatasetRevision.created_by: dict[str, Any] | None +attribute posthog.ai.evaluations.DatasetRevision.dataset_id: str +attribute posthog.ai.evaluations.DatasetRevision.id: str +attribute posthog.ai.evaluations.DatasetRevision.revision: int +attribute posthog.ai.evaluations.DatasetRevision.team_id: int +attribute posthog.ai.evaluations.DatasetRevisionPage.count: int +attribute posthog.ai.evaluations.DatasetRevisionPage.next_offset: int | None +attribute posthog.ai.evaluations.DatasetRevisionPage.results: tuple[DatasetRevision, ...] +attribute posthog.ai.evaluations.DatasetSnapshot.dataset: Dataset +attribute posthog.ai.evaluations.DatasetSnapshot.revision: int +attribute posthog.ai.evaluations.DatasetSnapshot.revision_id: str attribute posthog.ai.evaluations.EvaluationAPIError.attr = attr attribute posthog.ai.evaluations.EvaluationAPIError.code = code attribute posthog.ai.evaluations.EvaluationAPIError.detail = deepcopy(detail) @@ -461,6 +517,7 @@ attribute posthog.ai.evaluations.NumericScorerConfig.max: float | None attribute posthog.ai.evaluations.NumericScorerConfig.min: float | None attribute posthog.ai.evaluations.NumericScorerConfig.passing_rule: NumericPassingRule | None attribute posthog.ai.evaluations.NumericScorerConfig.step: float | None +attribute posthog.ai.evaluations.OfflineEvaluations.datasets = Datasets(self._transport, project_id) attribute posthog.ai.evaluations.OfflineEvaluations.scorers = Scorers(self._transport, project_id) attribute posthog.ai.evaluations.ResultOptions.error_code: str | None attribute posthog.ai.evaluations.ResultOptions.error_message: str | None @@ -1128,6 +1185,8 @@ class posthog.ai.anthropic.anthropic_providers.AsyncAnthropicBedrock(posthog_cli class posthog.ai.anthropic.anthropic_providers.AsyncAnthropicVertex(posthog_client: Optional[PostHogClient] = None, **kwargs) class posthog.ai.claude_agent_sdk.client.PostHogClaudeSDKClient(options: Optional[ClaudeAgentOptions] = None, transport: Any = None, *, posthog_client: Optional[Client] = None, posthog_distinct_id: Optional[Union[str, Callable[[ResultMessage], Optional[str]]]] = None, posthog_trace_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: bool = False, posthog_groups: Optional[Dict[str, Any]] = None) class posthog.ai.claude_agent_sdk.processor.PostHogClaudeAgentProcessor(client: Optional[Client] = None, distinct_id: Optional[Union[str, Callable[[ResultMessage], Optional[str]]]] = None, privacy_mode: bool = False, groups: Optional[Dict[str, Any]] = None, properties: Optional[Dict[str, Any]] = None) +class posthog.ai.evaluations.AsyncDatasetSnapshot(dataset: Dataset, revision: int, revision_id: str, _datasets: AsyncDatasets) +class posthog.ai.evaluations.AsyncDatasets(transport: AsyncTransport, project_id: int) class posthog.ai.evaluations.AsyncExperiment(transport: AsyncTransport, project_id: int, id: str, *, receipt: ExperimentReceipt | None = None, submission: bytes | None = None) class posthog.ai.evaluations.AsyncOfflineEvaluations(*, project_id: int, secret_key: str, host: str | None = None, timeout: float = 15, max_retries: int = 3) class posthog.ai.evaluations.AsyncScorers(transport: AsyncTransport, project_id: int) @@ -1137,6 +1196,14 @@ class posthog.ai.evaluations.BulkUploadReceipt(chunks: tuple[UploadReceipt, ...] class posthog.ai.evaluations.CategoricalPassingRule class posthog.ai.evaluations.CategoricalScorerConfig class posthog.ai.evaluations.CategoricalScorerOption +class posthog.ai.evaluations.Dataset(id: str, name: str, description: str, metadata: dict[str, JSONValue], tags: tuple[str, ...], archived: bool, current_revision: int | None, current_revision_id: str | None, created_at: datetime, updated_at: datetime | None, created_by: dict[str, Any] | None, team_id: int) +class posthog.ai.evaluations.DatasetItem(id: str, dataset_id: str, client_item_id: str | None, version: int, version_id: str, dataset_revision: int, dataset_revision_id: str, archived: bool, input: JSONValue, expected_output: JSONValue, source_output: JSONValue, metadata: dict[str, JSONValue], source_trace_id: str | None, source_event_id: str | None, source_timestamp: datetime | None, created_at: datetime, updated_at: datetime | None, created_by: dict[str, Any] | None, version_created_at: datetime, version_created_by: dict[str, Any] | None, team_id: int) +class posthog.ai.evaluations.DatasetItemPage(count: int, results: tuple[DatasetItem, ...], next_offset: int | None) +class posthog.ai.evaluations.DatasetPage(count: int, results: tuple[Dataset, ...], next_offset: int | None) +class posthog.ai.evaluations.DatasetRevision(id: str, dataset_id: str, revision: int, created_at: datetime, created_by: dict[str, Any] | None, team_id: int) +class posthog.ai.evaluations.DatasetRevisionPage(count: int, results: tuple[DatasetRevision, ...], next_offset: int | None) +class posthog.ai.evaluations.DatasetSnapshot(dataset: Dataset, revision: int, revision_id: str, _datasets: Datasets) +class posthog.ai.evaluations.Datasets(transport: SyncTransport, project_id: int) class posthog.ai.evaluations.EvaluationAPIError(*, status: int | None = None, code: str | None = None, detail: Any = None, attr: str | None = None, errors: Any = None, response: dict[str, Any] | None = None, retry_after: float | None = None, persistence: Literal['rejected', 'unknown', 'not_sent'] = 'unknown', submission: dict[str, Any] | None = None) class posthog.ai.evaluations.EvaluationItem(*, id: str | UUID | None = None, input: JSONValue | _Unset = _UNSET, output: JSONValue | _Unset = _UNSET, expected_output: JSONValue | _Unset = _UNSET, metadata: dict[str, JSONValue] | None | _Unset = _UNSET, case_key: str | None = None, trial: str | None = None, dataset_item_identifier: str | None = None, dataset_item_version_identifier: str | None = None, dataset_item_version_id: str | UUID | None = None, application_trace_id: str | None = None, payload: dict[str, JSONValue] | _Unset = _UNSET) class posthog.ai.evaluations.EvaluationResult(*, item: EvaluationItem | str | UUID, scorer_version_id: str | UUID, value: ScoreValue | None = None, status: ResultStatus = 'ok', error_code: str | None = None, reasoning: str | None | _Unset = _UNSET, error_message: str | None | _Unset = _UNSET, metadata: dict[str, JSONValue] | None | _Unset = _UNSET, evaluator_trace_id: str | None = None, evaluated_at: datetime | str | None = None, payload: dict[str, JSONValue] | _Unset = _UNSET) @@ -1501,6 +1568,22 @@ method posthog.ai.claude_agent_sdk.client.PostHogClaudeSDKClient.receive_respons method posthog.ai.claude_agent_sdk.client.PostHogClaudeSDKClient.set_model(model: Optional[str] = None) -> None method posthog.ai.claude_agent_sdk.client.PostHogClaudeSDKClient.set_permission_mode(mode: str) -> None method posthog.ai.claude_agent_sdk.processor.PostHogClaudeAgentProcessor.query(*, prompt: Any, options: Optional[ClaudeAgentOptions] = None, transport: Any = None, posthog_distinct_id: Optional[Union[str, Callable[[ResultMessage], Optional[str]]]] = None, posthog_trace_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: Optional[bool] = None, posthog_groups: Optional[Dict[str, Any]] = None) +method posthog.ai.evaluations.AsyncDatasetSnapshot.iter_items(*, page_size: int = 25) +method posthog.ai.evaluations.AsyncDatasets.archive(dataset_id: str | UUID) +method posthog.ai.evaluations.AsyncDatasets.archive_item(item_id: str | UUID, *, base_version: int) +method posthog.ai.evaluations.AsyncDatasets.create(*, name: str, description: str | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET, tags: Sequence[str] | _Unset = _UNSET) +method posthog.ai.evaluations.AsyncDatasets.create_item(dataset_id: str | UUID, *, input: JSONValue, expected_output: JSONValue | _Unset = _UNSET, source_output: JSONValue | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET, client_item_id: str | None = None, source_trace_id: str | None = None, source_event_id: str | None = None, source_timestamp: datetime | str | None = None) +method posthog.ai.evaluations.AsyncDatasets.get(dataset_id: str | UUID) +method posthog.ai.evaluations.AsyncDatasets.get_item(item_id: str | UUID, *, revision: int | None = None) +method posthog.ai.evaluations.AsyncDatasets.list(*, limit: int = 50, offset: int = 0, search: str | None = None, archived: bool = False, order_by: DatasetOrder = '-created_at', tags: Sequence[str] | None = None, tag_match: Literal['any', 'all'] = 'any') +method posthog.ai.evaluations.AsyncDatasets.list_item_versions(item_id: str | UUID, *, limit: int = 25, offset: int = 0) +method posthog.ai.evaluations.AsyncDatasets.list_items(dataset_id: str | UUID, *, revision: int | None = None, archived: bool = False, limit: int = 25, offset: int = 0) +method posthog.ai.evaluations.AsyncDatasets.list_revisions(dataset_id: str | UUID, *, limit: int = 50, offset: int = 0) +method posthog.ai.evaluations.AsyncDatasets.restore(dataset_id: str | UUID) +method posthog.ai.evaluations.AsyncDatasets.restore_item(item_id: str | UUID, *, base_version: int, source_version: int | None = None) +method posthog.ai.evaluations.AsyncDatasets.snapshot(dataset_id: str | UUID, *, revision: int | None = None) +method posthog.ai.evaluations.AsyncDatasets.update(dataset_id: str | UUID, *, name: str | _Unset = _UNSET, description: str | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET, tags: Sequence[str] | _Unset = _UNSET) +method posthog.ai.evaluations.AsyncDatasets.update_item(item_id: str | UUID, *, base_version: int, input: JSONValue | _Unset = _UNSET, expected_output: JSONValue | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET) method posthog.ai.evaluations.AsyncExperiment.complete() method posthog.ai.evaluations.AsyncExperiment.fail() method posthog.ai.evaluations.AsyncExperiment.get_tags() @@ -1517,6 +1600,23 @@ method posthog.ai.evaluations.AsyncScorers.get_version(scorer_id: str | UUID, ve method posthog.ai.evaluations.AsyncScorers.list(*, limit: int = 100, offset: int = 0, search: str | None = None, kind: ScorerKind | None = None, archived: bool | None = None, order_by: ScorerOrder | None = None, tags: Sequence[str] | None = None, tag_match: Literal['any', 'all'] = 'any') method posthog.ai.evaluations.AsyncScorers.list_versions(scorer_id: str | UUID, *, limit: int = 50, cursor: str | None = None) method posthog.ai.evaluations.AsyncScorers.update(scorer_id: str | UUID, *, name: str | _Omitted = _OMITTED, description: str | None | _Omitted = _OMITTED, archived: bool | _Omitted = _OMITTED, tags: Sequence[str] | _Omitted = _OMITTED) +method posthog.ai.evaluations.DatasetSnapshot.iter_items(*, page_size: int = 25) +method posthog.ai.evaluations.Datasets.archive(dataset_id: str | UUID) +method posthog.ai.evaluations.Datasets.archive_item(item_id: str | UUID, *, base_version: int) +method posthog.ai.evaluations.Datasets.create(*, name: str, description: str | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET, tags: Sequence[str] | _Unset = _UNSET) +method posthog.ai.evaluations.Datasets.create_item(dataset_id: str | UUID, *, input: JSONValue, expected_output: JSONValue | _Unset = _UNSET, source_output: JSONValue | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET, client_item_id: str | None = None, source_trace_id: str | None = None, source_event_id: str | None = None, source_timestamp: datetime | str | None = None) +method posthog.ai.evaluations.Datasets.get(dataset_id: str | UUID) +method posthog.ai.evaluations.Datasets.get_item(item_id: str | UUID, *, revision: int | None = None) +method posthog.ai.evaluations.Datasets.list(*, limit: int = 50, offset: int = 0, search: str | None = None, archived: bool = False, order_by: DatasetOrder = '-created_at', tags: Sequence[str] | None = None, tag_match: Literal['any', 'all'] = 'any') +method posthog.ai.evaluations.Datasets.list_item_versions(item_id: str | UUID, *, limit: int = 25, offset: int = 0) +method posthog.ai.evaluations.Datasets.list_items(dataset_id: str | UUID, *, revision: int | None = None, archived: bool = False, limit: int = 25, offset: int = 0) +method posthog.ai.evaluations.Datasets.list_revisions(dataset_id: str | UUID, *, limit: int = 50, offset: int = 0) +method posthog.ai.evaluations.Datasets.restore(dataset_id: str | UUID) +method posthog.ai.evaluations.Datasets.restore_item(item_id: str | UUID, *, base_version: int, source_version: int | None = None) +method posthog.ai.evaluations.Datasets.snapshot(dataset_id: str | UUID, *, revision: int | None = None) +method posthog.ai.evaluations.Datasets.update(dataset_id: str | UUID, *, name: str | _Unset = _UNSET, description: str | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET, tags: Sequence[str] | _Unset = _UNSET) +method posthog.ai.evaluations.Datasets.update_item(item_id: str | UUID, *, base_version: int, input: JSONValue | _Unset = _UNSET, expected_output: JSONValue | _Unset = _UNSET, metadata: dict[str, JSONValue] | _Unset = _UNSET) +method posthog.ai.evaluations.EvaluationItem.from_dataset_item(item: DatasetItem, *, output: JSONValue | _Unset = _UNSET, id: str | UUID | None = None, case_key: str | None = None, trial: str | None = None, application_trace_id: str | None = None) method posthog.ai.evaluations.EvaluationItem.from_dict(data: Mapping[str, Any]) method posthog.ai.evaluations.EvaluationItem.to_dict() method posthog.ai.evaluations.EvaluationResult.from_dict(data: Mapping[str, Any], *, item: EvaluationItem | None = None) From d86b016ec0962b8eb3956306a451e67d73ee53d7 Mon Sep 17 00:00:00 2001 From: Radu Raicea Date: Tue, 20 Oct 2026 16:01:50 -0400 Subject: [PATCH 2/3] chore(llma): share evaluation request and response helpers --- posthog/ai/evaluations/_async_datasets.py | 59 ++++------- .../ai/evaluations/_dataset_serialization.py | 97 +++++++++++++++---- posthog/ai/evaluations/_datasets.py | 59 ++++------- posthog/ai/evaluations/_scorers.py | 25 ++--- posthog/ai/evaluations/_serialization.py | 22 ++++- posthog/ai/evaluations/_tags.py | 14 ++- posthog/test/ai/test_evaluations_datasets.py | 16 ++- posthog/test/ai/test_evaluations_scorers.py | 14 ++- 8 files changed, 175 insertions(+), 131 deletions(-) diff --git a/posthog/ai/evaluations/_async_datasets.py b/posthog/ai/evaluations/_async_datasets.py index 46a9cba0a..8cb1d86a9 100644 --- a/posthog/ai/evaluations/_async_datasets.py +++ b/posthog/ai/evaluations/_async_datasets.py @@ -9,14 +9,16 @@ from ._dataset_serialization import ( _body, + _create_item_body, _dataset, _integer, _item, + _item_list_params, + _list_params, _page, _page_params, _parse, _revision, - _tag_params, ) from ._dataset_types import ( Dataset, @@ -30,9 +32,6 @@ JSONValue, _UNSET, _Unset, - _boolean, - _identifier, - _timestamp, _uuid, ) from ._transport import AsyncTransport @@ -104,16 +103,9 @@ async def list( Search matches names, descriptions, and metadata. Keep filters and pass ``next_offset`` to continue. Maximum limit is 100; tags accept 20 names. """ - params = _page_params(limit, offset, 100) - params.update(_tag_params(tags, tag_match)) - params["archived"] = "true" if _boolean(archived) else "false" - if order_by not in ("created_at", "-created_at", "updated_at", "-updated_at"): - raise ValueError("Unsupported dataset order.") - params["order_by"] = order_by - if search is not None: - if not isinstance(search, str) or not search.strip(): - raise ValueError("search must be a nonempty string.") - params["search"] = search + params = _list_params( + limit, offset, search, archived, order_by, tags, tag_match + ) return _parse( lambda data: DatasetPage( *_page(data, lambda row: _dataset(row, self._project_id), offset) @@ -203,11 +195,7 @@ async def list_items( by default; archived=True selects only archived items. Maximum limit is 25. """ identity = _uuid(dataset_id) - params = _page_params(limit, offset, 25) - params["dataset"] = identity - params["archived"] = "true" if _boolean(archived) else "false" - if revision is not None: - params["revision"] = _integer(revision) + params = _item_list_params(identity, revision, archived, limit, offset) return _parse( lambda data: DatasetItemPage( *_page( @@ -262,29 +250,16 @@ async def create_item( source event ID additionally identifies an event within that trace. """ identity = _uuid(dataset_id) - client_item_id = _identifier(client_item_id, "client_item_id") - source_trace_id = _identifier(source_trace_id, "source_trace_id") - source_event_id = _identifier(source_event_id, "source_event_id") - if (source_trace_id is None) != (source_timestamp is None) or ( - source_trace_id is None and source_event_id is not None - ): - raise ValueError( - "Source provenance requires both source_trace_id and source_timestamp." - ) - body = _body( - dict( - dataset=identity, - input=input, - expected_output=expected_output, - source_output=source_output, - metadata=metadata, - client_item_id=client_item_id, - source_trace_id=source_trace_id, - source_event_id=source_event_id, - source_timestamp=None - if source_timestamp is None - else _timestamp(source_timestamp), - ) + body = _create_item_body( + identity, + input=input, + expected_output=expected_output, + source_output=source_output, + metadata=metadata, + client_item_id=client_item_id, + source_trace_id=source_trace_id, + source_event_id=source_event_id, + source_timestamp=source_timestamp, ) return _parse( lambda data: _item( diff --git a/posthog/ai/evaluations/_dataset_serialization.py b/posthog/ai/evaluations/_dataset_serialization.py index ae7ff7614..df68820d9 100644 --- a/posthog/ai/evaluations/_dataset_serialization.py +++ b/posthog/ai/evaluations/_dataset_serialization.py @@ -3,11 +3,11 @@ from __future__ import annotations import json +from datetime import datetime from typing import Any, Callable, Sequence, TypeVar, cast from urllib.parse import parse_qs, urlsplit -from ._dataset_types import Dataset, DatasetItem, DatasetRevision -from ._errors import EvaluationAPIError +from ._dataset_types import Dataset, DatasetItem, DatasetOrder, DatasetRevision from ._serialization import ( JSONValue, _Unset, @@ -16,10 +16,12 @@ _encode, _identifier, _object, + _parse_response, + _timestamp, _uuid, _validate_json, ) -from ._tags import _response_tags, _tag_names +from ._tags import _response_tags, _tag_names, _tag_params _T = TypeVar("_T") @@ -143,15 +145,9 @@ def _item( def _parse(parser: Callable[[dict[str, Any]], _T], data: dict[str, Any]) -> _T: - try: - return parser(data) - except (KeyError, TypeError, ValueError, OverflowError): - raise EvaluationAPIError( - code="invalid_response", - detail="The dataset API returned an invalid response.", - response=data, - persistence="unknown", - ) from None + return _parse_response( + parser, data, detail="The dataset API returned an invalid response." + ) def _page( @@ -180,6 +176,40 @@ def _page_params(limit: int, offset: int, maximum: int) -> dict[str, str | int]: return {"limit": limit, "offset": _integer(offset, 0)} +def _list_params( + limit: int, + offset: int, + search: str | None, + archived: bool, + order_by: DatasetOrder, + tags: Sequence[str] | None, + tag_match: str, +) -> dict[str, str | int]: + params = _page_params(limit, offset, 100) + params.update(_tag_params(tags, tag_match)) + params["tag_match"] = tag_match + params["archived"] = "true" if _boolean(archived) else "false" + if order_by not in ("created_at", "-created_at", "updated_at", "-updated_at"): + raise ValueError("Unsupported dataset order.") + params["order_by"] = order_by + if search is not None: + if not isinstance(search, str) or not search.strip(): + raise ValueError("search must be a nonempty string.") + params["search"] = search + return params + + +def _item_list_params( + dataset_id: str, revision: int | None, archived: bool, limit: int, offset: int +) -> dict[str, str | int]: + params = _page_params(limit, offset, 25) + params["dataset"] = dataset_id + params["archived"] = "true" if _boolean(archived) else "false" + if revision is not None: + params["revision"] = _integer(revision) + return params + + def _body(fields: dict[str, Any]) -> bytes: data = { key: value for key, value in fields.items() if not isinstance(value, _Unset) @@ -200,10 +230,39 @@ def _body(fields: dict[str, Any]) -> bytes: return _encode(data) -def _tag_params(tags: Sequence[str] | None, tag_match: str) -> dict[str, str | int]: - if tag_match not in ("any", "all"): - raise ValueError("tag_match must be any or all.") - params: dict[str, str | int] = {"tag_match": tag_match} - if tags is not None: - params["tags"] = json.dumps(_tag_names(tags, limit=20), ensure_ascii=False) - return params +def _create_item_body( + dataset_id: str, + *, + input: JSONValue, + expected_output: JSONValue | _Unset, + source_output: JSONValue | _Unset, + metadata: dict[str, JSONValue] | _Unset, + client_item_id: str | None, + source_trace_id: str | None, + source_event_id: str | None, + source_timestamp: datetime | str | None, +) -> bytes: + client_item_id = _identifier(client_item_id, "client_item_id") + source_trace_id = _identifier(source_trace_id, "source_trace_id") + source_event_id = _identifier(source_event_id, "source_event_id") + if (source_trace_id is None) != (source_timestamp is None) or ( + source_trace_id is None and source_event_id is not None + ): + raise ValueError( + "Source provenance requires both source_trace_id and source_timestamp." + ) + return _body( + dict( + dataset=dataset_id, + input=input, + expected_output=expected_output, + source_output=source_output, + metadata=metadata, + client_item_id=client_item_id, + source_trace_id=source_trace_id, + source_event_id=source_event_id, + source_timestamp=None + if source_timestamp is None + else _timestamp(source_timestamp), + ) + ) diff --git a/posthog/ai/evaluations/_datasets.py b/posthog/ai/evaluations/_datasets.py index 571325cdf..d92b7d5b1 100644 --- a/posthog/ai/evaluations/_datasets.py +++ b/posthog/ai/evaluations/_datasets.py @@ -9,14 +9,16 @@ from ._dataset_serialization import ( _body, + _create_item_body, _dataset, _integer, _item, + _item_list_params, + _list_params, _page, _page_params, _parse, _revision, - _tag_params, ) from ._dataset_types import ( Dataset, @@ -30,9 +32,6 @@ JSONValue, _UNSET, _Unset, - _boolean, - _identifier, - _timestamp, _uuid, ) from ._transport import SyncTransport @@ -104,16 +103,9 @@ def list( Search matches names, descriptions, and metadata. Keep filters and pass ``next_offset`` to continue. Maximum limit is 100; tags accept 20 names. """ - params = _page_params(limit, offset, 100) - params.update(_tag_params(tags, tag_match)) - params["archived"] = "true" if _boolean(archived) else "false" - if order_by not in ("created_at", "-created_at", "updated_at", "-updated_at"): - raise ValueError("Unsupported dataset order.") - params["order_by"] = order_by - if search is not None: - if not isinstance(search, str) or not search.strip(): - raise ValueError("search must be a nonempty string.") - params["search"] = search + params = _list_params( + limit, offset, search, archived, order_by, tags, tag_match + ) return _parse( lambda data: DatasetPage( *_page(data, lambda row: _dataset(row, self._project_id), offset) @@ -203,11 +195,7 @@ def list_items( by default; archived=True selects only archived items. Maximum limit is 25. """ identity = _uuid(dataset_id) - params = _page_params(limit, offset, 25) - params["dataset"] = identity - params["archived"] = "true" if _boolean(archived) else "false" - if revision is not None: - params["revision"] = _integer(revision) + params = _item_list_params(identity, revision, archived, limit, offset) return _parse( lambda data: DatasetItemPage( *_page( @@ -260,29 +248,16 @@ def create_item( source event ID additionally identifies an event within that trace. """ identity = _uuid(dataset_id) - client_item_id = _identifier(client_item_id, "client_item_id") - source_trace_id = _identifier(source_trace_id, "source_trace_id") - source_event_id = _identifier(source_event_id, "source_event_id") - if (source_trace_id is None) != (source_timestamp is None) or ( - source_trace_id is None and source_event_id is not None - ): - raise ValueError( - "Source provenance requires both source_trace_id and source_timestamp." - ) - body = _body( - dict( - dataset=identity, - input=input, - expected_output=expected_output, - source_output=source_output, - metadata=metadata, - client_item_id=client_item_id, - source_trace_id=source_trace_id, - source_event_id=source_event_id, - source_timestamp=None - if source_timestamp is None - else _timestamp(source_timestamp), - ) + body = _create_item_body( + identity, + input=input, + expected_output=expected_output, + source_output=source_output, + metadata=metadata, + client_item_id=client_item_id, + source_trace_id=source_trace_id, + source_event_id=source_event_id, + source_timestamp=source_timestamp, ) return _parse( lambda data: _item( diff --git a/posthog/ai/evaluations/_scorers.py b/posthog/ai/evaluations/_scorers.py index 1d9e15483..4dd2aa450 100644 --- a/posthog/ai/evaluations/_scorers.py +++ b/posthog/ai/evaluations/_scorers.py @@ -9,7 +9,6 @@ from urllib.parse import parse_qs, urlsplit from uuid import UUID -from ._errors import EvaluationAPIError from ._scorer_types import ( BooleanScorerConfig, CategoricalScorerConfig, @@ -22,7 +21,8 @@ ScorerVersion, ScorerVersionPage, ) -from ._tags import _response_tags, _tag_names +from ._serialization import _parse_response +from ._tags import _response_tags, _tag_names, _tag_params from ._transport import AsyncTransport, SyncTransport @@ -155,15 +155,9 @@ def _version_page(data: dict[str, Any]) -> ScorerVersionPage: def _parse( parser: Callable[[dict[str, Any]], _ResultT], data: dict[str, Any] ) -> _ResultT: - try: - return parser(data) - except (KeyError, TypeError, ValueError, OverflowError): - raise EvaluationAPIError( - code="invalid_response", - detail="The scorer API returned an invalid response.", - response=data, - persistence="unknown", - ) from None + return _parse_response( + parser, data, detail="The scorer API returned an invalid response." + ) def _encode(data: dict[str, Any]) -> bytes: @@ -235,8 +229,7 @@ def _list_params( tags: Sequence[str] | None, tag_match: Literal["any", "all"], ) -> dict[str, str | int]: - if tag_match not in ("any", "all"): - raise ValueError("tag_match must be any or all.") + tag_params = _tag_params(tags, tag_match) params: dict[str, str | int] = { "limit": _integer(limit, minimum=1), "offset": _integer(offset), @@ -251,11 +244,7 @@ def _list_params( params["archived"] = "true" if archived else "false" if order_by is not None: params["order_by"] = order_by - if tags is not None: - params["tags"] = json.dumps(_tag_names(tags, limit=20), ensure_ascii=False) - params["tag_match"] = tag_match - elif tag_match != "any": - params["tag_match"] = tag_match + params.update(tag_params) return params diff --git a/posthog/ai/evaluations/_serialization.py b/posthog/ai/evaluations/_serialization.py index b0d8bf448..4e09bb8fa 100644 --- a/posthog/ai/evaluations/_serialization.py +++ b/posthog/ai/evaluations/_serialization.py @@ -3,9 +3,11 @@ import json import math from datetime import datetime, timezone -from typing import Any, Iterable, Union, cast +from typing import Any, Callable, Iterable, TypeVar, Union, cast from uuid import UUID +from ._errors import EvaluationAPIError + JSONValue = Union[ None, bool, int, float, str, list["JSONValue"], dict[str, "JSONValue"] ] @@ -22,6 +24,24 @@ def __repr__(self) -> str: _MAX_REQUEST_BYTES = 5 * 1024 * 1024 _MAX_ITEM_BYTES = 1024 * 1024 _MAX_RESULT_BYTES = 256 * 1024 +_ResponseT = TypeVar("_ResponseT") + + +def _parse_response( + parser: Callable[[dict[str, Any]], _ResponseT], + data: dict[str, Any], + *, + detail: str, +) -> _ResponseT: + try: + return parser(data) + except (KeyError, TypeError, ValueError, OverflowError): + raise EvaluationAPIError( + code="invalid_response", + detail=detail, + response=data, + persistence="unknown", + ) from None def _uuid(value: object) -> str: diff --git a/posthog/ai/evaluations/_tags.py b/posthog/ai/evaluations/_tags.py index be1448f70..b78f0ca17 100644 --- a/posthog/ai/evaluations/_tags.py +++ b/posthog/ai/evaluations/_tags.py @@ -1,7 +1,8 @@ -"""Shared validation for mutable scorer and experiment tags.""" +"""Shared validation and filters for scorer, dataset, and experiment tags.""" from __future__ import annotations +import json from collections.abc import Sequence from typing import Any @@ -26,3 +27,14 @@ def _response_tags(value: Any) -> tuple[str, ...]: if not isinstance(value, list) or any(not isinstance(tag, str) for tag in value): raise ValueError("Expected a list of tag names.") return tuple(value) + + +def _tag_params(tags: Sequence[str] | None, tag_match: str) -> dict[str, str | int]: + if tag_match not in ("any", "all"): + raise ValueError("tag_match must be any or all.") + params: dict[str, str | int] = {} + if tags is not None: + params["tags"] = json.dumps(_tag_names(tags, limit=20), ensure_ascii=False) + if tags is not None or tag_match != "any": + params["tag_match"] = tag_match + return params diff --git a/posthog/test/ai/test_evaluations_datasets.py b/posthog/test/ai/test_evaluations_datasets.py index 504cc491b..b5ed783ce 100644 --- a/posthog/test/ai/test_evaluations_datasets.py +++ b/posthog/test/ai/test_evaluations_datasets.py @@ -272,7 +272,11 @@ def test_archive_and_restore_use_explicit_transitions(client, scope): ) -def test_filtered_dataset_pages_keep_continuation_and_normalized_tags(client): +@pytest.mark.parametrize("tags", [None, [], [" Quality "]]) +@pytest.mark.parametrize("tag_match", ["any", "all"]) +def test_filtered_dataset_pages_keep_continuation_and_normalized_tags( + client, tags, tag_match +): evaluations, transport = client transport.request.return_value = page_response( dataset_response(tags=["quality"]), next_offset=50, count=51 @@ -282,8 +286,8 @@ def test_filtered_dataset_pages_keep_continuation_and_normalized_tags(client): search="capital", archived=True, order_by="updated_at", - tags=[" Quality "], - tag_match="all", + tags=tags, + tag_match=tag_match, ) assert page.next_offset == 50 assert page.count == 51 @@ -294,8 +298,8 @@ def test_filtered_dataset_pages_keep_continuation_and_normalized_tags(client): "search": "capital", "archived": "true", "order_by": "updated_at", - "tags": '[" Quality "]', - "tag_match": "all", + **({"tags": json.dumps(tags)} if tags is not None else {}), + "tag_match": tag_match, } @@ -444,6 +448,8 @@ def test_invalid_item_responses_cannot_enter_a_pinned_run(client, overrides): invoke(evaluations.datasets.list_items, DATASET_ID, revision=3) assert caught.value.code == "invalid_response" assert caught.value.persistence == "unknown" + assert caught.value.detail == "The dataset API returned an invalid response." + assert caught.value.response == transport.request.return_value @pytest.mark.parametrize( diff --git a/posthog/test/ai/test_evaluations_scorers.py b/posthog/test/ai/test_evaluations_scorers.py index 9fa40e918..54cf97bac 100644 --- a/posthog/test/ai/test_evaluations_scorers.py +++ b/posthog/test/ai/test_evaluations_scorers.py @@ -214,7 +214,7 @@ def test_invalid_scorer_response_tags_raise(scorers, tags): @pytest.mark.parametrize("tag_match", ["any", "all"]) -@pytest.mark.parametrize("tags", [[], ("quality", 'a,b"c', "品質")]) +@pytest.mark.parametrize("tags", [None, [], ("quality", 'a,b"c', "品質")]) def test_scorer_list_tag_filters_use_json_and_preserve_other_filters( scorers, tags, tag_match ): @@ -229,8 +229,14 @@ def test_scorer_list_tag_filters_use_json_and_preserve_other_filters( namespace.list, tags=tags, tag_match=tag_match, offset=3, search="score" ) params = transport.request.call_args.kwargs["params"] - assert json.loads(params["tags"]) == list(tags) - assert params["tag_match"] == tag_match + if tags is None: + assert "tags" not in params + else: + assert json.loads(params["tags"]) == list(tags) + if tags is None and tag_match == "any": + assert "tag_match" not in params + else: + assert params["tag_match"] == tag_match assert params["offset"] == 3 assert params["search"] == "score" assert page.results[0].tags == ("quality",) @@ -418,6 +424,8 @@ def test_malformed_success_preserves_unknown_persistence(scorers, override): ) assert caught.value.code == "invalid_response" assert caught.value.persistence == "unknown" + assert caught.value.detail == "The scorer API returned an invalid response." + assert caught.value.response == transport.request.return_value assert "Quality" not in str(caught.value) transport.request.assert_called_once() From 2936009887d4b23c9eb4156f77e06734e9513e78 Mon Sep 17 00:00:00 2001 From: Radu Raicea Date: Wed, 21 Oct 2026 10:26:03 -0400 Subject: [PATCH 3/3] fix(llma): address dataset sdk review feedback --- posthog/ai/evaluations/_async_datasets.py | 10 ++- .../ai/evaluations/_dataset_serialization.py | 35 +++------- posthog/ai/evaluations/_datasets.py | 10 ++- posthog/ai/evaluations/_scorers.py | 29 ++------- posthog/ai/evaluations/_serialization.py | 27 +++++++- posthog/test/ai/test_evaluations_datasets.py | 65 ++++++++++++++++++- 6 files changed, 120 insertions(+), 56 deletions(-) diff --git a/posthog/ai/evaluations/_async_datasets.py b/posthog/ai/evaluations/_async_datasets.py index 8cb1d86a9..80cc6418c 100644 --- a/posthog/ai/evaluations/_async_datasets.py +++ b/posthog/ai/evaluations/_async_datasets.py @@ -11,7 +11,6 @@ _body, _create_item_body, _dataset, - _integer, _item, _item_list_params, _list_params, @@ -32,6 +31,7 @@ JSONValue, _UNSET, _Unset, + _integer, _uuid, ) from ._transport import AsyncTransport @@ -387,6 +387,14 @@ async def snapshot( return AsyncDatasetSnapshot( dataset, selected, dataset.current_revision_id, self ) + # Revisions are consecutive and newest first. Concurrent writes can + # shift their offsets, so verify the result and scan on a miss. + page = await self.list_revisions( + dataset.id, limit=1, offset=dataset.current_revision - selected + ) + for candidate in page.results: + if candidate.revision == selected: + return AsyncDatasetSnapshot(dataset, selected, candidate.id, self) offset = 0 while True: page = await self.list_revisions(dataset.id, limit=100, offset=offset) diff --git a/posthog/ai/evaluations/_dataset_serialization.py b/posthog/ai/evaluations/_dataset_serialization.py index df68820d9..c95fb5aea 100644 --- a/posthog/ai/evaluations/_dataset_serialization.py +++ b/posthog/ai/evaluations/_dataset_serialization.py @@ -5,7 +5,6 @@ import json from datetime import datetime from typing import Any, Callable, Sequence, TypeVar, cast -from urllib.parse import parse_qs, urlsplit from ._dataset_types import Dataset, DatasetItem, DatasetOrder, DatasetRevision from ._serialization import ( @@ -15,8 +14,11 @@ _date, _encode, _identifier, + _integer, + _next_offset, _object, _parse_response, + _text, _timestamp, _uuid, _validate_json, @@ -26,20 +28,8 @@ _T = TypeVar("_T") -def _integer(value: Any, minimum: int = 1) -> int: - if type(value) is not int or value < minimum: - raise ValueError(f"Expected an integer greater than or equal to {minimum}.") - return value - - -def _text(value: Any) -> str: - if not isinstance(value, str): - raise ValueError("Expected a string.") - return value - - def _json(value: Any) -> JSONValue: - return cast(JSONValue, json.loads(_encode(value))) + return cast(JSONValue, json.loads(_encode(value, sort_keys=False))) def _metadata(value: Any) -> dict[str, JSONValue]: @@ -157,23 +147,17 @@ def _page( if not isinstance(rows, list): raise ValueError("Expected a list of dataset results.") results = tuple(parser(_object(row)) for row in rows) - count = _integer(data["count"], 0) - next_offset = None - if data["next"] is not None: - # Extract only the continuation. Never send credentials to a response URL. - offsets = parse_qs(urlsplit(_text(data["next"])).query).get("offset") - if offsets is None or len(offsets) != 1: - raise ValueError("Missing dataset page offset.") - next_offset = _integer(int(offsets[0]), offset + 1) - if not results: - raise ValueError("An empty page must not have a continuation.") + count = _integer(data["count"], minimum=0) + next_offset = _next_offset(data["next"], minimum=offset + 1) + if next_offset is not None and not results: + raise ValueError("An empty page must not have a continuation.") return count, results, next_offset def _page_params(limit: int, offset: int, maximum: int) -> dict[str, str | int]: if _integer(limit) > maximum: raise ValueError(f"Page limit must be between 1 and {maximum}.") - return {"limit": limit, "offset": _integer(offset, 0)} + return {"limit": limit, "offset": _integer(offset, minimum=0)} def _list_params( @@ -187,7 +171,6 @@ def _list_params( ) -> dict[str, str | int]: params = _page_params(limit, offset, 100) params.update(_tag_params(tags, tag_match)) - params["tag_match"] = tag_match params["archived"] = "true" if _boolean(archived) else "false" if order_by not in ("created_at", "-created_at", "updated_at", "-updated_at"): raise ValueError("Unsupported dataset order.") diff --git a/posthog/ai/evaluations/_datasets.py b/posthog/ai/evaluations/_datasets.py index d92b7d5b1..cfe359a2e 100644 --- a/posthog/ai/evaluations/_datasets.py +++ b/posthog/ai/evaluations/_datasets.py @@ -11,7 +11,6 @@ _body, _create_item_body, _dataset, - _integer, _item, _item_list_params, _list_params, @@ -32,6 +31,7 @@ JSONValue, _UNSET, _Unset, + _integer, _uuid, ) from ._transport import SyncTransport @@ -381,6 +381,14 @@ def snapshot( and dataset.current_revision_id is not None ): return DatasetSnapshot(dataset, selected, dataset.current_revision_id, self) + # Revisions are consecutive and newest first. Concurrent writes can + # shift their offsets, so verify the result and scan on a miss. + page = self.list_revisions( + dataset.id, limit=1, offset=dataset.current_revision - selected + ) + for candidate in page.results: + if candidate.revision == selected: + return DatasetSnapshot(dataset, selected, candidate.id, self) offset = 0 while True: page = self.list_revisions(dataset.id, limit=100, offset=offset) diff --git a/posthog/ai/evaluations/_scorers.py b/posthog/ai/evaluations/_scorers.py index 4dd2aa450..a99bd7ddf 100644 --- a/posthog/ai/evaluations/_scorers.py +++ b/posthog/ai/evaluations/_scorers.py @@ -6,7 +6,6 @@ import json from datetime import datetime from typing import Any, Callable, Literal, Sequence, TypeVar, cast, overload -from urllib.parse import parse_qs, urlsplit from uuid import UUID from ._scorer_types import ( @@ -21,7 +20,7 @@ ScorerVersion, ScorerVersionPage, ) -from ._serialization import _parse_response +from ._serialization import _integer, _next_offset, _parse_response, _text from ._tags import _response_tags, _tag_names, _tag_params from ._transport import AsyncTransport, SyncTransport @@ -47,18 +46,6 @@ def _uuid(value: str | UUID) -> str: raise ValueError("Scorer and version IDs must be UUIDs.") from None -def _integer(value: Any, *, minimum: int = 0) -> int: - if type(value) is not int or value < minimum: - raise ValueError("Expected an integer within the supported range.") - return value - - -def _text(value: Any) -> str: - if not isinstance(value, str): - raise ValueError("Expected a string.") - return value - - def _optional_text(value: Any) -> str | None: return None if value is None else _text(value) @@ -129,24 +116,18 @@ def _results(data: dict[str, Any]) -> list[dict[str, Any]]: def _scorer_page(data: dict[str, Any]) -> ScorerPage: next_url = _optional_text(data["next"]) - next_offset = None - if next_url is not None: - offsets = parse_qs(urlsplit(next_url).query).get("offset") - if offsets is None or len(offsets) != 1: - raise ValueError("Missing scorer page offset.") - next_offset = _integer(int(offsets[0])) return ScorerPage( - count=_integer(data["count"]), + count=_integer(data["count"], minimum=0), results=tuple(_scorer(row) for row in _results(data)), next=next_url, previous=_optional_text(data["previous"]), - next_offset=next_offset, + next_offset=_next_offset(next_url), ) def _version_page(data: dict[str, Any]) -> ScorerVersionPage: return ScorerVersionPage( - count=_integer(data["count"]), + count=_integer(data["count"], minimum=0), results=tuple(_version(row) for row in _results(data)), next_cursor=_optional_text(data["next_cursor"]), ) @@ -232,7 +213,7 @@ def _list_params( tag_params = _tag_params(tags, tag_match) params: dict[str, str | int] = { "limit": _integer(limit, minimum=1), - "offset": _integer(offset), + "offset": _integer(offset, minimum=0), } if search is not None: params["search"] = search diff --git a/posthog/ai/evaluations/_serialization.py b/posthog/ai/evaluations/_serialization.py index 4e09bb8fa..83aadb280 100644 --- a/posthog/ai/evaluations/_serialization.py +++ b/posthog/ai/evaluations/_serialization.py @@ -4,6 +4,7 @@ import math from datetime import datetime, timezone from typing import Any, Callable, Iterable, TypeVar, Union, cast +from urllib.parse import parse_qs, urlsplit from uuid import UUID from ._errors import EvaluationAPIError @@ -44,6 +45,28 @@ def _parse_response( ) from None +def _integer(value: Any, *, minimum: int = 1) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"Expected an integer greater than or equal to {minimum}.") + return value + + +def _text(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("Expected a string.") + return value + + +def _next_offset(value: Any, *, minimum: int = 0) -> int | None: + if value is None: + return None + # Extract only the continuation. Never send credentials to a response URL. + offsets = parse_qs(urlsplit(_text(value)).query).get("offset") + if offsets is None or len(offsets) != 1: + raise ValueError("Missing page offset.") + return _integer(int(offsets[0]), minimum=minimum) + + def _uuid(value: object) -> str: if not isinstance(value, (str, UUID)): raise ValueError("Provide a UUID or a valid UUID string.") @@ -71,14 +94,14 @@ def _identifier(value: object, field: str, limit: int = 255) -> str | None: return value -def _encode(value: object) -> bytes: +def _encode(value: object, *, sort_keys: bool = True) -> bytes: try: return json.dumps( value, ensure_ascii=False, allow_nan=False, separators=(",", ":"), - sort_keys=True, + sort_keys=sort_keys, ).encode("utf-8") except (TypeError, ValueError, UnicodeError) as exc: raise ValueError( diff --git a/posthog/test/ai/test_evaluations_datasets.py b/posthog/test/ai/test_evaluations_datasets.py index b5ed783ce..e90d34005 100644 --- a/posthog/test/ai/test_evaluations_datasets.py +++ b/posthog/test/ai/test_evaluations_datasets.py @@ -299,10 +299,26 @@ def test_filtered_dataset_pages_keep_continuation_and_normalized_tags( "archived": "true", "order_by": "updated_at", **({"tags": json.dumps(tags)} if tags is not None else {}), - "tag_match": tag_match, + **({"tag_match": tag_match} if tags is not None or tag_match != "any" else {}), } +@pytest.mark.parametrize( + "field", ["input", "expected_output", "source_output", "metadata"] +) +def test_item_json_preserves_response_key_order_without_aliasing(client, field): + evaluations, transport = client + value = {"z": [{"y": 1, "b": 2}], "a": 3} + transport.request.return_value = item_response(**{field: value}) + + item = invoke(evaluations.datasets.get_item, ITEM_ID) + copied = getattr(item, field) + assert list(copied) == ["z", "a"] + assert list(copied["z"][0]) == ["y", "b"] + value["z"][0]["y"] = 99 + assert copied["z"][0]["y"] == 1 + + def test_item_history_and_historical_get_keep_stable_and_version_ids_separate(client): evaluations, transport = client transport.request.return_value = page_response( @@ -346,10 +362,51 @@ def test_snapshot_pins_before_iteration_and_never_follows_remote_pagination_urls assert transport.request.call_args.kwargs["params"]["offset"] == 1 +def test_historical_snapshot_looks_up_expected_offset_without_scanning(client): + evaluations, transport = client + transport.request.side_effect = [ + dataset_response(current_revision=10_000), + page_response(revision_response(), count=10_000), + ] + + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID, revision=1) + + assert snapshot.revision == 1 and snapshot.revision_id == OLD_REVISION_ID + assert transport.request.call_count == 2 + assert transport.request.call_args.args == ( + "GET", + BASE + DATASET_ID + "/revisions/", + ) + assert transport.request.call_args.kwargs["params"] == {"limit": 1, "offset": 9999} + + +def test_historical_snapshot_falls_back_when_concurrent_writes_shift_offset(client): + evaluations, transport = client + transport.request.side_effect = [ + dataset_response(), + page_response(revision_response(revision=2), next_offset=3, count=4), + page_response( + revision_response(revision=4), + revision_response(revision=3), + revision_response(revision=2), + revision_response(), + ), + ] + + snapshot = invoke(evaluations.datasets.snapshot, DATASET_ID, revision=1) + + assert snapshot.revision == 1 and snapshot.revision_id == OLD_REVISION_ID + assert [call.kwargs["params"] for call in transport.request.call_args_list[1:]] == [ + {"limit": 1, "offset": 2}, + {"limit": 100, "offset": 0}, + ] + + def test_historical_snapshot_resolves_revision_uuid_across_pages(client): evaluations, transport = client transport.request.side_effect = [ dataset_response(), + page_response(count=2), page_response( revision_response(id=REVISION_ID, revision=3), next_offset=1, count=2 ), @@ -360,6 +417,9 @@ def test_historical_snapshot_resolves_revision_uuid_across_pages(client): assert snapshot.revision_id == OLD_REVISION_ID assert transport.request.call_args.kwargs["params"]["offset"] == 1 assert snapshot.dataset.current_revision == 3 + assert [ + call.kwargs["params"]["offset"] for call in transport.request.call_args_list[1:] + ] == [2, 0, 1] assert len(collect(snapshot.iter_items())) == 1 assert transport.request.call_args.kwargs["params"]["revision"] == 1 @@ -384,11 +444,12 @@ def test_missing_historical_revision_stops_once_history_passes_it(client): evaluations, transport = client transport.request.side_effect = [ dataset_response(), + page_response(revision_response(), count=2), page_response(revision_response(), next_offset=1, count=2), ] with pytest.raises(ValueError, match="revision 2 does not exist"): invoke(evaluations.datasets.snapshot, DATASET_ID, revision=2) - assert transport.request.call_count == 2 + assert transport.request.call_count == 3 def test_snapshot_propagates_partial_read_failure_and_can_be_iterated_again(client):