docent-python 0.1.50a0__tar.gz → 0.1.52a0__tar.gz

{docent_python-0.1.50a0 → docent_python-0.1.52a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.50a0
+Version: 0.1.52a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues

docent_python-0.1.52a0/docent/data_models/feedback.py

@@ -0,0 +1,393 @@
+"""Data structures for run-centric feedback elicitation and user context inference."""
+
+import json
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+from docent.data_models.citation import InlineCitation
+from docent.judges.util.voting import OutputDistribution
+
+
+def _stable_json(value: Any) -> str:
+    return json.dumps(value, sort_keys=True)
+
+
+def _indent_lines(lines: list[str], indent: int) -> list[str]:
+    prefix = " " * max(0, indent)
+    return [f"{prefix}{line}" if line else "" for line in lines]
+
+
+def _tag_block(tag: str, body_lines: list[str], indent: int) -> list[str]:
+    lines = [f"<{tag}>"]
+    if body_lines:
+        lines.extend(_indent_lines(body_lines, indent))
+    else:
+        lines.extend(_indent_lines(["N/A"], indent))
+    lines.append(f"</{tag}>")
+    return lines
+
+
+def _text_or_na(text: str | None) -> str:
+    if text is None:
+        return "N/A"
+    stripped = text.strip()
+    return stripped if stripped else "N/A"
+
+
+def _text_lines_or_na(text: str | None) -> list[str]:
+    return _text_or_na(text).splitlines()
+
+
+def _render_citations_block(citations: list[InlineCitation], indent: int) -> list[str]:
+    citation_payload = [citation.model_dump(mode="json") for citation in citations]
+    citation_text = _stable_json(citation_payload) if citation_payload else "N/A"
+    return _tag_block("Citations", [citation_text], indent)
+
+
+def _render_user_distribution_block(
+    user_distribution: OutputDistribution | None,
+    indent: int,
+) -> list[str]:
+    distribution_text = (
+        _stable_json(user_distribution.model_dump(mode="json"))
+        if user_distribution is not None
+        else "N/A"
+    )
+    return _tag_block("Estimated user distribution p_u", [distribution_text], indent)
+
+
+def _render_user_distribution_reasoning_block(
+    reasoning: str | None,
+    reasoning_citations: list[InlineCitation] | None,
+    indent: int,
+) -> list[str]:
+    body_lines = _text_lines_or_na(reasoning)
+    body_lines.extend(_render_citations_block(reasoning_citations or [], indent))
+    return _tag_block("p_u reasoning", body_lines, indent)
+
+
+class LabelingRequestFocusItem(BaseModel):
+    """Specific rubric-related question the human labeler should inspect."""
+
+    question: str
+    citations: list[InlineCitation] = Field(default_factory=list[InlineCitation])
+    sample_answers: list[str] = Field(default_factory=list[str])
+
+    def to_str(self, indent: int = 0) -> str:
+        """Render focus item in a deterministic LLM-facing format."""
+        lines: list[str] = []
+
+        # Render the question and its citations as one nested block.
+        question_lines = _text_lines_or_na(self.question)
+        question_lines.extend(_render_citations_block(self.citations, indent))
+        lines.extend(_tag_block("Question", question_lines, indent))
+
+        sample_answers_lines = (
+            [
+                f"Answer {sample_idx}: {sample_answer}"
+                for sample_idx, sample_answer in enumerate(self.sample_answers, start=1)
+            ]
+            if self.sample_answers
+            else ["N/A"]
+        )
+        lines.extend(_tag_block("Sample Answers", sample_answers_lines, indent))
+        return "\n".join(lines)
+
+
+class QAPair(BaseModel):
+    """A single review-focus answer captured for one run."""
+
+    # What the user was shown
+    focus_index: int
+
+    # What the user responded
+    answer: str
+    explanation: str | None = None
+
+    # The user could have skipped this question and provided nothing
+    status: Literal["answered", "skipped"]
+    timestamp: datetime = Field(default_factory=datetime.now)
+
+    def to_str(self, labeling_request: "LabelingRequest", indent: int = 0) -> str:
+        """Render QA pair in a deterministic LLM-facing format."""
+        if self.focus_index < 0 or self.focus_index >= len(labeling_request.review_focus):
+            raise ValueError(
+                f"focus_index={self.focus_index} is out of bounds for review_focus length "
+                f"{len(labeling_request.review_focus)}"
+            )
+        focus_item = labeling_request.review_focus[self.focus_index]
+        lines = focus_item.to_str(indent=indent).splitlines()
+        lines.append(f"User answer: {_text_or_na(self.answer)}")
+        lines.append(f"User explanation: {_text_or_na(self.explanation)}")
+        return "\n".join(lines)
+
+
+class LabelingRequest(BaseModel):
+    """Structured labeling request shown to the user."""
+
+    title: str
+    review_context: str
+    review_context_citations: list[InlineCitation] = Field(default_factory=list[InlineCitation])
+    review_focus: list[LabelingRequestFocusItem] = Field(
+        default_factory=list[LabelingRequestFocusItem]
+    )
+    user_distribution: OutputDistribution | None = None
+    user_distribution_reasoning: str | None = None
+
+    def to_str(self, indent: int = 0) -> str:
+        """Render labeling request in a deterministic LLM-facing format."""
+        body_lines: list[str] = [f"Title: {_text_or_na(self.title)}"]
+
+        review_context_lines = _text_lines_or_na(self.review_context)
+        review_context_lines.extend(_render_citations_block(self.review_context_citations, indent))
+        body_lines.extend(_tag_block("Review Context", review_context_lines, indent))
+
+        review_focus_lines: list[str] = []
+        if self.review_focus:
+            for focus_idx, focus_item in enumerate(self.review_focus, start=1):
+                focus_lines = focus_item.to_str(indent=indent).splitlines()
+                review_focus_lines.extend(_tag_block(f"Focus {focus_idx}", focus_lines, indent))
+        else:
+            review_focus_lines.append("N/A")
+        body_lines.extend(_tag_block("Review Focus", review_focus_lines, indent))
+
+        body_lines.extend(_render_user_distribution_block(self.user_distribution, indent))
+        body_lines.extend(
+            _render_user_distribution_reasoning_block(
+                self.user_distribution_reasoning,
+                reasoning_citations=None,
+                indent=indent,
+            )
+        )
+
+        lines = _tag_block("Labeling Request", body_lines, indent)
+        return "\n".join(lines)
+
+
+class LabeledRun(BaseModel):
+    """A human label for one agent run."""
+
+    agent_run_id: str
+    timestamp: datetime = Field(default_factory=datetime.now)
+
+    # What the user responded
+    label_value: dict[str, Any]
+    explanation: str | None = None
+
+    def to_str(
+        self,
+        labeling_request: LabelingRequest | None = None,
+        indent: int = 0,
+    ) -> str:
+        """Render user label in a deterministic LLM-facing format."""
+        body_lines = [
+            f"User label: {_stable_json(self.label_value)}",
+            f"User explanation: {_text_or_na(self.explanation)}",
+        ]
+        if labeling_request is None:
+            return "\n".join(_tag_block("Label", body_lines, indent))
+
+        body_lines.extend(
+            _render_user_distribution_block(labeling_request.user_distribution, indent)
+        )
+        body_lines.extend(
+            _render_user_distribution_reasoning_block(
+                labeling_request.user_distribution_reasoning,
+                reasoning_citations=None,
+                indent=indent,
+            )
+        )
+        return "\n".join(_tag_block("Label", body_lines, indent))
+
+
+class AgentRunFeedbackContext(BaseModel):
+    """All feedback collected for a single agent run."""
+
+    feedback_context_id: str | None = None
+    agent_run_id: str
+    round: int
+    created_at: datetime = Field(default_factory=datetime.now)
+    last_updated: datetime = Field(default_factory=datetime.now)
+
+    # What the user was shown
+    labeling_request: LabelingRequest
+
+    # What the user responded
+    qa_pairs: list[QAPair] = Field(default_factory=list[QAPair])
+    label: LabeledRun | None = None
+
+    @model_validator(mode="after")
+    def validate_nested_agent_run_ids(self) -> "AgentRunFeedbackContext":
+        """Ensure nested run IDs are consistent with the top-level run ID."""
+        if self.label is not None and self.label.agent_run_id != self.agent_run_id:
+            raise ValueError("label.agent_run_id must match agent_run_id")
+        return self
+
+    def to_str(self, indent: int = 0) -> str:
+        """Render full feedback entry in a deterministic LLM-facing format."""
+        lines = self.labeling_request.to_str(indent=indent).splitlines()
+
+        qa_lines: list[str] = []
+        if not self.qa_pairs:
+            qa_lines.append("N/A")
+        else:
+            for qa_idx, qa_pair in enumerate(self.qa_pairs, start=1):
+                qa_entry_lines = qa_pair.to_str(
+                    labeling_request=self.labeling_request,
+                    indent=indent,
+                ).splitlines()
+                qa_lines.extend(_tag_block(f"QA {qa_idx}", qa_entry_lines, indent))
+        lines.extend(_tag_block("Question Answer Pairs", qa_lines, indent))
+
+        if self.label is None:
+            label_body_lines = [
+                "User label: N/A",
+                "User explanation: N/A",
+            ]
+            label_body_lines.extend(
+                _render_user_distribution_block(self.labeling_request.user_distribution, indent)
+            )
+            label_body_lines.extend(
+                _render_user_distribution_reasoning_block(
+                    self.labeling_request.user_distribution_reasoning,
+                    reasoning_citations=None,
+                    indent=indent,
+                )
+            )
+            lines.extend(_tag_block("Label", label_body_lines, indent))
+        else:
+            lines.extend(
+                self.label.to_str(
+                    labeling_request=self.labeling_request,
+                    indent=indent,
+                ).splitlines()
+            )
+        return "\n".join(lines)
+
+
+class FeedbackContext(BaseModel):
+    """Feedback context returned by the feedback REST API."""
+
+    feedback_context_id: str
+    feedback_session_id: str
+    agent_run_id: str
+    labeling_request: LabelingRequest
+    created_at: datetime
+    updated_at: datetime
+
+
+class FeedbackContextsResponse(BaseModel):
+    """Round-scoped feedback contexts returned by the feedback REST API."""
+
+    current_round: int
+    contexts: list[FeedbackContext] = Field(default_factory=list[FeedbackContext])
+
+
+FeedbackJobStatus = Literal["pending", "running", "cancelling", "canceled", "completed"]
+
+
+class StartFeedbackContextsJobResponse(BaseModel):
+    """Response for enqueueing or reusing a feedback contexts job."""
+
+    job_id: str
+
+
+class FeedbackContextsJobStateResponse(BaseModel):
+    """Current feedback contexts job status and round-scoped contexts."""
+
+    job_id: str | None
+    job_status: FeedbackJobStatus | None
+    current_round: int
+    contexts: list[FeedbackContext] = Field(default_factory=list[FeedbackContext])
+
+
+class UserData(BaseModel):
+    """User Data (U) for user-context inference and downstream evaluation."""
+
+    initial_rubric: str
+    agent_run_feedbacks: list[AgentRunFeedbackContext] = Field(
+        default_factory=lambda: list[AgentRunFeedbackContext]()
+    )
+    created_at: datetime = Field(default_factory=datetime.now)
+    last_updated: datetime = Field(default_factory=datetime.now)
+
+    def upsert_run_feedback(self, agent_run_feedback: AgentRunFeedbackContext) -> None:
+        """Insert or replace feedback for an agent run ID, updating timestamps."""
+        now = datetime.now()
+        upserted_feedback = agent_run_feedback.model_copy(deep=True)
+        upserted_feedback.last_updated = now
+
+        for idx, existing in enumerate(self.agent_run_feedbacks):
+            if existing.agent_run_id != upserted_feedback.agent_run_id:
+                continue
+            upserted_feedback.created_at = existing.created_at
+            self.agent_run_feedbacks[idx] = upserted_feedback
+            self.last_updated = now
+            return
+
+        self.agent_run_feedbacks.append(upserted_feedback)
+        self.last_updated = now
+
+    def validate_against_agreement_keys(self, agreement_keys: set[str]) -> None:
+        """Validate stored labels and p_u outcomes against rubric agreement keys."""
+        for feedback in self.agent_run_feedbacks:
+            run_id = feedback.agent_run_id
+
+            label = feedback.label
+            if label is not None:
+                invalid_label_keys = sorted(set(label.label_value.keys()) - agreement_keys)
+                if invalid_label_keys:
+                    raise ValueError(
+                        "Run "
+                        f"{run_id} has label_value keys outside rubric agreement keys: "
+                        + ", ".join(invalid_label_keys)
+                    )
+
+            user_distribution = feedback.labeling_request.user_distribution
+            if user_distribution is None:
+                continue
+
+            for outcome_idx, outcome in enumerate(user_distribution.outcomes, start=1):
+                invalid_output_keys = sorted(set(outcome.output.keys()) - agreement_keys)
+                if invalid_output_keys:
+                    raise ValueError(
+                        "Run "
+                        f"{run_id} has user_distribution outcome #{outcome_idx} keys outside "
+                        "rubric agreement keys: " + ", ".join(invalid_output_keys)
+                    )
+                for key, value in outcome.output.items():
+                    if isinstance(value, (str, bool, int, float)):
+                        continue
+                    raise ValueError(
+                        "Run "
+                        f"{run_id} has user_distribution outcome #{outcome_idx} non-scalar "
+                        f"value for key '{key}': {type(value).__name__}"
+                    )
+
+    def iter_answered_qa_entries(self) -> Iterator[tuple[AgentRunFeedbackContext, QAPair]]:
+        """Iterate answered QA pairs with their parent run feedback."""
+        for feedback in self.agent_run_feedbacks:
+            for qa_pair in feedback.qa_pairs:
+                if qa_pair.status == "answered":
+                    yield feedback, qa_pair
+
+    def iter_skipped_qa_entries(self) -> Iterator[tuple[AgentRunFeedbackContext, QAPair]]:
+        """Iterate skipped QA pairs with their parent run feedback."""
+        for feedback in self.agent_run_feedbacks:
+            for qa_pair in feedback.qa_pairs:
+                if qa_pair.status == "skipped":
+                    yield feedback, qa_pair
+
+    def iter_labeled_entries(self) -> Iterator[tuple[AgentRunFeedbackContext, LabeledRun]]:
+        """Iterate labeled run entries with their parent run feedback."""
+        for feedback in self.agent_run_feedbacks:
+            if feedback.label is None:
+                continue
+            yield feedback, feedback.label
+
+
+# Backward-compatible alias used by older callers/scripts.
+AgentRunFeedback = AgentRunFeedbackContext

docent_python-0.1.52a0/docent/judges/util/voting.py

@@ -0,0 +1,351 @@
+import json
+import math
+from collections import Counter
+from itertools import product
+from typing import Any, TypedDict, cast
+
+import numpy as np
+from pydantic import BaseModel, Field
+
+AgreementValue = str | bool | int | float
+
+
+class EstimateWithCI(TypedDict):
+    mean: float
+    var: float
+    n: int
+    ci_95: float
+
+
+JudgeOutputDistribution = dict[AgreementValue, EstimateWithCI]
+
+
+class DistributionOutcome(BaseModel):
+    """Single outcome and probability mass for a predictive distribution."""
+
+    output: dict[str, Any]
+    probability: float
+
+
+class OutputDistribution(BaseModel):
+    """Probability distribution over rubric-compliant outputs."""
+
+    outcomes: list[DistributionOutcome] = Field(default_factory=list[DistributionOutcome])
+
+
+def _stable_json_dict(value: dict[str, Any]) -> str:
+    return json.dumps(value, sort_keys=True, separators=(",", ":"))
+
+
+def normalize_output_distribution(distribution: OutputDistribution) -> OutputDistribution:
+    """Normalize probabilities and merge duplicate outcomes by canonical JSON output."""
+    if not distribution.outcomes:
+        return OutputDistribution()
+
+    merged: dict[str, DistributionOutcome] = {}
+    for outcome in distribution.outcomes:
+        key = _stable_json_dict(outcome.output)
+        existing = merged.get(key)
+        if existing is None:
+            merged[key] = DistributionOutcome(
+                output=outcome.output,
+                probability=max(0.0, outcome.probability),
+            )
+            continue
+
+        existing.probability += max(0.0, outcome.probability)
+
+    merged_outcomes = list(merged.values())
+    if not merged_outcomes:
+        return OutputDistribution()
+
+    total_probability = sum(item.probability for item in merged_outcomes)
+    if total_probability <= 0:
+        uniform_prob = 1.0 / len(merged_outcomes)
+        for item in merged_outcomes:
+            item.probability = uniform_prob
+    else:
+        for item in merged_outcomes:
+            item.probability = item.probability / total_probability
+
+    merged_outcomes.sort(key=lambda item: item.probability, reverse=True)
+
+    return OutputDistribution(outcomes=merged_outcomes)
+
+
+def assert_agreement_only_output_schema(schema: dict[str, Any]) -> list[str]:
+    """Validate agreement-only schema contract for elicitation and entropy workflows.
+
+    Contract:
+    - Every top-level property must be an agreement key (enum or boolean),
+    - ``additionalProperties`` must be explicitly ``False``, and
+    - At least one agreement key must exist.
+    """
+    # if schema.get("additionalProperties") is not False:
+    #     raise ValueError(
+    #         "Rubric output_schema must set additionalProperties to false for "
+    #         "agreement-only entropy workflows."
+    #     )
+
+    properties_obj = schema.get("properties")
+    if not isinstance(properties_obj, dict):
+        raise ValueError("Rubric output_schema must define an object-valued properties field.")
+    properties = cast(dict[str, Any], properties_obj)
+
+    agreement_keys = get_agreement_keys(schema)
+    if not agreement_keys:
+        raise ValueError(
+            "Rubric output_schema must include at least one top-level agreement key "
+            "(enum or boolean)."
+        )
+
+    non_agreement_keys = sorted(set(properties.keys()) - set(agreement_keys))
+    if non_agreement_keys:
+        details: list[str] = []
+        for key in non_agreement_keys:
+            field_obj = properties.get(key, {})
+            if isinstance(field_obj, dict):
+                field_schema = cast(dict[str, Any], field_obj)
+                field_type = field_schema.get("type")
+                field_type_text = field_type if isinstance(field_type, str) else "unknown"
+                has_enum = "enum" in field_schema
+                details.append(f"{key} (type={field_type_text}, enum={has_enum})")
+            else:
+                details.append(f"{key} (type=invalid-schema)")
+        raise ValueError(
+            "Rubric output_schema includes non-agreement top-level keys: " + ", ".join(details)
+        )
+
+    return agreement_keys
+
+
+def get_agreement_keys(schema: dict[str, Any]) -> list[str]:
+    """Get top-level schema keys that support agreement computations.
+
+    This includes top-level enum and boolean fields.
+    """
+    agreement_keys: list[str] = []
+
+    properties = schema.get("properties", {})
+    assert isinstance(properties, dict)
+    properties = cast(dict[str, Any], properties)
+
+    for key, field_schema in properties.items():
+        assert isinstance(field_schema, dict)
+        field_schema = cast(dict[str, Any], field_schema)
+
+        field_type = field_schema.get("type")
+        assert isinstance(field_type, str)
+
+        if field_type == "boolean":
+            agreement_keys.append(key)
+        elif "enum" in field_schema:
+            agreement_keys.append(key)
+
+    return agreement_keys
+
+
+def get_agreement_key_options(
+    schema: dict[str, Any],
+    agreement_keys: list[str] | None = None,
+) -> dict[str, list[AgreementValue]]:
+    """Return possible output options for each agreement key from schema."""
+    if agreement_keys is None:
+        agreement_keys = get_agreement_keys(schema)
+
+    properties = schema.get("properties", {})
+    assert isinstance(properties, dict)
+    properties = cast(dict[str, Any], properties)
+
+    key_options: dict[str, list[AgreementValue]] = {}
+    for key in agreement_keys:
+        field_schema_obj = properties.get(key, {})
+        assert isinstance(field_schema_obj, dict)
+        field_schema = cast(dict[str, Any], field_schema_obj)
+
+        field_type = field_schema.get("type")
+        assert isinstance(field_type, str)
+
+        if field_type == "boolean":
+            key_options[key] = [True, False]
+            continue
+
+        if "enum" in field_schema:
+            enum_values = field_schema.get("enum")
+            assert isinstance(enum_values, list)
+            options: list[AgreementValue] = []
+            for enum_value in cast(list[object], enum_values):
+                assert isinstance(enum_value, (str, bool, int, float))
+                options.append(enum_value)
+            key_options[key] = options
+            continue
+
+        key_options[key] = []
+
+    return key_options
+
+
+def enumerate_agreement_output_space(
+    schema: dict[str, Any],
+    agreement_keys: list[str] | None = None,
+    max_outcomes: int | None = None,
+) -> list[dict[str, AgreementValue]]:
+    """Enumerate all possible agreement-only output objects from schema key options."""
+    if agreement_keys is None:
+        agreement_keys = get_agreement_keys(schema)
+    if not agreement_keys:
+        raise ValueError("Cannot enumerate output space: no agreement keys were found.")
+
+    if max_outcomes is not None and max_outcomes <= 0:
+        raise ValueError("max_outcomes must be > 0 when provided.")
+
+    key_options = get_agreement_key_options(schema, agreement_keys)
+    ordered_options: list[list[AgreementValue]] = []
+    total_outcomes = 1
+
+    for key in agreement_keys:
+        options = key_options.get(key, [])
+        if not options:
+            raise ValueError(
+                f"Cannot enumerate output space: agreement key '{key}' has no valid options."
+            )
+        ordered_options.append(options)
+        total_outcomes *= len(options)
+        if max_outcomes is not None and total_outcomes > max_outcomes:
+            raise ValueError(
+                "Cannot enumerate output space: "
+                f"{total_outcomes} outcomes exceeds max_outcomes={max_outcomes}."
+            )
+
+    return [
+        {key: value for key, value in zip(agreement_keys, combo)}
+        for combo in product(*ordered_options)
+    ]
+
+
+def find_modal_result(indep_results: list[dict[str, Any]], agreement_keys: list[str]):
+    """Find the result that best matches modal values across agreement keys.
+
+    Args:
+        indep_results: List of independent results to analyze
+        agreement_keys: Keys to measure agreement on
+
+    Returns:
+        Tuple of (max_idx, agt_key_modes_and_counts) where:
+        - max_idx is the index of the result that best matches modal values
+        - agt_key_modes_and_counts maps each key to (modal_value, count) or None if no values exist for that key
+
+    Raises:
+        ValueError: If no results are provided
+    """
+    if not indep_results:
+        raise ValueError("No results to score")
+
+    # For each agreement key, compute the mode and count (or None, if no values exist for that key)
+    agt_key_modes_and_counts: dict[str, tuple[str | bool | int, int] | None] = {}
+    for key in agreement_keys:
+        key_modes = Counter(v for r in indep_results if (v := r.get(key)) is not None)
+        if most_common_one := key_modes.most_common(1):
+            agt_key_modes_and_counts[key] = most_common_one[0]
+        else:
+            agt_key_modes_and_counts[key] = None
+
+    # Score each rollout based on how many agreement keys they match
+    # If there is no mode for a key, or if a certain result doesn't have that key, it doesn't count.
+    # TODO(mengk): This may bias towards results that have more keys.
+    indep_result_scores: list[int] = []
+    for r in indep_results:
+        score = 0
+        for key in agreement_keys:
+            mode_and_count = agt_key_modes_and_counts[key]
+            if mode_and_count and r.get(key) == mode_and_count[0]:
+                score += 1
+        indep_result_scores.append(score)
+
+    # Argmax
+    max_idx = indep_result_scores.index(max(indep_result_scores))
+
+    return max_idx, agt_key_modes_and_counts
+
+
+def compute_output_distributions(
+    eq_weighted_outputs: list[dict[str, Any]],
+    output_schema: dict[str, Any],
+    agreement_keys: list[str],
+) -> dict[str, JudgeOutputDistribution]:
+    """Estimate per-key output value distributions from equally weighted outputs.
+
+    For each key in ``agreement_keys``, this function:
+    1. Enumerates allowed values from the schema (``enum`` or ``boolean``),
+    2. Counts observed non-null values in ``eq_weighted_outputs``,
+    3. Normalizes counts into empirical probabilities, and
+    4. Computes per-value summary statistics:
+       - ``mean``: empirical probability ``p``
+       - ``var``: Bernoulli variance ``p * (1 - p)``
+       - ``n``: number of observed (non-null) values for that key
+       - ``ci_95``: normal-approximation 95% CI half-width ``1.96 * sqrt(var / n)``
+
+    Optional fields that are missing in some results are skipped (not counted toward ``n``).
+    If no values are observed for a key, all value probabilities are set to ``0.0`` and ``n=0``.
+
+    Args:
+        eq_weighted_outputs: Output objects to aggregate, each counted with equal weight.
+        output_schema: JSON schema used to validate/define allowed output values.
+        agreement_keys: Keys to include in the distribution computation.
+
+    Returns:
+        Mapping from agreement key to per-value distribution estimates.
+
+    Raises:
+        AssertionError: If an observed value is not one of the schema-derived possible values.
+    """
+
+    key_options = get_agreement_key_options(output_schema, agreement_keys)
+    raw_counts: dict[str, dict[AgreementValue, int]] = {
+        key: {value: 0 for value in key_options.get(key, [])} for key in agreement_keys
+    }
+    # Collect counts for each possible value
+    for result in eq_weighted_outputs:
+        for key in agreement_keys:
+            if (value := result.get(key)) is not None:  # Could be none if the key is optional
+                assert value in raw_counts[key], (
+                    "this should never happen; the value must be in possible values, since judge results have been validated against the schema"
+                )
+                raw_counts[key][value] += 1
+
+    distributions: dict[str, JudgeOutputDistribution] = {}
+    for agt_key in agreement_keys:
+        distributions[agt_key] = {}
+
+        # First normalize the counts to get probabilities
+        counts = raw_counts[agt_key]
+        total = sum(counts.values())
+        probs = {value: (count / total) if total > 0 else 0.0 for value, count in counts.items()}
+
+        for output_key, value in probs.items():
+            mean, estimate_var = value, (value * (1 - value))
+            # TODO(mengk): change to the wilson score interval
+            ci_95 = float(1.96 * np.sqrt(estimate_var / total)) if total > 0 else 0.0
+            estimate: EstimateWithCI = {
+                "mean": mean,
+                "var": estimate_var,
+                "n": total,
+                "ci_95": ci_95,
+            }
+            distributions[agt_key][output_key] = estimate
+
+    return distributions
+
+
+def compute_entropy(distribution: OutputDistribution) -> float:
+    """Compute Shannon entropy over normalized outcomes in nats."""
+    normalized = normalize_output_distribution(distribution)
+    if not normalized.outcomes:
+        return 0.0
+
+    entropy = 0.0
+    for outcome in normalized.outcomes:
+        probability = outcome.probability
+        if probability > 0:
+            entropy -= probability * math.log(probability)
+    return entropy

{docent_python-0.1.50a0 → docent_python-0.1.52a0}/docent/sdk/client.py

@@ -21,6 +21,11 @@ from tqdm import tqdm
 from docent._llm_util.providers.preference_types import ModelOption
 from docent._log_util.logger import LoggerAdapter, get_logger
 from docent.data_models.agent_run import AgentRun
+from docent.data_models.feedback import (
+    AgentRunFeedbackContext,
+    FeedbackContextsJobStateResponse,
+    StartFeedbackContextsJobResponse,
+)
 from docent.data_models.judge import Label
 from docent.judges.util.meta_schema import validate_judge_result_schema
 from docent.loaders import load_inspect
@@ -962,6 +967,113 @@ class Docent:
         clustering_state = self.get_clustering_state(collection_id, rubric_id)
         return clustering_state.get("assignments", {})
 
+    def create_feedback_session(
+        self,
+        collection_id: str,
+        rubric_id: str,
+        rubric_version: int,
+    ) -> str:
+        """Create a feedback session for a specific rubric version."""
+        url = f"{self._api_url}/feedback/{collection_id}/session"
+        payload = {
+            "rubric_id": rubric_id,
+            "rubric_version": rubric_version,
+        }
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+        return response.json()["feedback_session_id"]
+
+    def start_feedback_contexts_job(
+        self,
+        collection_id: str,
+        feedback_session_id: str,
+        num_samples: int = 50,
+        top_n: int = 20,
+        seed: int = 0,
+        candidate_pool_limit: int = 1_000,
+        where_clause: str | None = None,
+        increment_round: bool = False,
+    ) -> StartFeedbackContextsJobResponse:
+        """Start or reuse a background job to compute feedback contexts for a session."""
+        payload = {
+            "feedback_session_id": feedback_session_id,
+            "num_samples": num_samples,
+            "top_n": top_n,
+            "seed": seed,
+            "candidate_pool_limit": candidate_pool_limit,
+            "where_clause": where_clause,
+            "increment_round": increment_round,
+        }
+        url = f"{self._api_url}/feedback/{collection_id}/contexts/start"
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+        return StartFeedbackContextsJobResponse.model_validate(response.json())
+
+    def get_feedback_contexts(
+        self,
+        collection_id: str,
+        feedback_session_id: str,
+    ) -> FeedbackContextsJobStateResponse:
+        """Get feedback contexts state for a session, including job status and current round data."""
+        payload = {
+            "feedback_session_id": feedback_session_id,
+        }
+        url = f"{self._api_url}/feedback/{collection_id}/contexts/state"
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+        return FeedbackContextsJobStateResponse.model_validate(response.json())
+
+    def get_agent_run_feedback_contexts_by_session(
+        self,
+        collection_id: str,
+        feedback_session_id: str,
+    ) -> list[AgentRunFeedbackContext]:
+        """Get all persisted AgentRun feedback contexts for a feedback session.
+
+        Returns fully hydrated AgentRun feedback context objects from the database,
+        including QA pairs and labels (if present), across all rounds in the session.
+        """
+        url = f"{self._api_url}/feedback/{collection_id}/session/{feedback_session_id}/contexts"
+        response = self._session.get(url)
+        self._handle_response_errors(response)
+        return [AgentRunFeedbackContext.model_validate(item) for item in response.json()]
+
+    def add_feedback_qa(
+        self,
+        collection_id: str,
+        feedback_context_id: str,
+        focus_index: int,
+        answer: str,
+        explanation: str | None = None,
+    ) -> None:
+        """Submit a QA answer for a feedback context."""
+        payload = {
+            "feedback_context_id": feedback_context_id,
+            "focus_index": focus_index,
+            "answer": answer,
+            "explanation": explanation,
+        }
+        url = f"{self._api_url}/feedback/{collection_id}/qa"
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+
+    def upsert_feedback_label(
+        self,
+        collection_id: str,
+        feedback_context_id: str,
+        label_value: dict[str, Any],
+        explanation: str | None = None,
+    ) -> None:
+        """Create or update a label for a feedback context."""
+        payload = {
+            "feedback_context_id": feedback_context_id,
+            "label_value": label_value,
+            "explanation": explanation,
+        }
+        url = f"{self._api_url}/feedback/{collection_id}/label"
+        response = self._session.put(url, json=payload)
+        self._handle_response_errors(response)
+
     def create_label_set(
         self,
         collection_id: str,

{docent_python-0.1.50a0 → docent_python-0.1.52a0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "docent-python"
 description = "Docent SDK"
-version = "0.1.50-alpha"
+version = "0.1.52-alpha"
 authors = [
   { name="Transluce", email="info@transluce.org" },
 ]

docent_python-0.1.50a0/docent/judges/util/voting.py

@@ -1,140 +0,0 @@
-from collections import Counter
-from typing import Any, TypedDict, cast
-
-import numpy as np
-
-
-class EstimateWithCI(TypedDict):
-    mean: float
-    var: float
-    n: int
-    ci_95: float
-
-
-JudgeOutputDistribution = dict[str | bool | int | float, EstimateWithCI]
-
-
-def get_agreement_keys(schema: dict[str, Any]) -> list[str]:
-    """Get list of top-level keys in schema that we want to measure agreement on.
-
-    This includes enum and bool fields.
-
-    Args:
-        schema: JSON schema dict
-
-    Returns:
-        List of field names (keys) that should be used for measuring agreement
-    """
-    agreement_keys: list[str] = []
-
-    properties = schema.get("properties", {})
-    assert isinstance(properties, dict)
-    properties = cast(dict[str, Any], properties)
-
-    for key, field_schema in properties.items():
-        assert isinstance(field_schema, dict)
-        field_schema = cast(dict[str, Any], field_schema)
-
-        field_type = field_schema.get("type")
-        assert isinstance(field_type, str)
-
-        # Include boolean fields
-        if field_type == "boolean":
-            agreement_keys.append(key)
-        # Include enum fields (strings and numbers must be in this category)
-        elif "enum" in field_schema:
-            agreement_keys.append(key)
-
-    return agreement_keys
-
-
-def find_modal_result(indep_results: list[dict[str, Any]], agreement_keys: list[str]):
-    """Find the result that best matches modal values across agreement keys.
-
-    Args:
-        indep_results: List of independent results to analyze
-        agreement_keys: Keys to measure agreement on
-
-    Returns:
-        Tuple of (max_idx, agt_key_modes_and_counts) where:
-        - max_idx is the index of the result that best matches modal values
-        - agt_key_modes_and_counts maps each key to (modal_value, count) or None if no values exist for that key
-
-    Raises:
-        ValueError: If no results are provided
-    """
-    if not indep_results:
-        raise ValueError("No results to score")
-
-    # For each agreement key, compute the mode and count (or None, if no values exist for that key)
-    agt_key_modes_and_counts: dict[str, tuple[str | bool | int, int] | None] = {}
-    for key in agreement_keys:
-        key_modes = Counter(v for r in indep_results if (v := r.get(key)) is not None)
-        if most_common_one := key_modes.most_common(1):
-            agt_key_modes_and_counts[key] = most_common_one[0]
-        else:
-            agt_key_modes_and_counts[key] = None
-
-    # Score each rollout based on how many agreement keys they match
-    # If there is no mode for a key, or if a certain result doesn't have that key, it doesn't count.
-    # TODO(mengk): This may bias towards results that have more keys.
-    indep_result_scores: list[int] = []
-    for r in indep_results:
-        score = 0
-        for key in agreement_keys:
-            mode_and_count = agt_key_modes_and_counts[key]
-            if mode_and_count and r.get(key) == mode_and_count[0]:
-                score += 1
-        indep_result_scores.append(score)
-
-    # Argmax
-    max_idx = indep_result_scores.index(max(indep_result_scores))
-
-    return max_idx, agt_key_modes_and_counts
-
-
-def compute_output_distributions(
-    indep_results: list[dict[str, Any]], output_schema: dict[str, Any], agreement_keys: list[str]
-):
-    def _get_possible_values(key: str) -> list[str | bool | int | float]:
-        if "enum" in output_schema.get("properties", {}).get(key, {}):
-            return output_schema.get("properties", {}).get(key, {}).get("enum", [])
-        elif output_schema.get("properties", {}).get(key, {}).get("type") == "boolean":
-            return [True, False]
-        else:
-            return []
-
-    raw_counts: dict[str, dict[str | bool | int | float, int]] = {
-        key: {value: 0 for value in _get_possible_values(key)} for key in agreement_keys
-    }
-    # Collect counts for each possible value
-    for result in indep_results:
-        for key in agreement_keys:
-            if (value := result.get(key)) is not None:  # Could be none if the key is optional
-                assert value in raw_counts[key], (
-                    "this should never happen; the value must be in possible values, since judge results have been validated against the schema"
-                )
-                raw_counts[key][value] += 1
-
-    distributions: dict[str, JudgeOutputDistribution] = {}
-    for agt_key in agreement_keys:
-        distributions[agt_key] = {}
-
-        # First normalize the counts to get probabilities
-        counts = raw_counts[agt_key]
-        total = sum(counts.values())
-        probs = {value: (count / total) if total > 0 else 0.0 for value, count in counts.items()}
-
-        for output_key, value in probs.items():
-            mean, estimate_var = value, (value * (1 - value))
-            # TODO(mengk): change to the wilson score interval
-            ci_95 = float(1.96 * np.sqrt(estimate_var / total)) if total > 0 else 0.0
-            estimate: EstimateWithCI = {
-                "mean": mean,
-                "var": estimate_var,
-                "n": total,
-                "ci_95": ci_95,
-            }
-            distributions[agt_key][output_key] = estimate
-
-    return distributions