evalkit-py 0.2.0.dev12__py3-none-any.whl

evalkit/core/scorecard.py

@@ -0,0 +1,255 @@
+"""Aggregation of evaluation rows into the scorecard report.
+
+A row's bucket is resolved when it is written (`bucket_for` plus the
+exclusions below); bucket means are computed here. `answer_rubric_score` is the
+`content` bucket's only member; every other name the rubric evaluator
+publishes is a diagnostic on it and stays out of the bucket — recall =
+coverage x accuracy identically, so taking all three would double-count every
+missed rubric; the per-type rates partition the very rubrics the score
+already counts; and citation is a separate axis whose denominator exists only
+on cases that ask for sources.
+"""
+
+from __future__ import annotations
+
+import statistics
+from collections.abc import Iterable
+
+from pydantic import BaseModel, Field
+
+from evalkit.store.base import AggregateRows
+
+GAUGE_PREFIX = "gauge_"
+
+ANSWER_RUBRIC_ORDER: tuple[str, ...] = (
+    "answer_rubric_score",
+    "answer_rubric_coverage",
+    "answer_rubric_accuracy",
+    "answer_rubric_factual_retrieval",
+    "answer_rubric_derived_computation",
+    "answer_rubric_analysis_interpretation",
+    "answer_rubric_citation",
+)
+
+ANSWER_RUBRIC_DIAGNOSTICS = frozenset(ANSWER_RUBRIC_ORDER) - {"answer_rubric_score"}
+
+EXCLUDED_FROM_BUCKETS = ANSWER_RUBRIC_DIAGNOSTICS
+
+BUCKET_ORDER: tuple[str, ...] = (
+    "content",
+    "trajectory",
+    "presentation",
+    "efficiency",
+    "custom",
+)
+
+BUCKETS: dict[str, frozenset[str]] = {
+    "content": frozenset({"answer_rubric_score"}),
+    "trajectory": frozenset(
+        {"trajectory_rubric_score", "tool_trajectory_match", "required_tools_called"}
+    ),
+    "presentation": frozenset({"presentation_judge"}),
+    "efficiency": frozenset({"max_response_time"}),
+}
+
+_BUCKET_BY_EVALUATOR: dict[str, str] = {
+    name: bucket for bucket, members in BUCKETS.items() for name in members
+}
+
+# One representative evaluator per bucket the inline score blocks show: every
+# bucket with a single member. `custom` has no fixed member and is left out.
+BUCKET_EVALUATOR = {
+    bucket: next(iter(members)) for bucket, members in BUCKETS.items() if len(members) == 1
+}
+COST_GAUGE = "gauge_agent_cost_usd"
+
+
+def is_gauge(name: str) -> bool:
+    return name.startswith(GAUGE_PREFIX)
+
+
+def gauge_places(name: str) -> int:
+    """Display precision for a gauge value — dollar gauges need decimals,
+    count/token gauges render as whole numbers."""
+    return 4 if name.endswith("_cost_usd") else 0
+
+
+def bucket_for(name: str, source: str = "builtin") -> str | None:
+    if source.startswith("custom:"):
+        return "custom"
+    return _BUCKET_BY_EVALUATOR.get(name)
+
+
+def order_evaluator_columns(present: Iterable[str]) -> list[str]:
+    present = set(present)
+    head = [name for name in ANSWER_RUBRIC_ORDER if name in present]
+    tail = sorted(present.difference(head))
+    return head + tail
+
+
+class CaseRow(BaseModel):
+    case_id: str
+    scores: dict[str, float | None] = Field(default_factory=dict)
+    explanations: dict[str, str | None] = Field(default_factory=dict)
+    gauges: dict[str, float | None] = Field(default_factory=dict)
+    errors: dict[str, str] = Field(default_factory=dict)
+    run_errors: list[str] = Field(default_factory=list)
+    repetitions: int = 1
+    overall: float | None = None
+    baseline_overall: float | None = None
+    delta: float | None = None
+
+
+class EvaluatorStat(BaseModel):
+    name: str
+    score: float | None = None
+    cases: int = 0
+    errors: int = 0
+    bucket: str | None = None
+    source: str = "builtin"
+
+
+class GaugeStat(BaseModel):
+    name: str
+    mean: float | None = None
+    median: float | None = None
+    count: int = 0
+
+
+class Report(BaseModel):
+    experiment_id: str
+    name: str
+    dataset: str
+    metadata: dict = Field(default_factory=dict)
+    cases: list[CaseRow]
+    evaluators: list[EvaluatorStat]
+    buckets: dict[str, float | None]
+    gauges: list[GaugeStat]
+    evaluator_columns: list[str]
+    baseline_id: str | None = None
+
+
+def _mean(values: list[float]) -> float | None:
+    return sum(values) / len(values) if values else None
+
+
+def _aggregate_rows(agg: AggregateRows) -> dict[str, CaseRow]:
+    rows = {
+        case_id: CaseRow(case_id=case_id, repetitions=n)
+        for case_id, n in agg.repetitions_seen.items()
+    }
+    bucketed: dict[str, list[float]] = {}
+    for case_id, name, mean, explanation, bucket in agg.case_scores:
+        rows[case_id].scores[name] = mean
+        rows[case_id].explanations[name] = explanation
+        if bucket is not None:
+            bucketed.setdefault(case_id, []).append(mean)
+    for case_id, name, mean in agg.case_gauges:
+        rows[case_id].gauges[name] = mean
+    for case_id, name, error in agg.case_errors:
+        rows[case_id].errors[name] = error
+    for case_id, error in agg.run_errors:
+        rows[case_id].run_errors.append(error)
+    for case_id, row in rows.items():
+        row.overall = _mean(bucketed.get(case_id, []))
+    return rows
+
+
+def assemble_report(agg: AggregateRows, baseline: AggregateRows | None = None) -> Report:
+    exp = agg.experiment
+    return _report(
+        Report(
+            experiment_id=str(exp.id),
+            name=f"#{exp.number}",
+            dataset=exp.dataset,
+            metadata={
+                **exp.config.model_dump(mode="json"),
+                "agent": exp.agent,
+                "created_by": exp.created_by,
+                "status": exp.status.value,
+            },
+            cases=[],
+            evaluators=[],
+            buckets={},
+            gauges=[],
+            evaluator_columns=[],
+        ),
+        _aggregate_rows(agg),
+        agg.sources,
+        {name: bucket for _, name, _, _, bucket in agg.case_scores},
+        None if baseline is None else (str(baseline.experiment.id), _aggregate_rows(baseline)),
+    )
+
+
+def _report(
+    header: Report,
+    rows: dict[str, CaseRow],
+    sources: dict[str, str],
+    stored_buckets: dict[str, str | None],
+    baseline: tuple[str, dict[str, CaseRow]] | None,
+) -> Report:
+    evaluator_names = {
+        name
+        for row in rows.values()
+        for name in list(row.scores) + list(row.errors)
+        if not is_gauge(name)
+    }
+    evaluators = []
+    for name in order_evaluator_columns(evaluator_names):
+        case_values = [
+            value for row in rows.values() if (value := row.scores.get(name)) is not None
+        ]
+        evaluators.append(
+            EvaluatorStat(
+                name=name,
+                score=_mean(case_values),
+                cases=len(case_values),
+                errors=sum(1 for row in rows.values() if name in row.errors),
+                bucket=stored_buckets.get(name),
+                source=sources.get(name, "builtin"),
+            )
+        )
+
+    buckets: dict[str, float | None] = {}
+    for bucket in BUCKET_ORDER:
+        members = [
+            stat.score
+            for stat in evaluators
+            if stat.bucket == bucket
+            and stat.score is not None
+            and stat.name not in EXCLUDED_FROM_BUCKETS
+        ]
+        buckets[bucket] = _mean(members)
+
+    gauge_names = {name for row in rows.values() for name in row.gauges}
+    gauges = []
+    for name in sorted(gauge_names):
+        values = [value for row in rows.values() if (value := row.gauges.get(name)) is not None]
+        gauges.append(
+            GaugeStat(
+                name=name,
+                mean=_mean(values),
+                median=statistics.median(values) if values else None,
+                count=len(values),
+            )
+        )
+
+    baseline_id = None
+    if baseline is not None:
+        baseline_id, baseline_rows = baseline
+        for case_id, row in rows.items():
+            base = baseline_rows.get(case_id)
+            if base is not None and base.overall is not None and row.overall is not None:
+                row.baseline_overall = base.overall
+                row.delta = round(row.overall - base.overall, 6)
+
+    return header.model_copy(
+        update={
+            "cases": sorted(rows.values(), key=lambda r: r.case_id),
+            "evaluators": evaluators,
+            "buckets": buckets,
+            "gauges": gauges,
+            "evaluator_columns": order_evaluator_columns(evaluator_names),
+            "baseline_id": baseline_id,
+        }
+    )

evalkit/core/traces/convert.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+import json
+from bisect import bisect_right
+from datetime import datetime
+from typing import Any, NamedTuple
+
+from evalkit.contracts.atif import (
+    AtifAgent,
+    FinalMetrics,
+    Observation,
+    ObservationResult,
+    Step,
+    StepMetrics,
+    SubagentTrajectoryRef,
+    ToolCall,
+    Trajectory,
+)
+from evalkit.contracts.esf import EsfSpan, MappingDocument
+from evalkit.core.traces.mapping import gen_ai_messages
+
+ATIF_VERSION = "ATIF-v1.8"
+
+
+def _as_text(value: Any) -> str | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, default=str)
+
+
+def _as_dict(value: dict[str, Any] | str) -> dict[str, Any]:
+    if isinstance(value, dict):
+        return value
+    try:
+        parsed = json.loads(value)
+    except ValueError:
+        return {"raw": value}
+    return parsed if isinstance(parsed, dict) else {"raw": value}
+
+
+def _output_text(value: Any) -> str | None:
+    messages = gen_ai_messages(value)
+    if messages is None:
+        return _as_text(value)
+    texts = [m.text for m in messages if m.role == "assistant" and m.text]
+    return texts[-1] if texts else None
+
+
+class _Tree:
+    def __init__(self, spans: list[EsfSpan]):
+        self.by_id = {s.span_id: s for s in spans}
+        self.children: dict[str | None, list[EsfSpan]] = {}
+        for s in sorted(spans, key=lambda x: x.started_at):
+            parent = s.parent_id if s.parent_id in self.by_id else None
+            self.children.setdefault(parent, []).append(s)
+
+    def roots(self) -> list[EsfSpan]:
+        return self.children.get(None, [])
+
+    def subtree(self, span_id: str) -> list[EsfSpan]:
+        out: list[EsfSpan] = []
+        stack = [span_id]
+        while stack:
+            current = stack.pop()
+            for child in self.children.get(current, []):
+                out.append(child)
+                stack.append(child.span_id)
+        return out
+
+    def ancestors(self, span: EsfSpan) -> list[EsfSpan]:
+        out: list[EsfSpan] = []
+        parent = span.parent_id
+        while parent in self.by_id:
+            out.append(self.by_id[parent])
+            parent = self.by_id[parent].parent_id
+        return out
+
+
+def _answer_of(root: EsfSpan | None, spans: list[EsfSpan], rule: str) -> str | None:
+    if rule == "root":
+        return _output_text(root.output) if root is not None else None
+    name, _, key = rule.partition(".")
+    ordered = sorted(spans, key=lambda x: x.started_at)
+    named = next((s for s in ordered if s.name == name), None)
+    if named is None:
+        return None
+    if key and isinstance(named.output, dict):
+        return _as_text(named.output.get(key))
+    return _output_text(named.output)
+
+
+def _window(llms: list[EsfSpan], at: datetime) -> int | None:
+    index = bisect_right([s.started_at for s in llms], at) - 1
+    return index if index >= 0 else None
+
+
+def _result(
+    span: EsfSpan, call_id: str | None, refs: dict[str, list[SubagentTrajectoryRef]]
+) -> ObservationResult:
+    return ObservationResult(
+        source_call_id=call_id,
+        content=_output_text(span.output) if span.kind == "agent" else _as_text(span.output),
+        subagent_trajectory_ref=refs.get(span.span_id),
+        extra={"is_error": span.error is not None},
+    )
+
+
+def _place(
+    steps: list[Step], outcomes: list[EsfSpan], llms: list[EsfSpan]
+) -> tuple[dict[int, list[tuple[str | None, EsfSpan]]], list[EsfSpan]]:
+    owner = {c.tool_call_id: k for k, step in enumerate(steps) for c in step.tool_calls or []}
+    placed: dict[int, list[tuple[str | None, EsfSpan]]] = {}
+    claimed: set[str] = set()
+    by_time: list[EsfSpan] = []
+    for span in sorted(outcomes, key=lambda s: s.started_at):
+        call_id = span.tool_call_id
+        if call_id and call_id in owner and call_id not in claimed:
+            claimed.add(call_id)
+            placed.setdefault(owner[call_id], []).append((call_id, span))
+        else:
+            by_time.append(span)
+    pending = {
+        k: [c.tool_call_id for c in step.tool_calls or [] if c.tool_call_id not in claimed]
+        for k, step in enumerate(steps)
+    }
+    unplaced: list[EsfSpan] = []
+    for span in by_time:
+        k = _window(llms, span.started_at)
+        if k is None:
+            unplaced.append(span)
+            continue
+        placed.setdefault(k, []).append((pending[k].pop(0) if pending[k] else None, span))
+    return placed, unplaced
+
+
+def _steps(
+    llms: list[EsfSpan],
+    outcomes: list[EsfSpan],
+    refs: dict[str, list[SubagentTrajectoryRef]],
+) -> tuple[list[Step], list[EsfSpan]]:
+    steps: list[Step] = []
+    agent_steps: list[Step] = []
+    llms = sorted(llms, key=lambda s: s.started_at)
+    for i, call in enumerate(llms):
+        if i == 0 and call.input_messages:
+            for m in call.input_messages:
+                if m.role == "user":
+                    steps.append(Step(step_id=0, source="user", message=m.text or ""))
+                elif m.role == "system":
+                    steps.append(Step(step_id=0, source="system", message=m.text or ""))
+        assistant = next(
+            (m for m in reversed(call.output_messages or []) if m.role == "assistant"), None
+        )
+        tool_calls = [
+            ToolCall(tool_call_id=c.id, function_name=c.name, arguments=_as_dict(c.arguments))
+            for c in (assistant.tool_calls if assistant else None) or []
+        ]
+        step = Step(
+            step_id=0,
+            timestamp=call.started_at.isoformat(),
+            source="agent",
+            message=(assistant.text if assistant else None) or "",
+            reasoning_content=assistant.reasoning if assistant else None,
+            model_name=call.model,
+            metrics=StepMetrics(
+                prompt_tokens=call.prompt_tokens,
+                completion_tokens=call.completion_tokens,
+                cached_tokens=call.cached_tokens,
+                cost_usd=call.cost_usd,
+                extra={"cache_creation_input_tokens": call.cache_write_tokens}
+                if call.cache_write_tokens
+                else None,
+            ),
+            tool_calls=tool_calls or None,
+        )
+        steps.append(step)
+        agent_steps.append(step)
+    placed, unplaced = _place(agent_steps, outcomes, llms)
+    for k, pairs in placed.items():
+        order = {c.tool_call_id: j for j, c in enumerate(agent_steps[k].tool_calls or [])}
+        pairs.sort(
+            key=lambda p: (order.get(p[0], len(order)) if p[0] else len(order), p[1].started_at)
+        )
+        agent_steps[k].observation = Observation(
+            results=[_result(span, call_id, refs) for call_id, span in pairs]
+        )
+    return steps, unplaced
+
+
+def _add[Num: (int, float)](total: Num | None, value: Num | None) -> Num | None:
+    if total is None:
+        return value
+    if value is None:
+        return total
+    return total + value
+
+
+def _final_metrics(steps: list[Step], subagents: list[Trajectory]) -> FinalMetrics:
+    final = FinalMetrics(total_steps=len(steps))
+    for m in (s.metrics for s in steps if s.metrics is not None):
+        final.total_prompt_tokens = _add(final.total_prompt_tokens, m.prompt_tokens)
+        final.total_completion_tokens = _add(final.total_completion_tokens, m.completion_tokens)
+        final.total_cached_tokens = _add(final.total_cached_tokens, m.cached_tokens)
+        final.total_cost_usd = _add(final.total_cost_usd, m.cost_usd)
+    for sub in (s.final_metrics for s in subagents if s.final_metrics is not None):
+        final.total_prompt_tokens = _add(final.total_prompt_tokens, sub.total_prompt_tokens)
+        final.total_completion_tokens = _add(
+            final.total_completion_tokens, sub.total_completion_tokens
+        )
+        final.total_cached_tokens = _add(final.total_cached_tokens, sub.total_cached_tokens)
+        final.total_cost_usd = _add(final.total_cost_usd, sub.total_cost_usd)
+    return final
+
+
+class _Split(NamedTuple):
+    tree: _Tree
+    original_root: EsfSpan | None
+    main_root: EsfSpan | None
+    top_nested: list[EsfSpan]
+    main: list[EsfSpan]
+
+
+def _split(kept: list[EsfSpan]) -> _Split:
+    tree = _Tree(kept)
+    roots = tree.roots()
+    original_root = roots[0] if len(roots) == 1 else None
+    main_root = original_root
+    if main_root is None or main_root.kind != "agent":
+        top_agents = [
+            s
+            for s in kept
+            if s.kind == "agent" and not any(a.kind == "agent" for a in tree.ancestors(s))
+        ]
+        if top_agents:
+            main_root = min(top_agents, key=lambda s: s.started_at)
+    nested_agent_ids = {s.span_id for s in kept if s.kind == "agent" and s is not main_root}
+    top_nested = [
+        s
+        for s in kept
+        if s.span_id in nested_agent_ids
+        and not any(a.span_id in nested_agent_ids for a in tree.ancestors(s))
+    ]
+    nested_ids = {n.span_id for n in top_nested}
+    nested_ids |= {s.span_id for n in top_nested for s in tree.subtree(n.span_id)}
+    main = [s for s in kept if s.span_id not in nested_ids]
+    return _Split(tree, original_root, main_root, top_nested, main)
+
+
+def _main_llms(main: list[EsfSpan]) -> list[EsfSpan]:
+    return [s for s in main if s.kind == "llm"]
+
+
+def _excluded(mapping: MappingDocument | None) -> set[str]:
+    return set(mapping.exclude_spans) if mapping is not None else set()
+
+
+def first_llm_spans(spans: list[EsfSpan], mapping: MappingDocument | None) -> list[EsfSpan]:
+    split = _split([s for s in spans if s.name not in _excluded(mapping)])
+    llms = sorted(_main_llms(split.main), key=lambda s: s.started_at)
+    nested = [
+        first
+        for n in split.top_nested
+        for first in first_llm_spans([n, *split.tree.subtree(n.span_id)], mapping)
+    ]
+    return llms[:1] + nested
+
+
+def to_atif(
+    spans: list[EsfSpan],
+    mapping: MappingDocument | None,
+    *,
+    agent_name: str,
+    session_id: str | None,
+    mapping_version: int | None,
+    incomplete_ended_at: datetime | None = None,
+) -> Trajectory:
+    kept = [s for s in spans if s.name not in _excluded(mapping)]
+    incomplete = sum(1 for s in kept if s.ended_at is None)
+    if incomplete_ended_at is not None:
+        kept = [
+            s.model_copy(update={"ended_at": incomplete_ended_at}) if s.ended_at is None else s
+            for s in kept
+        ]
+    split = _split(kept)
+    tree, original_root, main_root = split.tree, split.original_root, split.main_root
+    top_nested, main = split.top_nested, split.main
+    tools = [
+        s for s in main if s.kind == "tool" and not any(a.kind == "tool" for a in tree.ancestors(s))
+    ]
+    tool_ids = {t.span_id for t in tools}
+    refs: dict[str, list[SubagentTrajectoryRef]] = {}
+    outcomes = list(tools)
+    for n in top_nested:
+        ref = SubagentTrajectoryRef(trajectory_id=f"{n.trace_id}:{n.span_id}")
+        spawner = next((a for a in tree.ancestors(n) if a.span_id in tool_ids), n)
+        refs.setdefault(spawner.span_id, []).append(ref)
+        if spawner is n:
+            outcomes.append(n)
+    steps, unplaced = _steps(_main_llms(main), outcomes, refs)
+    rule = mapping.answer if mapping is not None else "root"
+    answer = _answer_of(original_root or main_root, kept, rule)
+    if answer and (not steps or steps[-1].tool_calls):
+        steps.append(Step(step_id=0, source="agent", message=answer))
+    for i, step in enumerate(steps, 1):
+        step.step_id = i
+    subagents = [
+        to_atif(
+            [n, *tree.subtree(n.span_id)],
+            mapping,
+            agent_name=n.agent_name or n.name,
+            session_id=session_id,
+            mapping_version=mapping_version,
+        )
+        for n in top_nested
+    ]
+    ordered_main = sorted(main, key=lambda x: x.started_at)
+    first_model = next((s.model for s in ordered_main if s.kind == "llm" and s.model), None)
+    extra: dict[str, Any] = {"trajectory_source": "otel", "mapping_version": mapping_version}
+    if incomplete:
+        extra["incomplete_spans"] = incomplete
+    unplaced_tools = [s.name for s in unplaced if s.kind == "tool"]
+    if unplaced_tools:
+        extra["unplaced_tool_spans"] = unplaced_tools
+    anchor = main_root or (kept[0] if kept else None)
+    return Trajectory(
+        schema_version=ATIF_VERSION,
+        session_id=session_id,
+        trajectory_id=f"{anchor.trace_id}:{anchor.span_id}" if anchor else None,
+        agent=AtifAgent(name=agent_name, model_name=first_model),
+        steps=steps,
+        final_metrics=_final_metrics(steps, subagents),
+        subagent_trajectories=subagents or None,
+        extra=extra,
+    )

evalkit/core/traces/importer.py

@@ -0,0 +1,113 @@
+from datetime import datetime
+from typing import Protocol
+
+from pydantic import BaseModel
+
+from evalkit.contracts.enums import TraceImportStatus
+from evalkit.contracts.esf import MappingDocument
+from evalkit.core.runner import error_text
+from evalkit.core.traces.mapping import session_id_of
+from evalkit.core.traces.raw import RawSpan
+from evalkit.core.traces.sources import TraceSourceClient
+
+IMPORT_LEASE_SECONDS = 180
+IMPORT_STOPPED_ERROR = "the server stopped while importing"
+ACTIVE_IMPORT = frozenset({TraceImportStatus.RUNNING, TraceImportStatus.STOPPING})
+
+
+class ImportJob(BaseModel):
+    id: int
+    agent_id: int
+    trace_source_id: int
+    project: str
+    until: datetime
+    resume_from: datetime
+    resume_trace_ids: list[str]
+
+
+class ImportStore(Protocol):
+    async def job(self, import_id: int) -> ImportJob: ...
+
+    async def advance(
+        self,
+        import_id: int,
+        holder: str,
+        *,
+        resume_from: datetime,
+        resume_trace_ids: list[str],
+        traces: int,
+        spans: int,
+        sample_session: str | None,
+    ) -> TraceImportStatus | None: ...
+
+    async def renew(self, import_id: int, holder: str) -> TraceImportStatus | None: ...
+
+    async def finish(
+        self, import_id: int, holder: str, status: TraceImportStatus, error: str | None = None
+    ) -> None: ...
+
+
+class SpanWriter(Protocol):
+    async def upsert(
+        self, agent_id: int, spans: list[RawSpan], *, external_trace_source_id: int | None = None
+    ) -> int: ...
+
+
+async def run_import(
+    import_id: int,
+    holder: str,
+    client: TraceSourceClient,
+    store: ImportStore,
+    spans: SpanWriter,
+    mapping: MappingDocument | None,
+) -> None:
+    job = await store.job(import_id)
+    boundary, written = job.resume_from, set(job.resume_trace_ids)
+    async for page in client.trace_pages(job.project, job.resume_from, job.until):
+        fresh = [tid for tid, start in page.traces if start != boundary or tid not in written]
+        span_count, sample = 0, None
+        if fresh:
+            async for batch in client.trace_spans(job.project, fresh):
+                for span in batch:
+                    span.session_id_hint = session_id_of(span.attributes, mapping)
+                await spans.upsert(
+                    job.agent_id, batch, external_trace_source_id=job.trace_source_id
+                )
+                span_count += len(batch)
+                roots = [s.session_id_hint for s in batch if s.parent_span_id is None]
+                sample = next((r for r in reversed(roots) if r), sample)
+                if await store.renew(import_id, holder) is None:
+                    return
+        at_last = {tid for tid, start in page.traces if start == page.last_start}
+        written = written | at_last if page.last_start == boundary else at_last
+        boundary = page.last_start
+        state = await store.advance(
+            import_id,
+            holder,
+            resume_from=boundary,
+            resume_trace_ids=sorted(written),
+            traces=len(fresh),
+            spans=span_count,
+            sample_session=sample,
+        )
+        if state is None:
+            return
+        if state is TraceImportStatus.STOPPING:
+            await store.finish(import_id, holder, TraceImportStatus.STOPPED)
+            return
+    await store.finish(import_id, holder, TraceImportStatus.DONE)
+
+
+async def guarded_import(
+    import_id: int,
+    holder: str,
+    client: TraceSourceClient,
+    store: ImportStore,
+    spans: SpanWriter,
+    mapping: MappingDocument | None,
+) -> None:
+    try:
+        await run_import(import_id, holder, client, store, spans, mapping)
+    except Exception as exc:
+        await store.finish(import_id, holder, TraceImportStatus.FAILED, error_text(exc))
+        raise