adaptive-harmony 0.1.23__py3-none-any.whl

adaptive_harmony/graders/templated_prompt_judge.py

@@ -0,0 +1,237 @@
+import html
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar
+
+from pybars import Compiler
+from pydantic import BaseModel, Field
+
+from adaptive_harmony import Grade, InferenceModel, StringThread, StringTurn
+from adaptive_harmony.core.structured_output import JsonParseError, render_schema
+from adaptive_harmony.core.utils import stringify_thread
+from adaptive_harmony.graders import BaseGrader
+from adaptive_harmony.graders.utils import (
+    FailedJudgeLog,
+    SuccessJudgeLog,
+    separate_context_from_last_user_turn,
+    validate_thread_last_assistant,
+)
+
+OPENAI_MODEL_FAMILIES_TEMPERATURE_1_ONLY = ["gpt-5", "o1", "o3", "o4"]
+
+
+def _turns_to_dicts(turns: list[StringTurn]) -> list[dict[str, str]]:
+    return [{"role": turn.role, "content": turn.content} for turn in turns]
+
+
+class BaseTemplatedPromptJudgeOutput(ABC, BaseModel):
+    @abstractmethod
+    def get_score(self) -> float:
+        pass
+
+    def get_reasoning(self) -> str | None:
+        return None
+
+
+class SimpleReasonedFloatOutput(BaseTemplatedPromptJudgeOutput):
+    """Sample default output format for simple float scores."""
+
+    score: float = Field(description="The numerical score for the sample")
+    reasoning: str = Field(description="Reasoning behind the score")
+
+    def get_score(self) -> float:
+        return self.score
+
+    def get_reasoning(self) -> str:
+        """Extract the reasoning from this output."""
+        return self.reasoning
+
+
+class BinaryJudgeOutput(BaseTemplatedPromptJudgeOutput):
+    """Output format for binary PASS/FAIL/NA judges."""
+
+    reasoning: str = Field(description="Reasoning string to support the rationale behind the score")
+    score: str = Field(description="The literal score for the sample", pattern="^(PASS|FAIL|NA)$")
+
+    SCORES_MAP: ClassVar[dict[str, float]] = {"PASS": 1.0, "FAIL": 0.0}
+
+    def get_score(self) -> float:
+        """Convert PASS/FAIL/NA to float. NA raises an exception."""
+        if self.score not in self.SCORES_MAP:
+            from adaptive_harmony.graders.exceptions import IgnoreScoreException
+
+            raise IgnoreScoreException(f"Non applicable score: {self.reasoning}")
+        return self.SCORES_MAP[self.score]
+
+    def get_reasoning(self) -> str:
+        return self.reasoning
+
+
+class TemplatedPromptJudgeGrader[T: BaseTemplatedPromptJudgeOutput](BaseGrader[SuccessJudgeLog | FailedJudgeLog]):
+    r"""
+    A flexible grader that uses Handlebars templates for system and user prompts.
+    Generic over T, which must inherit from BaseTemplatedPromptJudgeOutput.
+    This enforces that all output models implement get_score() and get_reasoning() methods.
+    Templates have access to comprehensive context extracted from the StringThread:
+    - output_schema: Expected output structure schema as a string
+    - turns: List of all turns as dicts with "role" and "content" keys
+    - metadata: Thread metadata dict
+    - context_turns: All turns without the assistant's completion (includes system prompt)
+    - context_str: Context formatted as string ('turn.role:\nturn.content\n' for every turn)
+    - context_turns_without_last_user: Same as context_turns, but without the last user turn
+    - context_str_without_last_user: Same as context_str, but without the last user turn
+    - last_user_turn_content: Content of the last user turn
+    - completion: Assistant's completion
+    Example Handlebars templates:
+    System: "You are a judge. Evaluate responses based on: {{criteria}}.
+             Always output the following JSON schema, with no preamble or postamble: {{output_schema}}"
+    User: "Conversation context:\n {{context_str_without_last_user}}
+           Question:\n {{last_user_turn_content}}
+           Response:\n {{completion}}"
+    Advanced examples:
+    - Conditionals: "{{#if metadata.domain}}Domain: {{metadata.domain}}{{/if}}"
+    - Loops with element index: "{{#each user_turns}}Turn {{@index}}: {{content}}{{/each}}"
+    - Built-in vars: "{{@index}}, {{@key}}, {{@first}}, {{@last}}"
+    """
+
+    def __init__(
+        self,
+        grader_key: str,
+        model: InferenceModel,
+        system_template: str,
+        user_template: str,
+        output_model: type[T],
+        template_variables: dict[str, Any] | None = None,
+        temperature: float = 0.0,
+        grader_id: str | None = None,
+    ):
+        """
+        Initialize the templated prompt judge.
+        Args:
+            grader_key: Unique identifier for this grader
+            model: Live InferenceModel to use as judge for grading
+            system_template: Handlebars template string for system prompt
+            user_template: Handlebars template string for user prompt
+            output_model: Pydantic model class that inherits from BaseOutput
+                         (must implement get_score() and get_reasoning() methods)
+            template_variables: Additional variables to make available in templates
+            grader_id: Optional grader ID (defaults to grader_key)
+        """
+        super().__init__(grader_key)
+        self._logs: list[SuccessJudgeLog | FailedJudgeLog] = []  # type: ignore[assignment]
+
+        # Set temperature to 1.0 if model_key is an OpenAI model in the temperature-1-only list
+        model_path: str = model.get_builder_args().get("path")  # type: ignore[assignment]
+        if model_path.startswith("openai://"):
+            model_name = model_path.removeprefix("openai://").split("?")[0]
+            if any(model_name.startswith(model) for model in OPENAI_MODEL_FAMILIES_TEMPERATURE_1_ONLY):
+                temperature = 1.0
+        self.model = model.temperature(temperature)
+        self.grader_id_or_key = grader_id or grader_key
+
+        # Template setup
+        self.compiler = Compiler()
+        self.system_template = self.compiler.compile(system_template)
+        self.user_template = self.compiler.compile(user_template)
+        self.template_variables = template_variables or {}
+
+        # Output configuration
+        self.output_model = output_model
+
+    @classmethod
+    def render_template(cls, template: str, thread: StringThread, output_model: type[BaseModel]):
+        compiler = Compiler()
+        compiled_template = compiler.compile(template)
+        context = cls.extract_template_context(thread, output_model)
+        return html.unescape(compiled_template(context))
+
+    @classmethod
+    def extract_template_context(
+        cls,
+        thread: StringThread,
+        output_model: type[BaseModel],
+        template_variables: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Extract context from StringThread for template rendering"""
+        validate_thread_last_assistant(thread)
+
+        turns = thread.get_turns()
+        context_without_last_user, last_user_turn = separate_context_from_last_user_turn(
+            thread, include_system_prompt=True
+        )
+        context = [turn for turn in turns[:-1]]
+
+        # Build comprehensive context
+        context = {
+            # Core thread data
+            "metadata": thread.metadata,
+            "turns": _turns_to_dicts(turns),
+            # Context and key turns
+            "context_turns": _turns_to_dicts(context),
+            "context_str": stringify_thread(StringThread(context)),
+            "context_turns_without_last_user": _turns_to_dicts(context_without_last_user),
+            "context_str_without_last_user": stringify_thread(StringThread(context_without_last_user)),
+            "last_user_turn_content": last_user_turn,
+            "completion": thread.last_content(),
+            "output_schema": render_schema(output_model),
+            # Additional template variables
+            **(template_variables or {}),
+        }
+
+        return context
+
+    def _build_judge_prompt(self, thread: StringThread) -> StringThread:
+        """Build the judging prompt using Handlebars templates"""
+        context = self.extract_template_context(thread, self.output_model, self.template_variables)
+
+        # Render templates using pybars with helpers
+        system_content = html.unescape(self.system_template(context))
+        user_content = html.unescape(self.user_template(context))
+
+        # Build prompt thread
+        judge_thread = StringThread().system(system_content).user(user_content)
+        return judge_thread
+
+    async def grade(self, sample: StringThread) -> Grade:
+        """Grade a sample using the templated prompt"""
+
+        judging_prompt = self._build_judge_prompt(sample)
+        str_prompt = stringify_thread(judging_prompt, sep=f"\n\n{'-' * 10}\n\n")
+
+        try:
+            _, parsed_output = await self.model.generate_and_validate(judging_prompt, self.output_model)
+        except JsonParseError as e:
+            self.add_log({"prompt": str_prompt, "error": f"{str(e)}\n\nCOMPLETION:\n{e.completion}"})
+            raise
+        except Exception as e:
+            self.add_log({"prompt": str_prompt, "error": str(e)})
+            raise
+
+        # Extract score and reasoning using the abstract methods
+        score = parsed_output.get_score()
+        reasoning = parsed_output.get_reasoning() or ""
+
+        grade = Grade(value=score, grader_key=self.grader_id_or_key, reasoning=reasoning)
+        self.add_log({"score": score, "prompt": str_prompt, "reasoning": reasoning})
+
+        return grade
+
+    def get_logs(self, clear: bool = False, log_all_samples: bool = False) -> dict[str, float | Any]:
+        """Get aggregated logs from all grading calls"""
+        # Get base statistics
+        logs = super().get_logs(clear=False)
+
+        # Get sample logs
+        successfully_scored_samples = [log for log in self._logs if "score" in log]
+        failed_scored_samples = [log for log in self._logs if "error" in log]
+
+        if not log_all_samples and successfully_scored_samples:
+            # Limit samples for display
+            successfully_scored_samples = successfully_scored_samples[:10]
+
+        sample_logs = self.get_sample_tables(successfully_scored_samples, failed_scored_samples)
+        logs.update(sample_logs)
+
+        if clear:
+            self.clear_logs()
+
+        return logs

adaptive_harmony/graders/utils.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypedDict
+
+from adaptive_harmony import StringThread
+
+if TYPE_CHECKING:
+    from adaptive_harmony import StringTurn
+
+
+def validate_thread_last_assistant(thread: StringThread):
+    turns = thread.get_turns()
+    assert len(turns) > 0, "The thread must have at least one turn"
+    assert turns[-1].role == "assistant", "The last turn must be an assistant turn"
+
+
+def separate_context_from_last_user_turn(
+    thread: StringThread,
+    include_system_prompt: bool = False,
+) -> tuple[list[StringTurn], str | None]:
+    """
+    Separates turns into context and last user turn.
+    Includes system prompt in context if include_system_prompt is True.
+    If there is no user turn, last user turn is None.
+    """
+    validate_thread_last_assistant(thread)
+    turns = thread.get_turns()
+
+    # Possibly remove system prompt
+    if not include_system_prompt and turns[0].role == "system":
+        turns = turns[1:]
+
+    # Find last user turn
+    user_question = None
+    user_question_idx = -1
+    for i, turn in enumerate(turns):
+        if turn.role == "user":
+            user_question = turn.content
+            user_question_idx = i
+
+    if user_question is None:
+        # Last turn is guaranteed to be assistant due to validate_thread_last_assistant
+        context_turns = turns[:-1]
+    else:
+        context_turns = turns[:user_question_idx]
+
+    return context_turns, user_question
+
+
+class SuccessJudgeLog(TypedDict):
+    prompt: str
+    reasoning: str
+    score: float
+
+
+class FailedJudgeLog(TypedDict):
+    prompt: str
+    error: str | None
+
+
+def sample_score_distribution(success_samples: list[SuccessJudgeLog], max_n_samples: int = 15) -> list[SuccessJudgeLog]:
+    # sort samples by score for percentile-based sampling
+    sorted_samples = sorted(success_samples, key=lambda x: x["score"])
+    total_samples = len(sorted_samples)
+
+    if total_samples >= max_n_samples:
+        # sample max_n_samples samples distributed across percentiles
+        indices = []
+        for i in range(max_n_samples):
+            # calculate percentile position (0% to 100% spread across 15 samples)
+            percentile = i / (max_n_samples - 1)  # 14 intervals for 15 samples
+            index = int(percentile * (total_samples - 1))
+            indices.append(index)
+
+        subset_successfully_scored_samples = [sorted_samples[i] for i in indices]
+    else:
+        subset_successfully_scored_samples = sorted_samples
+
+    return subset_successfully_scored_samples

adaptive_harmony/logging_table.py

@@ -0,0 +1 @@
+from harmony_client.logging_table import Table as Table