docent-python 0.1.14a0__py3-none-any.whl → 0.1.28a0__py3-none-any.whl

docent/judges/runner.py

@@ -0,0 +1,129 @@
+from typing import Protocol, Sequence, runtime_checkable
+
+import anyio
+from tqdm.auto import tqdm
+
+from docent._llm_util.llm_svc import BaseLLMService
+from docent._log_util import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.judges import (
+    JudgeResult,
+    JudgeResultCompletionCallback,
+    Rubric,
+)
+from docent.judges.impl import build_judge
+
+logger = get_logger(__name__)
+
+
+@runtime_checkable
+class AgentRunResolver(Protocol):
+    async def __call__(self) -> AgentRun | None: ...
+
+
+AgentRunInput = AgentRun | AgentRunResolver
+
+
+async def _resolve_agent_run(agent_run_input: AgentRunInput) -> AgentRun | None:
+    if isinstance(agent_run_input, AgentRun):
+        return agent_run_input
+    else:
+        return await agent_run_input()
+
+
+async def run_rubric(
+    agent_runs: Sequence[AgentRunInput],
+    rubric: Rubric,
+    llm_svc: BaseLLMService,
+    callback: JudgeResultCompletionCallback | None = None,
+    *,
+    n_rollouts_per_input: int | list[int] = 1,
+    show_progress: bool = True,
+) -> list[JudgeResult | None]:
+    if not agent_runs:
+        raise ValueError("agent_runs must be a non-empty sequence")
+    if rubric.n_rollouts_per_input <= 0:
+        raise ValueError("rubric.n_rollouts_per_input must be greater than 0")
+
+    # Normalize n_rollouts_per_input to a list
+    if isinstance(n_rollouts_per_input, int):
+        if n_rollouts_per_input < 0:
+            raise ValueError("n_rollouts_per_input must be non-negative")
+        rollouts_per_run = [n_rollouts_per_input] * len(agent_runs)
+    else:
+        rollouts_per_run = n_rollouts_per_input
+        if len(rollouts_per_run) != len(agent_runs):
+            raise ValueError("n_rollouts_per_input list must match agent_runs length")
+        if any(n < 0 for n in rollouts_per_run):
+            raise ValueError("All values in n_rollouts_per_input must be non-negative")
+
+    judge = build_judge(rubric, llm_svc)
+
+    total_rollouts = sum(rollouts_per_run)
+    logger.info(
+        "Running rubric %s version %s against %d agent runs with %d total rollouts",
+        rubric.id,
+        rubric.version,
+        len(agent_runs),
+        total_rollouts,
+    )
+
+    agent_results: list[list[JudgeResult | None]] = [[] for _ in agent_runs]
+    progress_bar = tqdm(
+        total=total_rollouts,
+        desc=f"Rubric {rubric.id}",
+        disable=not show_progress,
+    )
+
+    # NOTE(mengk): using a (2 * llm max concurrency) semaphore is a hack to avoid
+    #   hammering _resolve_agent_run, which makes expensive DB calls, when they aren't going to be
+    #   immediately processed by the LLMService anyways.
+    # TODO(mengk): We should eventually implement a more idiomatic solution to this.
+    #   It's related to the idea of a global concurrency limiter.
+    run_judge_semaphore = anyio.Semaphore(llm_svc.max_concurrency * 2)
+
+    async def _run_single_judge(index: int, agent_run_input: AgentRunInput):
+        async with run_judge_semaphore:
+            rollout_results: list[JudgeResult | None] = []
+
+            if rollouts_per_run[index] == 0:
+                agent_results[index] = []
+                if callback is not None:
+                    await callback(index, None)
+                return
+
+            agent_run = await _resolve_agent_run(agent_run_input)
+            if agent_run is None:
+                if callback is not None:
+                    await callback(index, None)
+                return
+
+            for _ in range(rollouts_per_run[index]):
+                result = await judge(agent_run)
+                rollout_results.append(result)
+                progress_bar.update()
+
+            agent_results[index] = rollout_results
+
+            if callback is not None:
+                # Filter out None results for the callback
+                valid_results = [r for r in rollout_results if r is not None]
+                await callback(index, valid_results if valid_results else None)
+
+    try:
+        async with anyio.create_task_group() as tg:
+            for index, agent_run in enumerate(agent_runs):
+                tg.start_soon(_run_single_judge, index, agent_run)
+    finally:
+        progress_bar.close()
+
+    flattened_results = [result for rollouts in agent_results for result in rollouts]
+    successful = sum(result is not None for result in flattened_results)
+    logger.info(
+        "Finished rubric %s: produced %d/%d judge results",
+        rubric.id,
+        successful,
+        len(flattened_results),
+    )
+
+    return flattened_results

docent/judges/stats.py

@@ -0,0 +1,205 @@
+from typing import Iterator, List, Tuple
+
+from scipy import stats
+
+from docent._log_util import get_logger
+
+logger = get_logger(__name__)
+
+
+def print_stats_with_intervals(name: str, mean: float, std: float, confidence_levels: list[float]):
+    """Print statistics with confidence intervals at multiple confidence levels.
+
+    Args:
+        name: Name of the statistic
+        mean: Mean value
+        std: Standard deviation
+        confidence_levels: List of confidence levels (e.g., [0.90, 0.95, 0.99])
+    """
+    intervals_str = ", ".join(
+        [
+            f"{int(level*100)}% interval: [{mean - stats.norm.ppf((1+level)/2) * std:.4f}, {mean + stats.norm.ppf((1+level)/2) * std:.4f}]"  # type: ignore
+            for level in confidence_levels
+        ]
+    )
+    print(f"{name} mean: {mean:.4f}, std: {std:.4f}, {intervals_str}")
+
+
+def _bounded_compositions(total: int, parts: int, bound: int) -> Iterator[Tuple[int, ...]]:
+    """
+    Yield all tuples (x1,...,x_parts) of nonnegative ints summing to `total`
+    with each xk <= bound.
+    """
+
+    # Recursive backtracking with pruning by remaining capacity.
+    def rec(k: int, remaining: int, prefix: List[int]) -> Iterator[Tuple[int, ...]]:
+        if k == parts:
+            if remaining == 0:
+                yield tuple(prefix)
+            return
+        # Max we can put here is min(bound, remaining - min_needed_for_rest)
+        # The min needed for the rest is 0; also cannot exceed remaining.
+        max_here = min(bound, remaining)
+        # Optional pruning: ensure the rest can absorb what's left (always true since min=0)
+        for x in range(max_here + 1):
+            prefix.append(x)
+            yield from rec(k + 1, remaining - x, prefix)
+            prefix.pop()
+
+    yield from rec(0, total, [])
+
+
+def plurality_vectors(m: int, K: int, i: int) -> Iterator[Tuple[int, ...]]:
+    """
+    Generate all count vectors n = (n1,...,nm) of nonnegative integers with
+    sum(n) = K and STRICT plurality at index i:
+       n[i] > n[j] for all j != i.
+
+    Yields vectors in no particular order.
+    """
+    if not (0 <= i < m):
+        raise ValueError("i must be in [0, m).")
+    if m < 2 or K < 1:
+        return  # nothing to yield in degenerate cases
+
+    for ni in range(1, K + 1):  # at least 1 vote for the winner
+        rest_total = K - ni
+        cap = ni - 1  # strict plurality: others must be <= ni-1
+        # If cap < 0 but rest_total > 0, impossible
+        if cap < 0 and rest_total > 0:
+            continue
+        # Build the other m-1 counts under the cap
+        for others in _bounded_compositions(rest_total, m - 1, cap):
+            # Stitch back in ni at position i
+            vec = list(others[:i]) + [ni] + list(others[i:])
+            yield tuple(vec)
+
+
+def p_mode(n: int, p_v: list[float], idx: int) -> float:
+    """Probability that the modal sample of sampling Multinom(n, p_v) is the idxth one."""
+    count_vecs = plurality_vectors(len(p_v), n, idx)
+    return sum(stats.multinomial.pmf(vec, n, p_v) for vec in count_vecs)  # type: ignore
+
+
+# async def analyze_majority_judge(
+#     rubric: Rubric,
+#     agent_runs: list[AgentRun],
+#     matched_labels: dict[str, dict[str, Any]],  # agent_run_id -> gold label obj
+#     results_path: Path,
+#     samples_per_agent_run: int = 10,
+#     maj_k: int = 5,  # Does not affect data collection
+#     max_llm_concurrency: int = 100,
+# ):
+#     # if rubric.n_rollouts_per_input != 1:
+#     #     raise ValueError("You should use n_rollouts_per_input=1")
+
+#     if not results_path.exists():
+#         logger.info(f"Evaluating rubrics and saving results to {results_path}")
+
+#         max_conc_per_rubric = min(
+#             max_llm_concurrency, len(agent_runs) * rubric.n_rollouts_per_input
+#         )
+#         max_parallel_rubrics = max(1, max_llm_concurrency // max_conc_per_rubric)
+#         logger.info(
+#             f"Evaluating {samples_per_agent_run} samples per agent run, {max_conc_per_rubric} concurrent LLM calls per rubric, {max_parallel_rubrics} parallel rubrics"
+#         )
+
+#         await evaluate_multiple_rubrics(
+#             rubrics=[rubric] * samples_per_agent_run,
+#             agent_runs=agent_runs,
+#             llm_svc=SimpleLLMService(),
+#             output_path=results_path,
+#             max_concurrent_llm_calls_per_rubric=max_conc_per_rubric,
+#             max_parallel_rubrics=max_parallel_rubrics,
+#         )
+#     else:
+#         logger.info(f"Found existing results at {results_path}, loading them")
+
+#     rows = load_rubric_results_from_file(results_path)
+
+#     # Parse results into a flat dataframe
+#     parsed_results: list[dict[str, Any]] = []
+#     for row in rows:
+#         config_key = row.rubric.model_dump_json(
+#             exclude={
+#                 "rubric_text",
+#                 "id",
+#                 "version",
+#                 "system_prompt_template",
+#                 "citation_instructions",
+#                 "output_schema",
+#             }
+#         )
+#         for result in row.judge_results:
+#             if result is None:
+#                 print("found null")
+#                 continue
+#                 # raise ValueError("result is None")
+
+#             gold_label = matched_labels[result.agent_run_id]["label"]
+#             judge_label = result.output["label"]
+#             md = result.result_metadata or {}
+
+#             parsed_results.append(
+#                 {
+#                     "config_key": config_key,
+#                     "agent_run_id": result.agent_run_id,
+#                     "gold_label": gold_label,
+#                     "judge_label": judge_label,
+#                     "final_results": md.get("final_results"),
+#                     "is_correct": gold_label == judge_label,
+#                 }
+#             )
+
+#     df = pd.DataFrame(parsed_results)
+#     potential_labels = ["cheating", "future leakage", "not suspicious"]
+
+#     def _get_pred_dist(x: pd.Series):
+#         x = x.tolist()
+#         counts = [x.count(l) for l in potential_labels]
+#         assert np.sum(counts) == len(x)
+#         return {l: counts[i] / len(x) for i, l in enumerate(potential_labels)}
+
+#     n_ars = len(df.groupby("agent_run_id").count())
+#     p_correct = (
+#         df.groupby("agent_run_id")
+#         .agg(
+#             {
+#                 "gold_label": lambda x: x.iloc[0],
+#                 "judge_label": _get_pred_dist,
+#                 "is_correct": np.mean,
+#             }
+#         )
+#         .rename(columns={"judge_label": "pred_dist", "is_correct": "p_correct_naive"})
+#     )
+#     p_correct["p_correct_naive_var"] = p_correct["p_correct_naive"].apply(lambda x: x * (1 - x))
+
+#     p_correct["p_correct_majority"] = p_correct.apply(
+#         lambda row: p_mode(
+#             maj_k,
+#             [row["pred_dist"][l] for l in potential_labels],
+#             potential_labels.index(row["gold_label"]),
+#         ),
+#         axis=1,
+#     )
+#     p_correct["p_correct_majority_var"] = p_correct["p_correct_majority"].apply(
+#         lambda x: x * (1 - x)
+#     )
+#     p_correct.sort_values(by="p_correct_majority_var", ascending=False, inplace=True)
+
+#     overall_naive_mean = p_correct["p_correct_naive"].mean()
+#     overall_naive_std = np.sqrt(p_correct["p_correct_naive_var"].sum() / n_ars**2)
+#     overall_majority_mean = p_correct["p_correct_majority"].mean()
+#     overall_majority_std = np.sqrt(p_correct["p_correct_majority_var"].sum() / n_ars**2)
+
+#     confidence_levels = [0.5, 0.95]
+#     print_stats_with_intervals(
+#         "Overall naive", overall_naive_mean, overall_naive_std, confidence_levels
+#     )
+#     print_stats_with_intervals(
+#         f"Overall majority (k={maj_k})",
+#         overall_majority_mean,
+#         overall_majority_std,
+#         confidence_levels,
+#     )
+#     return p_correct

docent/judges/types.py

@@ -0,0 +1,311 @@
+import enum
+import json
+from string import Formatter
+from typing import Any, Callable, Literal, Protocol
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, field_serializer, field_validator
+
+from docent._llm_util.providers.preference_types import PUBLIC_PROVIDER_PREFERENCES, ModelOption
+from docent._log_util import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.data_models.citation import parse_citations
+from docent.data_models.transcript import TEXT_RANGE_CITE_INSTRUCTION
+from docent.judges.util.meta_schema import validate_judge_result_schema
+
+logger = get_logger(__name__)
+
+DEFAULT_JUDGE_SYSTEM_PROMPT_TEMPLATE = """
+Here is a rubric that we are using to judge transcripts of AI agent runs.
+
+Rubric:
+<rubric>
+{rubric}
+</rubric>
+
+Agent run:
+<agent_run>
+{agent_run}
+</agent_run>
+
+Your goal is to judge the agent run according to the criteria given in the rubric. Start by faithfully following the decision procedure in extremely careful detail, step by step.
+
+When you are finished, output your final adjudication, surrounded by <response>...</response> tags. The response must be a valid JSON string which can be parsed with python `json.loads` without any additional processing. Double quotes (`"`) in the middle of a string in the JSON object must be escaped with a backslash.
+
+The JSON object you produce must adhere to the following schema:
+{output_schema}
+
+{citation_instructions}
+""".strip()
+
+DEFAULT_MULTI_TURN_JUDGE_SYSTEM_PROMPT_TEMPLATE = """
+Here is a rubric that we are using to judge transcripts of AI agent runs.
+
+Rubric:
+<rubric>
+{rubric}
+</rubric>
+
+Agent run:
+<agent_run>
+{agent_run}
+</agent_run>
+
+Your goal is to judge the agent run according to the criteria given in the rubric. Start by faithfully following the decision procedure in extremely careful detail, step by step. You must execute **one step in the decision procedure per assistant message turn**. After each turn, output a complete and detailed recount of all actions you took, and everything you discovered. Then call the `step_finished` tool.
+
+When you are finished going through the decision procedure, output your final adjudication, surrounded by <response>...</response> tags. The response must be a valid JSON string which can be parsed with python `json.loads` without any additional processing. Double quotes (`"`) in the middle of a string in the JSON object must be escaped with a backslash.
+
+The JSON object you produce must adhere to the following schema:
+{output_schema}
+
+{citation_instructions}
+""".strip()
+
+DEFAULT_EXPOSED_REASONING_JUDGE_SYSTEM_PROMPT_TEMPLATE = """
+Here is a rubric that we are using to judge transcripts of AI agent runs.
+
+Rubric:
+<rubric>
+{rubric}
+</rubric>
+
+Agent run:
+<agent_run>
+{agent_run}
+</agent_run>
+
+Your goal is to judge the agent run according to the criteria given in the rubric. Start by faithfully following the decision procedure in extremely careful detail, step by step. You must *fully externalize* your reasoning work by outputting details in the assistant message, surrounded by <reasoning>...</reasoning> tags. The reasoning section can be as messy as you need. You should use *high* reasoning effort.
+
+When you are finished, output your final adjudication in the assistant message, surrounded by <response>...</response> tags. The response must be a valid JSON string which can be parsed with python `json.loads` without any additional processing. Double quotes (`"`) in the middle of a string in the JSON object must be escaped with a backslash.
+
+The JSON object you produce must adhere to the following schema:
+{output_schema}
+
+{citation_instructions}
+""".strip()
+
+DEFAULT_CITATION_INSTRUCTIONS = f"""
+For strings which require citations (according to the `citations: True` property), you must also follow these instructions:
+{TEXT_RANGE_CITE_INSTRUCTION}
+""".strip()
+
+DEFAULT_JUDGE_OUTPUT_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "label": {"type": "string", "enum": ["match", "no match"]},
+        "explanation": {"type": "string", "citations": True},
+    },
+    # Require these properties to be present
+    "required": ["label", "explanation"],
+    # Allow additional properties though, as their presence is not breaking
+}
+
+DEFAULT_JUDGE_MODEL = PUBLIC_PROVIDER_PREFERENCES.default_judge_models[0]
+
+
+class JudgeVariant(str, enum.Enum):
+    MAJORITY = "majority"
+    MULTI_REFLECT = "multi-reflect"
+
+
+class Rubric(BaseModel):
+    """TODO(mengk): this should really be called JudgeConfig,
+    but temporarily keeping this for consistency with docent_core."""
+
+    class Config:
+        frozen = True
+
+    # Primary key
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    version: int = 1
+
+    # What the judge actually does
+    rubric_text: str
+    n_rollouts_per_input: int = 1
+    judge_variant: JudgeVariant = JudgeVariant.MAJORITY
+    # TODO(mengk): add this to the database
+    # No need right now because multi-turn is still very experimental.
+    rollout_type: Literal["single_turn", "multi_turn"] = "single_turn"
+
+    # Default instructions for the judge
+    system_prompt_template: str = DEFAULT_JUDGE_SYSTEM_PROMPT_TEMPLATE
+    citation_instructions: str = DEFAULT_CITATION_INSTRUCTIONS
+    output_schema: dict[str, Any] = DEFAULT_JUDGE_OUTPUT_SCHEMA
+
+    # How to run the judge
+    judge_model: ModelOption = DEFAULT_JUDGE_MODEL
+
+    def materialize_system_prompt(self, agent_run: AgentRun) -> str:
+        """Construct the full prompt text for rubric evaluation.
+
+        This is the canonical implementation of prompt construction - use this function
+        anywhere you need to construct a rubric evaluation prompt (including cost estimation).
+        """
+
+        output_schema_text = json.dumps(self.output_schema, indent=2)
+
+        # We've already validated that the system prompt template has these keys
+        prompt = self.system_prompt_template.format(
+            rubric=self.rubric_text,
+            agent_run=agent_run.to_text_new(),
+            output_schema=output_schema_text,
+            # Only include citation instructions if the schema requests citations
+            citation_instructions=(
+                self.citation_instructions if _schema_requests_citations(self.output_schema) else ""
+            ),
+        ).strip()
+
+        return prompt
+
+    @field_validator("system_prompt_template")
+    @classmethod
+    def validate_system_prompt_template(cls, system_prompt_template: str):
+        # Extract all field names from the template
+        formatter = Formatter()
+        field_names = {
+            field_name
+            for _, field_name, _, _ in formatter.parse(system_prompt_template)
+            if field_name is not None
+        }
+
+        # Check for required fields
+        required_fields = {"agent_run", "output_schema", "rubric", "citation_instructions"}
+        missing_fields = required_fields - field_names
+
+        if missing_fields:
+            raise ValueError(
+                f"system_prompt_template must contain the following placeholders: {missing_fields}"
+            )
+
+        return system_prompt_template
+
+    @field_validator("output_schema")
+    @classmethod
+    def validate_output_schema(cls, output_schema: dict[str, Any]):
+        """
+        Raises:
+            jsonschema.ValidationError: If the schema is invalid
+            jsonschema.SchemaError: If the schema is not a valid 2020-12 schema
+        """
+        validate_judge_result_schema(output_schema)
+        return output_schema
+
+
+class MultiTurnRubric(Rubric):
+    system_prompt_template: str = DEFAULT_MULTI_TURN_JUDGE_SYSTEM_PROMPT_TEMPLATE
+    rollout_type: Literal["single_turn", "multi_turn"] = "multi_turn"
+
+
+class ExposedReasoningRubric(Rubric):
+    system_prompt_template: str = DEFAULT_EXPOSED_REASONING_JUDGE_SYSTEM_PROMPT_TEMPLATE
+
+
+class ResultType(enum.Enum):
+    """Enum for the type of result that a judge result can have."""
+
+    DIRECT_RESULT = "direct_result"
+    NEAR_MISS = "near_miss"
+
+
+class JudgeResult(BaseModel):
+    class Config:
+        frozen = True
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    agent_run_id: str
+    rubric_id: str
+    rubric_version: int
+
+    # Outputs
+    output: dict[str, Any]
+    result_metadata: dict[str, Any] | None = None
+    result_type: ResultType
+
+    # Deprecated
+    value: str | None = None
+
+    @field_serializer("result_type")
+    def serialize_result_type(self, result_type: ResultType) -> str:
+        return result_type.value
+
+
+class JudgeResultWithCitations(JudgeResult):
+    @classmethod
+    def from_judge_result(
+        cls, result: JudgeResult, schema: dict[str, Any]
+    ) -> "JudgeResultWithCitations":
+        """Judge result must be validated against the schema before calling this function!"""
+
+        def _parse_citation_string(output: str) -> dict[str, Any]:
+            text, citations = parse_citations(output)
+            return {"text": text, "citations": citations}
+
+        data = result.model_dump()
+        try:
+            data["output"] = traverse_schema_and_transform(
+                data["output"], schema, _parse_citation_string
+            )
+        except Exception as e:
+            logger.error(f"Failed to parse citations: {e}")
+            logger.error(f"Output: {data['output']}")
+            data["output"] = {"raw": data["output"]}
+        return cls(**data)
+
+
+class JudgeResultCompletionCallback(Protocol):
+    """Called when some batch of judge results is completed.
+    Supports batched calls for cases where many results are pre-computed.
+    This avoids invoking the callback separately for each datapoint.
+    """
+
+    async def __call__(
+        self,
+        batch_index: int,
+        judge_results: list[JudgeResult] | None,
+    ) -> None: ...
+
+
+def traverse_schema_and_transform(
+    output: Any,
+    schema: dict[str, Any],
+    citation_string_handler: Callable[[str], Any],
+) -> Any:
+    """Recursively traverse output based on schema, applying citation_string_handler to citation strings."""
+    if schema.get("type") == "string" and schema.get("citations"):  # type: ignore
+        return citation_string_handler(output)
+    elif schema.get("type") == "object":
+        properties: dict[str, Any] = schema.get("properties", {})
+        result: dict[str, Any] = {}
+        for key in properties:
+            if key in output:
+                result[key] = traverse_schema_and_transform(
+                    output[key], properties[key], citation_string_handler
+                )
+        return result
+    elif schema.get("type") == "array":
+        item_schema: dict[str, Any] = schema.get("items", {})
+        return [
+            traverse_schema_and_transform(item, item_schema, citation_string_handler)
+            for item in output
+        ]
+    else:
+        return output
+
+
+def _schema_requests_citations(schema: dict[str, Any]) -> bool:
+    """Check if any field in the schema requests citations by having 'citations': 'true'."""
+
+    def _check_field(field_schema: Any) -> bool:
+        if isinstance(field_schema, dict):
+            if field_schema.get("citations"):  # type: ignore
+                return True
+            for value in field_schema.values():  # type: ignore
+                if isinstance(value, dict) and _check_field(value):
+                    return True
+                elif isinstance(value, list):
+                    for item in value:  # type: ignore
+                        if isinstance(item, dict) and _check_field(item):
+                            return True
+        return False
+
+    return _check_field(schema)