diff --git a/pebble/orchestrator/schemas.py b/pebble/orchestrator/schemas.py new file mode 100644 index 00000000..83b7e8d2 --- /dev/null +++ b/pebble/orchestrator/schemas.py @@ -0,0 +1,221 @@ +"""Pydantic models for the Pebble chat orchestrator. + +The shapes here are the agent's DURABLE contract — what gets persisted +to ``bedrock.pebble_chat_scratchpad``, what the planner emits, what +the executor consumes. Tests assert against these. Frontend types +mirror the Plan / PlanStep shapes 1:1 so the agent's plan view +renders without translation. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from enum import Enum +from typing import Any, Optional +from uuid import UUID, uuid4 + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class StepType(str, Enum): + PLAN = "plan" + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + EVALUATION = "evaluation" + RENDER = "render" + CONFLICT = "conflict" + CHECKPOINT = "checkpoint" + ERROR = "error" + + +class EvalVerdict(str, Enum): + PASS = "pass" + RETRY = "retry" + ABORT = "abort" + + +# --------------------------------------------------------------------------- +# Plan + PlanStep +# --------------------------------------------------------------------------- + +class PlanStep(BaseModel): + """A single step in an agent plan. The planner emits these; the + executor invokes them in order, respecting depends_on. + """ + model_config = ConfigDict(frozen=True) + + step_id: UUID = Field(default_factory=uuid4) + tool: str + args: dict[str, Any] = Field(default_factory=dict) + expected_shape: str = "" # human-readable description + success_criteria: str = "" # what makes this step's result useful + depends_on: tuple[UUID, ...] = () # IDs of prior steps that must complete first + + @field_validator("tool") + @classmethod + def _tool_not_empty(cls, v: str) -> str: + if not v or not v.strip(): + raise ValueError("PlanStep.tool must be non-empty") + return v.strip() + + +class Plan(BaseModel): + """The full plan emitted by the planner before any tool executes. + + ``rationale`` is the planner's brief reasoning summary — for + transparency in the FE plan view. Not load-bearing for execution. + + ``estimated_cost_usd`` and ``estimated_tool_calls`` feed the budget + pre-flight: if a plan's estimate already exceeds the conversation + budget, the executor refuses to start and asks the user to narrow. + """ + model_config = ConfigDict(frozen=True) + + plan_id: UUID = Field(default_factory=uuid4) + user_query: str + steps: tuple[PlanStep, ...] + rationale: str = "" + estimated_cost_usd: float = 0.0 + estimated_tool_calls: int = 0 + created_at: datetime = Field( + default_factory=lambda: datetime.now(tz=timezone.utc), + ) + + @field_validator("user_query") + @classmethod + def _query_not_empty(cls, v: str) -> str: + if not v or not v.strip(): + raise ValueError("Plan.user_query must be non-empty") + return v.strip() + + @field_validator("steps") + @classmethod + def _steps_consistent(cls, v: tuple[PlanStep, ...]) -> tuple[PlanStep, ...]: + # depends_on must reference earlier step_ids only. + seen: set[UUID] = set() + for step in v: + for dep in step.depends_on: + if dep not in seen: + raise ValueError( + f"PlanStep {step.step_id} depends_on {dep} which is " + "not a prior step", + ) + seen.add(step.step_id) + return v + + +# --------------------------------------------------------------------------- +# Tool calls + results +# --------------------------------------------------------------------------- + +class ToolCall(BaseModel): + """Persisted record of an executor tool invocation.""" + model_config = ConfigDict(frozen=True) + + step_id: UUID + tool: str + args: dict[str, Any] + plan_step_id: Optional[UUID] = None # which plan step authorized this + invoked_at: datetime = Field( + default_factory=lambda: datetime.now(tz=timezone.utc), + ) + + +class ToolResult(BaseModel): + """Persisted record of a tool's output. Failure-by-default; + successful results carry data.""" + model_config = ConfigDict(frozen=True) + + step_id: UUID + tool: str + ok: bool + data: Optional[dict[str, Any]] = None + error: Optional[str] = None + duration_ms: int = 0 + cost_usd: float = 0.0 + tokens_in: int = 0 + tokens_out: int = 0 + citations: tuple[str, ...] = () # IDs the renderer can render as + + +# --------------------------------------------------------------------------- +# Evaluation +# --------------------------------------------------------------------------- + +class Evaluation(BaseModel): + """Output of the evaluator (Haiku-as-judge). + + ``cost_usd`` / ``tokens_in`` / ``tokens_out`` carry the LLM-call + accounting for the eval pass itself — same shape as + ``ToolResult``'s cost fields. The chat orchestrator surfaces these + in the ``eval_emitted`` SSE event so the frontend can render a + running cost / token tally per conversation. + """ + model_config = ConfigDict(frozen=True) + + plan_id: UUID + factuality: float = Field(ge=0.0, le=1.0) + completeness: float = Field(ge=0.0, le=1.0) + harm: str = Field(default="none", pattern=r"^(none|mild|severe)$") + verdict: EvalVerdict + rationale: str = "" + rejected_claims: tuple[str, ...] = () + cost_usd: float = 0.0 + tokens_in: int = 0 + tokens_out: int = 0 + + +# --------------------------------------------------------------------------- +# Final response shape (what the renderer emits, what the FE consumes) +# --------------------------------------------------------------------------- + +class Citation(BaseModel): + """A citation reference attached to a span of text.""" + model_config = ConfigDict(frozen=True) + + cite_id: str + entity_type: str # 'sf_account', 'pebble_profile', 'metric:stale_pipeline', ... + entity_id: str + title: str = "" + href: str = "" + + +class SuggestedAction(BaseModel): + """A write the agent proposes; the FE renders as a confirm card. + User confirmation re-issues the call with their JWT, NOT the + internal key. Until then the agent has not written anything. + """ + model_config = ConfigDict(frozen=True) + + action_id: UUID = Field(default_factory=uuid4) + kind: str # 'update_stage', 'create_task', 'send_email', ... + payload: dict[str, Any] + diff_preview: str # human-readable "Stage: A → B" + record_label: str # "Acme Corp · 006XYZ" for anti-mistake guard + rationale: str = "" + + +class ChartSpec(BaseModel): + """Recharts-shape JSON. The FE uses ``kind`` to pick a component.""" + model_config = ConfigDict(frozen=True) + + chart_id: UUID = Field(default_factory=uuid4) + kind: str = Field(pattern=r"^(line|bar|pie|area|scatter|funnel)$") + title: str = "" + data: list[dict[str, Any]] = Field(default_factory=list) + x_key: Optional[str] = None + y_keys: tuple[str, ...] = () + + +class FinalResponse(BaseModel): + """Stitched from the renderer; persisted as the conversation's + last scratchpad row of step_type=render.""" + model_config = ConfigDict(frozen=True) + + plan_id: UUID + text: str + citations: tuple[Citation, ...] = () + suggested_actions: tuple[SuggestedAction, ...] = () + charts: tuple[ChartSpec, ...] = () + degraded: bool = False + degradation_reason: Optional[str] = None diff --git a/pebble/tests/test_orchestrator_schemas.py b/pebble/tests/test_orchestrator_schemas.py new file mode 100644 index 00000000..ff050057 --- /dev/null +++ b/pebble/tests/test_orchestrator_schemas.py @@ -0,0 +1,217 @@ +"""Pydantic-shape tests for ``pebble.orchestrator.schemas``. + +The schemas are durable contracts (persisted, FE-shared, planner-shared) +so the constraints they encode must be enforced. +""" + +import os +import sys +from uuid import uuid4 + +import pytest +from pydantic import ValidationError + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..')) + +from pebble.orchestrator.schemas import ( + Citation, ChartSpec, EvalVerdict, Evaluation, FinalResponse, + Plan, PlanStep, StepType, SuggestedAction, ToolCall, ToolResult, +) + + +# --------------------------------------------------------------------------- +# PlanStep +# --------------------------------------------------------------------------- + +def test_plan_step_minimal_valid(): + step = PlanStep(tool="search_crm") + assert step.tool == "search_crm" + assert step.args == {} + assert step.depends_on == () + assert step.step_id # uuid auto-generated + + +def test_plan_step_strips_tool_whitespace(): + step = PlanStep(tool=" search_crm ") + assert step.tool == "search_crm" + + +def test_plan_step_rejects_empty_tool(): + with pytest.raises(ValidationError, match=r"non-empty"): + PlanStep(tool="") + + +def test_plan_step_rejects_whitespace_only_tool(): + with pytest.raises(ValidationError, match=r"non-empty"): + PlanStep(tool=" ") + + +def test_plan_step_frozen(): + """Frozen=True so steps can't be mutated after planning.""" + step = PlanStep(tool="search_crm") + with pytest.raises(ValidationError): + step.tool = "different" + + +# --------------------------------------------------------------------------- +# Plan +# --------------------------------------------------------------------------- + +def test_plan_minimal_valid(): + plan = Plan( + user_query="show open deals", + steps=(PlanStep(tool="search_crm"),), + ) + assert plan.user_query == "show open deals" + assert len(plan.steps) == 1 + + +def test_plan_strips_query_whitespace(): + plan = Plan(user_query=" show open deals ", steps=()) + assert plan.user_query == "show open deals" + + +def test_plan_rejects_empty_query(): + with pytest.raises(ValidationError, match=r"non-empty"): + Plan(user_query="", steps=()) + + +def test_plan_rejects_forward_reference_in_depends_on(): + """depends_on can only reference earlier steps. Forward refs are + a bug — the planner must emit topologically ordered steps.""" + later_id = uuid4() + earlier = PlanStep(tool="search_crm") + later = PlanStep(tool="get_record", depends_on=(later_id,)) + with pytest.raises(ValidationError, match=r"depends_on"): + Plan(user_query="x", steps=(earlier, later)) + + +def test_plan_accepts_valid_dependency_chain(): + a = PlanStep(tool="search_crm") + b = PlanStep(tool="get_record", depends_on=(a.step_id,)) + c = PlanStep(tool="generate_chart", depends_on=(a.step_id, b.step_id)) + plan = Plan(user_query="x", steps=(a, b, c)) + assert plan.steps[2].depends_on == (a.step_id, b.step_id) + + +def test_plan_frozen(): + plan = Plan(user_query="x", steps=(PlanStep(tool="search_crm"),)) + with pytest.raises(ValidationError): + plan.user_query = "different" + + +# --------------------------------------------------------------------------- +# ToolResult +# --------------------------------------------------------------------------- + +def test_tool_result_failure_default(): + r = ToolResult(step_id=uuid4(), tool="search_crm", ok=False, error="boom") + assert r.ok is False + assert r.data is None + assert r.error == "boom" + + +def test_tool_result_success_carries_data(): + r = ToolResult( + step_id=uuid4(), tool="search_crm", ok=True, + data={"items": [{"id": "001"}]}, + citations=("hit_1",), + duration_ms=42, + cost_usd=0.001, + ) + assert r.ok is True + assert r.data["items"][0]["id"] == "001" + assert r.citations == ("hit_1",) + + +def test_tool_result_frozen(): + r = ToolResult(step_id=uuid4(), tool="x", ok=True) + with pytest.raises(ValidationError): + r.ok = False + + +# --------------------------------------------------------------------------- +# Evaluation +# --------------------------------------------------------------------------- + +def test_evaluation_factuality_bounds(): + with pytest.raises(ValidationError): + Evaluation( + plan_id=uuid4(), factuality=1.5, completeness=0.5, + verdict=EvalVerdict.PASS, + ) + with pytest.raises(ValidationError): + Evaluation( + plan_id=uuid4(), factuality=-0.1, completeness=0.5, + verdict=EvalVerdict.PASS, + ) + + +def test_evaluation_harm_pattern(): + with pytest.raises(ValidationError): + Evaluation( + plan_id=uuid4(), factuality=1.0, completeness=1.0, + verdict=EvalVerdict.PASS, harm="catastrophic", # not in enum + ) + + +@pytest.mark.parametrize("harm", ["none", "mild", "severe"]) +def test_evaluation_harm_accepted_values(harm): + e = Evaluation( + plan_id=uuid4(), factuality=1.0, completeness=1.0, + verdict=EvalVerdict.PASS, harm=harm, + ) + assert e.harm == harm + + +@pytest.mark.parametrize("verdict", [EvalVerdict.PASS, EvalVerdict.RETRY, EvalVerdict.ABORT]) +def test_evaluation_verdict_enum(verdict): + e = Evaluation( + plan_id=uuid4(), factuality=0.5, completeness=0.5, verdict=verdict, + ) + assert e.verdict == verdict + + +# --------------------------------------------------------------------------- +# Citation, ChartSpec, SuggestedAction, FinalResponse +# --------------------------------------------------------------------------- + +def test_citation_minimal(): + c = Citation(cite_id="c1", entity_type="sf_account", entity_id="001ABC") + assert c.cite_id == "c1" + + +def test_chart_spec_kind_pattern(): + with pytest.raises(ValidationError): + ChartSpec(kind="bogus") + ChartSpec(kind="bar") + ChartSpec(kind="line") + ChartSpec(kind="pie") + + +def test_suggested_action_carries_rationale(): + a = SuggestedAction( + kind="update_stage", + payload={"opportunity_id": "006XYZ", "new_stage": "Closed Won"}, + diff_preview="Stage: Ask in Progress → Closed Won", + record_label="Acme · 006XYZ", + rationale="3 prior contacts moved this account to verbal commitment.", + ) + assert a.kind == "update_stage" + + +def test_final_response_immutable(): + fr = FinalResponse(plan_id=uuid4(), text="hi") + with pytest.raises(ValidationError): + fr.text = "different" + + +def test_step_type_enum_values(): + """The DB CHECK constraint on bedrock.pebble_chat_scratchpad.step_type + must accept exactly these values. If this set drifts, the migration + needs an update — this test is the canary.""" + expected = { + "plan", "tool_call", "tool_result", "evaluation", + "render", "conflict", "checkpoint", "error", + } + assert {st.value for st in StepType} == expected