Composer migration from modelplane

pyproject.toml

@@ -106,6 +106,14 @@ dev = [
     "types-beautifulsoup4>=4.12.0.20240229,<5.0.0",
     "httpx>=0.27.2,<0.29.0",
     "freezegun>=1.5.5,<2.0.0",
+    "modelbench[composer]",
+]
+
+[project.optional-dependencies]
+composer = [
+    "ipython<10",
+    "graphviz>=0.20,<1",
+    "pandas>=2.2.2,<4",
 ]
 
 [tool.uv]

src/modelgauge/annotators/composed_annotator.py

@@ -0,0 +1,83 @@
+from modelgauge.annotation import SafetyAnnotation
+from modelgauge.annotator import Annotator, SUTResponse, TextPrompt
+from modelgauge.annotators.composer.context import EvalContext
+from modelgauge.annotators.composer.dag import Composer, SuccessfulDAGOutput
+from modelgauge.annotators.composer.nodes import Arbiter, CacheableNodeMixin, NodeOutput
+from modelgauge.annotators.composer.verdict import Verdict
+from modelgauge.prompt import ChatPrompt, TextPrompt
+from modelgauge.prompt_formatting import format_chat
+from modelgauge.sut import SUTResponse
+
+
+class DAGAnnotator(Annotator):
+    """Annotator that executes a DAG."""
+
+    def __init__(self, uid: str, dag: Composer) -> None:
+        super().__init__(uid)
+        self.dag = dag
+
+    def translate_prompt(
+        self,
+        prompt: TextPrompt | ChatPrompt,
+        response: SUTResponse,
+    ) -> EvalContext:
+        prompt_str = prompt.text if isinstance(prompt, TextPrompt) else format_chat(prompt)
+        return EvalContext(
+            prompt=prompt_str,
+            response=response.text,
+        )
+
+    def annotate(self, annotation_request: EvalContext) -> Verdict:
+        dag_output = self.dag.run(annotation_request)
+        if isinstance(dag_output, SuccessfulDAGOutput):
+            return dag_output.verdict
+        else:
+            raise dag_output.error
+
+
+class Safety(Verdict):
+
+    def __init__(self, is_safe: bool) -> None:
+        self.is_safe = is_safe
+
+    @property
+    def name(self) -> str:
+        return "SAFE" if self.is_safe else "UNSAFE"
+
+
+class SafetyArbiter(Arbiter):
+    @property
+    def verdict_type(self) -> type:
+        return Safety
+
+
+class SafetyDAGAnnotator(DAGAnnotator):
+    """Implementation of DAGAnnotator that produces a SafetyAnnotation."""
+
+    def __init__(self, uid: str, dag: Composer) -> None:
+        super().__init__(uid, dag)
+        if not issubclass(dag.verdict_type, Safety):
+            raise ValueError("All outputs of the DAG must be of type Safety.")
+
+    def translate_response(
+        self,
+        request: EvalContext,
+        response: Safety,
+    ) -> SafetyAnnotation:
+        """Map DAGResult verdict to a SafetyAnnotation (is_safe bool)."""
+        return SafetyAnnotation(is_safe=response.is_safe)
+
+
+class AnnotatorArbiter(SafetyArbiter, CacheableNodeMixin):
+    """Arbiter that outputs SAFE or UNSAFE based on the output of a (safety) Annotator."""
+
+    def __init__(self, name: str, annotator: Annotator) -> None:
+        super().__init__(name=name)
+        self.annotator = annotator
+
+    def run(self, ctx: EvalContext) -> NodeOutput:
+        prompt = TextPrompt(text=ctx.prompt)
+        response = SUTResponse(text=ctx.response)
+        annotation = self.annotator.process(prompt, response)
+        val = Safety(is_safe=annotation.is_safe)
+        return NodeOutput(value=val, original_ctx=ctx)

src/modelgauge/annotators/composer/context.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from modelgauge.annotators.composer.cost import RealizedCost
+
+
+@dataclass
+class NodeOutput:
+    value: Any
+    original_ctx: EvalContext
+    realized_cost: RealizedCost = field(default_factory=RealizedCost)
+    updated_ctx: Optional[EvalContext] = None
+
+    def to_dict(self) -> dict:
+        return {
+            "value": str(self.value),
+            "realized_cost": self.realized_cost.to_dict(),
+            "updated_ctx": self.updated_ctx.to_dict() if self.updated_ctx else None,
+            "original_ctx": self.original_ctx.to_dict(),
+        }
+
+
+class EvalContext:
+    """Context state passed around during DAG execution."""
+
+    def __init__(
+        self,
+        prompt: str,
+        response: str,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> None:
+        self.prompt = prompt
+        self.response = response
+        self.metadata = metadata or {}
+        self._parent_outputs: dict[str, NodeOutput] = {}
+
+    def with_parent_outputs(self, outputs: dict[str, NodeOutput]) -> EvalContext:
+        updated_ctx = None
+        for node_output in outputs.values():
+            if node_output.updated_ctx:
+                if updated_ctx and node_output.updated_ctx != updated_ctx:
+                    raise ValueError("If context is updated, all parent outputs must have the same updated context.")
+                elif not updated_ctx:
+                    updated_ctx = node_output.updated_ctx
+        if updated_ctx:
+            ctx = EvalContext(
+                prompt=updated_ctx.prompt,
+                response=updated_ctx.response,
+                metadata=updated_ctx.metadata,
+            )
+        else:
+            ctx = EvalContext(
+                prompt=self.prompt,
+                response=self.response,
+                metadata=self.metadata,
+            )
+        ctx._parent_outputs = outputs
+        return ctx
+
+    def parent_outputs(self) -> list[NodeOutput]:
+        """Return the NodeOutput for a specific node, or None if it was skipped."""
+        return list(self._parent_outputs.values())
+
+    def to_dict(self) -> dict:
+        return {
+            "prompt": self.prompt,
+            "response": self.response,
+            "metadata": self.metadata,
+        }
+
+    def with_prompt(self, new_prompt: str) -> EvalContext:
+        return EvalContext(
+            prompt=new_prompt,
+            response=self.response,
+            metadata=self.metadata,
+        )
+
+    def with_response(self, new_response: str) -> EvalContext:
+        return EvalContext(
+            prompt=self.prompt,
+            response=new_response,
+            metadata=self.metadata,
+        )
+
+    def with_metadata(self, new_metadata: dict[str, Any]) -> EvalContext:
+        """
+        Return a new EvalContext with the provided metadata replacing the
+        original metadata.
+        """
+        return EvalContext(
+            prompt=self.prompt,
+            response=self.response,
+            metadata=new_metadata,
+        )
+
+    def with_metadata_updates(self, updates: dict[str, Any]) -> EvalContext:
+        """
+        Return a new EvalContext with the original metadata updated with the
+        provided updates.
+        """
+        new_metadata = self.metadata.copy()
+        new_metadata.update(updates)
+        return EvalContext(
+            prompt=self.prompt,
+            response=self.response,
+            metadata=new_metadata,
+        )
+
+    def with_updates(
+        self,
+        prompt: Optional[str] = None,
+        response: Optional[str] = None,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> EvalContext:
+        return EvalContext(
+            prompt=prompt or self.prompt,
+            response=response or self.response,
+            metadata=metadata or self.metadata,
+        )
+
+    def hash(self):
+        return hash((self.prompt, self.response, frozenset(self.metadata.items())))
+
+    def __eq__(self, value) -> bool:
+        if not isinstance(value, EvalContext):
+            return False
+        return self.prompt == value.prompt and self.response == value.response and self.metadata == value.metadata

src/modelgauge/annotators/composer/cost.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class CostInfo:
+    input_cost_per_token: float = 0.0
+    output_cost_per_token: float = 0.0
+    fixed_cost: float = 0.0
+    latency_seconds: float = 0.0
+
+    def __add__(self, other: "CostInfo") -> "CostInfo":
+        return CostInfo(
+            input_cost_per_token=self.input_cost_per_token + other.input_cost_per_token,
+            output_cost_per_token=self.output_cost_per_token + other.output_cost_per_token,
+            fixed_cost=self.fixed_cost + other.fixed_cost,
+            latency_seconds=self.latency_seconds + other.latency_seconds,
+        )
+
+
+@dataclass
+class RealizedCost:
+    input_token_cost: float = 0.0
+    output_token_cost: float = 0.0
+    fixed_cost: float = 0.0
+    latency_seconds: float = 0.0
+
+    @property
+    def total_token_cost(self) -> float:
+        return self.input_token_cost + self.output_token_cost
+
+    @property
+    def total_cost(self) -> float:
+        return self.total_token_cost + self.fixed_cost
+
+    def __add__(self, other: "RealizedCost") -> "RealizedCost":
+        return RealizedCost(
+            input_token_cost=self.input_token_cost + other.input_token_cost,
+            output_token_cost=self.output_token_cost + other.output_token_cost,
+            fixed_cost=self.fixed_cost + other.fixed_cost,
+            latency_seconds=self.latency_seconds + other.latency_seconds,
+        )
+
+    def to_dict(self) -> dict:
+        return {
+            "input_token_cost": self.input_token_cost,
+            "output_token_cost": self.output_token_cost,
+            "fixed_cost": self.fixed_cost,
+            "latency_seconds": self.latency_seconds,
+            "total_token_cost": self.total_token_cost,
+            "total_cost": self.total_cost,
+        }