Skip to content

reach.queries

Load, save, and validate labeled evaluation query sets and calculate digests.

Load, save, and validate labeled evaluation query sets and calculate digests.

Origin

Bases: StrEnum

Define origins of query sets.

Source code in src/reach/queries.py
class Origin(StrEnum):
    """Define origins of query sets."""

    AUTHORED = "authored"
    GENERATED = "generated"
    IMPORTED = "imported"

QuerySet

Bases: BaseModel

Represent a collection of labeled evaluation queries targeting a catalog.

Source code in src/reach/queries.py
class QuerySet(BaseModel):
    """Represent a collection of labeled evaluation queries targeting a catalog."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    catalog_id: str = ""
    notes: str = ""
    provenance: QuerySetProvenance = Field(default_factory=QuerySetProvenance)
    queries: tuple[Query, ...]

    @model_validator(mode="before")
    @classmethod
    def _normalize_legacy_query_ids(cls, data: object) -> object:
        """Migrate legacy 'id' keys in query mappings to 'query_id'."""
        if isinstance(data, Mapping) and isinstance(
            raw_queries := data.get("queries"), (list, tuple)
        ):
            normalized_queries = [
                (
                    {"query_id": q["id"], **{k: v for k, v in q.items() if k != "id"}}
                    if isinstance(q, Mapping) and "id" in q and "query_id" not in q
                    else q
                )
                for q in raw_queries
            ]
            return {**data, "queries": normalized_queries}
        return data

    @model_validator(mode="after")
    def _assert_unique_ids(self) -> Self:
        """Validate that all query IDs within the query set are unique."""
        seen: set[str] = set()
        duplicates: set[str] = set()
        for q in self.queries:
            if q.query_id in seen:
                duplicates.add(q.query_id)
            seen.add(q.query_id)
        if duplicates:
            msg = f"duplicate query ids: {sorted(duplicates)}"
            raise ValueError(msg)
        return self

    def for_skill(self, name: str) -> tuple[Query, ...]:
        """Return queries whose expected skill or truth label matches the specified name."""
        return tuple(q for q in self.queries if name in (q.expected_skill, q.truth_label))

    def covered_skills(self) -> frozenset[str]:
        """Return target skill names covered by positive benchmark queries."""
        return frozenset(
            q.expected_skill
            for q in self.queries
            if q.expected_skill is not None and q.kind != QueryKind.NEIGHBOR_NEGATIVE
        )

    def stale_skills(self, skills: Sequence[Skill]) -> frozenset[str]:
        """Return covered skill names whose markdown body differs from recorded provenance SHA."""
        from reach.generate import skill_body_digest

        if self.provenance is None or not self.provenance.skill_digests:
            return frozenset()
        covered = self.covered_skills()
        recorded = self.provenance.skill_digests
        return frozenset(
            s.name
            for s in skills
            if s.name in covered and s.name in recorded and recorded[s.name] != skill_body_digest(s)
        )

    def out_of_sync_skills(
        self,
        skills: Sequence[Skill],
    ) -> tuple[frozenset[str], frozenset[str]]:
        """Return (missing_skills, stale_skills) relative to the provided skill corpus."""
        covered = self.covered_skills()
        missing = frozenset(s.name for s in skills if s.name not in covered)
        stale = self.stale_skills(skills)
        return missing, stale

    def without_skills(self, skill_names: Iterable[str]) -> Self:
        """Return a copy with positive and adversarial queries for `skill_names` removed."""
        target_set = frozenset(skill_names)
        if not target_set:
            return self
        retained = tuple(q for q in self.queries if not _is_generated_for_skill(q, target_set))
        return self.model_copy(update={"queries": retained})

    def with_updated_digests(
        self,
        skills: Sequence[Skill],
        *,
        drafted_targets: Iterable[str] = (),
        covered_targets: Iterable[str] = (),
        bodies_hash: str | None = None,
        extra_updates: Mapping[str, object] | None = None,
    ) -> Self:
        """Return a copy with `provenance.skill_digests` updated for `drafted_targets`."""
        effective_covered = frozenset(covered_targets) or self.covered_skills()
        base_prov = self.provenance or QuerySetProvenance(origin=Origin.GENERATED)
        updated_prov = base_prov.with_updated_digests(
            skills,
            drafted_targets=drafted_targets,
            covered_targets=effective_covered,
            bodies_hash=bodies_hash,
            extra_updates=extra_updates,
        )
        return self.model_copy(update={"provenance": updated_prov})

covered_skills

covered_skills() -> frozenset[str]

Return target skill names covered by positive benchmark queries.

Source code in src/reach/queries.py
def covered_skills(self) -> frozenset[str]:
    """Return target skill names covered by positive benchmark queries."""
    return frozenset(
        q.expected_skill
        for q in self.queries
        if q.expected_skill is not None and q.kind != QueryKind.NEIGHBOR_NEGATIVE
    )

for_skill

for_skill(name: str) -> tuple[Query, ...]

Return queries whose expected skill or truth label matches the specified name.

Source code in src/reach/queries.py
def for_skill(self, name: str) -> tuple[Query, ...]:
    """Return queries whose expected skill or truth label matches the specified name."""
    return tuple(q for q in self.queries if name in (q.expected_skill, q.truth_label))

out_of_sync_skills

out_of_sync_skills(
    skills: Sequence[Skill],
) -> tuple[frozenset[str], frozenset[str]]

Return (missing_skills, stale_skills) relative to the provided skill corpus.

Source code in src/reach/queries.py
def out_of_sync_skills(
    self,
    skills: Sequence[Skill],
) -> tuple[frozenset[str], frozenset[str]]:
    """Return (missing_skills, stale_skills) relative to the provided skill corpus."""
    covered = self.covered_skills()
    missing = frozenset(s.name for s in skills if s.name not in covered)
    stale = self.stale_skills(skills)
    return missing, stale

stale_skills

stale_skills(skills: Sequence[Skill]) -> frozenset[str]

Return covered skill names whose markdown body differs from recorded provenance SHA.

Source code in src/reach/queries.py
def stale_skills(self, skills: Sequence[Skill]) -> frozenset[str]:
    """Return covered skill names whose markdown body differs from recorded provenance SHA."""
    from reach.generate import skill_body_digest

    if self.provenance is None or not self.provenance.skill_digests:
        return frozenset()
    covered = self.covered_skills()
    recorded = self.provenance.skill_digests
    return frozenset(
        s.name
        for s in skills
        if s.name in covered and s.name in recorded and recorded[s.name] != skill_body_digest(s)
    )

with_updated_digests

with_updated_digests(
    skills: Sequence[Skill],
    *,
    drafted_targets: Iterable[str] = (),
    covered_targets: Iterable[str] = (),
    bodies_hash: str | None = None,
    extra_updates: Mapping[str, object] | None = None,
) -> Self

Return a copy with provenance.skill_digests updated for drafted_targets.

Source code in src/reach/queries.py
def with_updated_digests(
    self,
    skills: Sequence[Skill],
    *,
    drafted_targets: Iterable[str] = (),
    covered_targets: Iterable[str] = (),
    bodies_hash: str | None = None,
    extra_updates: Mapping[str, object] | None = None,
) -> Self:
    """Return a copy with `provenance.skill_digests` updated for `drafted_targets`."""
    effective_covered = frozenset(covered_targets) or self.covered_skills()
    base_prov = self.provenance or QuerySetProvenance(origin=Origin.GENERATED)
    updated_prov = base_prov.with_updated_digests(
        skills,
        drafted_targets=drafted_targets,
        covered_targets=effective_covered,
        bodies_hash=bodies_hash,
        extra_updates=extra_updates,
    )
    return self.model_copy(update={"provenance": updated_prov})

without_skills

without_skills(skill_names: Iterable[str]) -> Self

Return a copy with positive and adversarial queries for skill_names removed.

Source code in src/reach/queries.py
def without_skills(self, skill_names: Iterable[str]) -> Self:
    """Return a copy with positive and adversarial queries for `skill_names` removed."""
    target_set = frozenset(skill_names)
    if not target_set:
        return self
    retained = tuple(q for q in self.queries if not _is_generated_for_skill(q, target_set))
    return self.model_copy(update={"queries": retained})

QuerySetProvenance

Bases: BaseModel

Record creation metadata, generator settings, per-skill SHA digests, and review status.

skill_digests maps each drafted skill name to a 12-character SHA-256 digest of its SKILL.md body (excluding YAML frontmatter description so description tuning does not invalidate benchmark queries). reach query draft --sync and reach sweep compare skill_digests against the workspace corpus to detect updated skills.

Source code in src/reach/queries.py
class QuerySetProvenance(BaseModel):
    """Record creation metadata, generator settings, per-skill SHA digests, and review status.

    `skill_digests` maps each drafted skill name to a 12-character SHA-256 digest of its
    `SKILL.md` body (excluding YAML frontmatter `description` so description tuning does
    not invalidate benchmark queries). `reach query draft --sync` and `reach sweep`
    compare `skill_digests` against the workspace corpus to detect updated skills.
    """

    model_config = ConfigDict(frozen=True)

    origin: Origin = Origin.AUTHORED
    recorded_at: AwareDatetime = Field(default_factory=lambda: datetime.now(UTC))
    tool_version: str = ""
    generator_model: str = ""
    generator_arm: str = ""
    queries_per_target: int | None = None
    rivals_in_view: int | None = None
    bodies_digest: str = ""
    skill_digests: dict[SkillNameKey, SkillDigestHex] = Field(default_factory=dict)
    config_fingerprint: str = ""
    source: str = ""
    reviewed: bool | None = None
    adversarial: bool | None = None
    adversarial_per_target: int | None = None

    def with_updated_digests(
        self,
        skills: Sequence[Skill],
        *,
        drafted_targets: Iterable[str] = (),
        covered_targets: Iterable[str] = (),
        bodies_hash: str | None = None,
        extra_updates: Mapping[str, object] | None = None,
    ) -> Self:
        """Return a validated copy with updated `skill_digests` for `drafted_targets`."""
        from reach.generate import bodies_digest, skill_body_digest

        drafted_set = frozenset(drafted_targets)
        covered_set = frozenset(covered_targets) | drafted_set
        merged_digests = dict(self.skill_digests)
        for skill in skills:
            if skill.name in drafted_set or (
                skill.name in covered_set and skill.name not in merged_digests
            ):
                merged_digests[skill.name] = skill_body_digest(skill)

        payload = self.model_dump()
        if extra_updates:
            payload.update(extra_updates)
        payload["skill_digests"] = dict(sorted(merged_digests.items()))
        payload["bodies_digest"] = bodies_hash if bodies_hash is not None else bodies_digest(skills)
        return type(self).model_validate(payload)

with_updated_digests

with_updated_digests(
    skills: Sequence[Skill],
    *,
    drafted_targets: Iterable[str] = (),
    covered_targets: Iterable[str] = (),
    bodies_hash: str | None = None,
    extra_updates: Mapping[str, object] | None = None,
) -> Self

Return a validated copy with updated skill_digests for drafted_targets.

Source code in src/reach/queries.py
def with_updated_digests(
    self,
    skills: Sequence[Skill],
    *,
    drafted_targets: Iterable[str] = (),
    covered_targets: Iterable[str] = (),
    bodies_hash: str | None = None,
    extra_updates: Mapping[str, object] | None = None,
) -> Self:
    """Return a validated copy with updated `skill_digests` for `drafted_targets`."""
    from reach.generate import bodies_digest, skill_body_digest

    drafted_set = frozenset(drafted_targets)
    covered_set = frozenset(covered_targets) | drafted_set
    merged_digests = dict(self.skill_digests)
    for skill in skills:
        if skill.name in drafted_set or (
            skill.name in covered_set and skill.name not in merged_digests
        ):
            merged_digests[skill.name] = skill_body_digest(skill)

    payload = self.model_dump()
    if extra_updates:
        payload.update(extra_updates)
    payload["skill_digests"] = dict(sorted(merged_digests.items()))
    payload["bodies_digest"] = bodies_hash if bodies_hash is not None else bodies_digest(skills)
    return type(self).model_validate(payload)

load_query_set

load_query_set(
    path: Path | str, *, catalog_id: str = "all"
) -> QuerySet

Load and validate a QuerySet from a JSON, YAML, JSONL, or CSV file.

Source code in src/reach/queries.py
def load_query_set(
    path: Path | str,
    *,
    catalog_id: str = "all",
) -> QuerySet:
    """Load and validate a QuerySet from a JSON, YAML, JSONL, or CSV file."""
    from reach.config import resolve_path

    resolved = resolve_path(path)
    content = resolved.read_text(encoding="utf-8")
    suffix = resolved.suffix.lower()
    if suffix in (".jsonl", ".csv"):
        from reach.exchange import Exchange, import_query_set

        fmt = Exchange.JSONL if suffix == ".jsonl" else Exchange.CSV
        return import_query_set(content, fmt, catalog_id=catalog_id, source=str(resolved))
    if suffix in (".yaml", ".yml"):
        import yaml

        return _apply_catalog_id_fallback(
            QuerySet.model_validate(yaml.safe_load(content)), catalog_id
        )
    if suffix == ".json":
        return _apply_catalog_id_fallback(read_model(QuerySet, resolved), catalog_id)
    try:
        return _apply_catalog_id_fallback(read_model(QuerySet, resolved), catalog_id)
    except (ValueError, ValidationError):
        from reach.exchange import Exchange, import_query_set

        return import_query_set(
            content, Exchange.JSONL, catalog_id=catalog_id, source=str(resolved)
        )

query_set_digest

query_set_digest(query_set: QuerySet) -> str

Compute a deterministic 12-character SHA-256 digest of query set content.

Source code in src/reach/queries.py
def query_set_digest(query_set: QuerySet) -> str:
    """Compute a deterministic 12-character SHA-256 digest of query set content."""
    return _digest_queries(query_set.catalog_id, _query_rows(query_set))

save_query_set

save_query_set(
    query_set: QuerySet,
    path: Path | str,
    *,
    fmt: str | None = None,
    mapping: Mapping[str, str] | None = None,
) -> Path

Serialize a QuerySet instance to disk in JSON, JSONL, or CSV format.

Source code in src/reach/queries.py
def save_query_set(
    query_set: QuerySet,
    path: Path | str,
    *,
    fmt: str | None = None,
    mapping: Mapping[str, str] | None = None,
) -> Path:
    """Serialize a QuerySet instance to disk in JSON, JSONL, or CSV format."""
    resolved = Path(path).expanduser().resolve()
    resolved.parent.mkdir(parents=True, exist_ok=True)
    target_fmt = fmt or (
        "jsonl"
        if resolved.suffix.lower() == ".jsonl"
        else "csv"
        if resolved.suffix.lower() == ".csv"
        else "json"
    )
    if target_fmt == "jsonl":
        from reach.exchange import Exchange, export_query_set

        resolved.write_text(
            export_query_set(query_set, Exchange.JSONL, mapping=mapping),
            encoding="utf-8",
        )
        return resolved
    if target_fmt == "csv":
        from reach.exchange import Exchange, export_query_set

        resolved.write_text(
            export_query_set(query_set, Exchange.CSV, mapping=mapping),
            encoding="utf-8",
        )
        return resolved
    return write_model(query_set, resolved)