weakincentives 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

weakincentives/tools/planning.py

@@ -0,0 +1,538 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Planning tool suite for session-scoped execution plans."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+from typing import Final, Literal, cast
+
+from ..prompt import SupportsDataclass
+from ..prompt.markdown import MarkdownSection
+from ..prompt.tool import Tool, ToolResult
+from ..session import Session, replace_latest, select_latest
+from ..session.session import DataEvent
+from .errors import ToolValidationError
+
+PlanStatus = Literal["active", "completed", "abandoned"]
+StepStatus = Literal["pending", "in_progress", "blocked", "done"]
+
+
+@dataclass(slots=True, frozen=True)
+class PlanStep:
+    step_id: str
+    title: str
+    details: str | None
+    status: StepStatus
+    notes: tuple[str, ...] = field(default_factory=tuple)
+
+
+@dataclass(slots=True, frozen=True)
+class Plan:
+    objective: str
+    status: PlanStatus
+    steps: tuple[PlanStep, ...] = field(default_factory=tuple)
+
+
+@dataclass(slots=True, frozen=True)
+class NewPlanStep:
+    title: str
+    details: str | None = None
+
+
+@dataclass(slots=True, frozen=True)
+class SetupPlan:
+    objective: str
+    initial_steps: tuple[NewPlanStep, ...] = field(default_factory=tuple)
+
+
+@dataclass(slots=True, frozen=True)
+class AddStep:
+    steps: tuple[NewPlanStep, ...]
+
+
+@dataclass(slots=True, frozen=True)
+class UpdateStep:
+    step_id: str
+    title: str | None = None
+    details: str | None = None
+
+
+@dataclass(slots=True, frozen=True)
+class MarkStep:
+    step_id: str
+    status: StepStatus
+    note: str | None = None
+
+
+@dataclass(slots=True, frozen=True)
+class ClearPlan:
+    pass
+
+
+@dataclass(slots=True, frozen=True)
+class ReadPlan:
+    pass
+
+
+@dataclass(slots=True, frozen=True)
+class _PlanningSectionParams:
+    """Placeholder params container for the planning tools section."""
+
+    pass
+
+
+_ASCII: Final[str] = "ascii"
+_MAX_OBJECTIVE_LENGTH: Final[int] = 240
+_MAX_TITLE_LENGTH: Final[int] = 160
+_MAX_DETAIL_LENGTH: Final[int] = 512
+_STEP_ID_PREFIX: Final[str] = "S"
+
+_PLANNING_SECTION_TEMPLATE: Final[str] = (
+    "Use planning tools for multi-step or stateful work that requires an"
+    " execution plan.\n"
+    "- Start with `planning_setup_plan` to set an objective (<=240 ASCII"
+    " characters) and optional initial steps.\n"
+    "- Keep steps concise (<=160 ASCII characters for titles, <=512 for"
+    " details).\n"
+    "- Extend plans with `planning_add_step` and refine steps with"
+    " `planning_update_step`.\n"
+    "- Track progress via `planning_mark_step` (pending, in_progress,"
+    " blocked, done).\n"
+    "- Inspect the latest plan using `planning_read_plan`.\n"
+    "- Use `planning_clear_plan` only when abandoning the objective.\n"
+    "Stay brief, ASCII-only, and skip planning for trivial single-step tasks."
+)
+
+
+class PlanningToolsSection(MarkdownSection[_PlanningSectionParams]):
+    """Prompt section exposing the planning tool suite."""
+
+    def __init__(self, *, session: Session) -> None:
+        self._session = session
+        session.register_reducer(Plan, replace_latest)
+        session.register_reducer(SetupPlan, _setup_plan_reducer, slice_type=Plan)
+        session.register_reducer(AddStep, _add_step_reducer, slice_type=Plan)
+        session.register_reducer(UpdateStep, _update_step_reducer, slice_type=Plan)
+        session.register_reducer(MarkStep, _mark_step_reducer, slice_type=Plan)
+        session.register_reducer(ClearPlan, _clear_plan_reducer, slice_type=Plan)
+
+        tools = _build_tools(session)
+        super().__init__(
+            title="Planning Tools",
+            key="planning.tools",
+            template=_PLANNING_SECTION_TEMPLATE,
+            default_params=_PlanningSectionParams(),
+            tools=tools,
+        )
+
+
+def _build_tools(
+    session: Session,
+) -> tuple[Tool[SupportsDataclass, SupportsDataclass], ...]:
+    suite = _PlanningToolSuite(session)
+    return (
+        Tool[SetupPlan, SetupPlan](
+            name="planning_setup_plan",
+            description="Create or replace the session plan.",
+            handler=suite.setup_plan,
+        ),
+        Tool[AddStep, AddStep](
+            name="planning_add_step",
+            description="Append one or more steps to the active plan.",
+            handler=suite.add_step,
+        ),
+        Tool[UpdateStep, UpdateStep](
+            name="planning_update_step",
+            description="Edit the title or details for an existing step.",
+            handler=suite.update_step,
+        ),
+        Tool[MarkStep, MarkStep](
+            name="planning_mark_step",
+            description="Update step status and optionally record a note.",
+            handler=suite.mark_step,
+        ),
+        Tool[ClearPlan, ClearPlan](
+            name="planning_clear_plan",
+            description="Mark the current plan as abandoned.",
+            handler=suite.clear_plan,
+        ),
+        Tool[ReadPlan, Plan](
+            name="planning_read_plan",
+            description="Return the latest plan snapshot.",
+            handler=suite.read_plan,
+        ),
+    )  # type: ignore[return-value]
+
+
+class _PlanningToolSuite:
+    """Collection of tool handlers bound to a session instance."""
+
+    def __init__(self, session: Session) -> None:
+        self._session = session
+
+    def setup_plan(self, params: SetupPlan) -> ToolResult[SetupPlan]:
+        objective = _normalize_required_text(
+            params.objective,
+            field_name="objective",
+            max_length=_MAX_OBJECTIVE_LENGTH,
+        )
+        initial_steps = _normalize_new_steps(params.initial_steps)
+        normalized = SetupPlan(objective=objective, initial_steps=initial_steps)
+        step_count = len(initial_steps)
+        message = (
+            f"Plan initialised with {step_count} step{'s' if step_count != 1 else ''}."
+        )
+        return ToolResult(message=message, value=normalized)
+
+    def add_step(self, params: AddStep) -> ToolResult[AddStep]:
+        plan = _require_plan(self._session)
+        _ensure_active(plan, "add steps to")
+        normalized_steps = _normalize_new_steps(params.steps)
+        if not normalized_steps:
+            raise ToolValidationError("Provide at least one step to add.")
+        message = (
+            "Queued"
+            f" {len(normalized_steps)} step{'s' if len(normalized_steps) != 1 else ''}"
+            " for addition."
+        )
+        return ToolResult(message=message, value=AddStep(steps=normalized_steps))
+
+    def update_step(self, params: UpdateStep) -> ToolResult[UpdateStep]:
+        plan = _require_plan(self._session)
+        _ensure_active(plan, "update steps in")
+        step_id = params.step_id.strip()
+        if not step_id:
+            raise ToolValidationError("Step ID must be provided.")
+        _ensure_ascii(step_id, "step_id")
+        updated_title = (
+            _normalize_required_text(
+                params.title,
+                field_name="title",
+                max_length=_MAX_TITLE_LENGTH,
+            )
+            if params.title is not None
+            else None
+        )
+        updated_details = (
+            _normalize_optional_text(
+                params.details,
+                field_name="details",
+                max_length=_MAX_DETAIL_LENGTH,
+            )
+            if params.details is not None
+            else None
+        )
+        if updated_title is None and updated_details is None:
+            raise ToolValidationError(
+                "Provide a new title or details to update a step."
+            )
+        _ensure_step_exists(plan, step_id)
+        normalized = UpdateStep(
+            step_id=step_id,
+            title=updated_title,
+            details=updated_details,
+        )
+        return ToolResult(message=f"Step {step_id} update queued.", value=normalized)
+
+    def mark_step(self, params: MarkStep) -> ToolResult[MarkStep]:
+        plan = _require_plan(self._session)
+        if plan.status == "abandoned":
+            raise ToolValidationError("Cannot mark steps on an abandoned plan.")
+        step_id = params.step_id.strip()
+        if not step_id:
+            raise ToolValidationError("Step ID must be provided.")
+        _ensure_ascii(step_id, "step_id")
+        _ensure_step_exists(plan, step_id)
+        note = _normalize_optional_text(
+            params.note,
+            field_name="note",
+            max_length=_MAX_DETAIL_LENGTH,
+            require_content=True,
+        )
+        normalized = MarkStep(step_id=step_id, status=params.status, note=note)
+        return ToolResult(
+            message=f"Step {step_id} marked as {params.status}.",
+            value=normalized,
+        )
+
+    def clear_plan(self, params: ClearPlan) -> ToolResult[ClearPlan]:
+        plan = _require_plan(self._session)
+        if plan.status == "abandoned":
+            raise ToolValidationError("Plan already abandoned.")
+        return ToolResult(message="Plan marked as abandoned.", value=params)
+
+    def read_plan(self, params: ReadPlan) -> ToolResult[Plan]:
+        del params
+        plan = select_latest(self._session, Plan)
+        if plan is None:
+            raise ToolValidationError("No plan is currently initialised.")
+        step_summary = _summarize_steps(plan.steps)
+        message = (
+            f"Objective: {plan.objective}\nStatus: {plan.status}\nSteps: {step_summary}"
+        )
+        return ToolResult(message=message, value=plan)
+
+
+def _setup_plan_reducer(
+    slice_values: tuple[Plan, ...], event: DataEvent
+) -> tuple[Plan, ...]:
+    params = cast(SetupPlan, event.value)
+    steps = tuple(
+        PlanStep(
+            step_id=_format_step_id(index + 1),
+            title=step.title,
+            details=step.details,
+            status="pending",
+            notes=(),
+        )
+        for index, step in enumerate(params.initial_steps)
+    )
+    plan = Plan(objective=params.objective, status="active", steps=steps)
+    return (plan,)
+
+
+def _add_step_reducer(
+    slice_values: tuple[Plan, ...], event: DataEvent
+) -> tuple[Plan, ...]:
+    previous = _latest_plan(slice_values)
+    if previous is None:
+        return slice_values
+    params = cast(AddStep, event.value)
+    existing = list(previous.steps)
+    next_index = _next_step_index(existing)
+    for step in params.steps:
+        next_index += 1
+        existing.append(
+            PlanStep(
+                step_id=_format_step_id(next_index),
+                title=step.title,
+                details=step.details,
+                status="pending",
+                notes=(),
+            )
+        )
+    updated = Plan(
+        objective=previous.objective,
+        status="active",
+        steps=tuple(existing),
+    )
+    return (updated,)
+
+
+def _update_step_reducer(
+    slice_values: tuple[Plan, ...], event: DataEvent
+) -> tuple[Plan, ...]:
+    previous = _latest_plan(slice_values)
+    if previous is None:
+        return slice_values
+    params = cast(UpdateStep, event.value)
+    updated_steps: list[PlanStep] = []
+    for step in previous.steps:
+        if step.step_id != params.step_id:
+            updated_steps.append(step)
+            continue
+        new_title = params.title if params.title is not None else step.title
+        new_details = params.details if params.details is not None else step.details
+        updated_steps.append(
+            PlanStep(
+                step_id=step.step_id,
+                title=new_title,
+                details=new_details,
+                status=step.status,
+                notes=step.notes,
+            )
+        )
+    updated_plan = Plan(
+        objective=previous.objective,
+        status=previous.status,
+        steps=tuple(updated_steps),
+    )
+    return (updated_plan,)
+
+
+def _mark_step_reducer(
+    slice_values: tuple[Plan, ...], event: DataEvent
+) -> tuple[Plan, ...]:
+    previous = _latest_plan(slice_values)
+    if previous is None:
+        return slice_values
+    params = cast(MarkStep, event.value)
+    updated_steps: list[PlanStep] = []
+    for step in previous.steps:
+        if step.step_id != params.step_id:
+            updated_steps.append(step)
+            continue
+        notes = step.notes
+        if params.note is not None:
+            notes = (*notes, params.note)
+        updated_steps.append(
+            PlanStep(
+                step_id=step.step_id,
+                title=step.title,
+                details=step.details,
+                status=params.status,
+                notes=notes,
+            )
+        )
+    plan_status: PlanStatus
+    if not updated_steps or all(step.status == "done" for step in updated_steps):
+        plan_status = "completed"
+    else:
+        plan_status = "active"
+    updated_plan = Plan(
+        objective=previous.objective,
+        status=plan_status,
+        steps=tuple(updated_steps),
+    )
+    return (updated_plan,)
+
+
+def _clear_plan_reducer(
+    slice_values: tuple[Plan, ...], event: DataEvent
+) -> tuple[Plan, ...]:
+    previous = _latest_plan(slice_values)
+    if previous is None:
+        return slice_values
+    del event
+    abandoned = Plan(objective=previous.objective, status="abandoned", steps=())
+    return (abandoned,)
+
+
+def _latest_plan(plans: tuple[Plan, ...]) -> Plan | None:
+    if not plans:
+        return None
+    return plans[-1]
+
+
+def _next_step_index(steps: Iterable[PlanStep]) -> int:
+    max_index = 0
+    for step in steps:
+        suffix = step.step_id[len(_STEP_ID_PREFIX) :]
+        try:
+            max_index = max(max_index, int(suffix))
+        except ValueError:
+            continue
+    return max_index
+
+
+def _format_step_id(index: int) -> str:
+    return f"{_STEP_ID_PREFIX}{index:03d}"
+
+
+def _normalize_new_steps(steps: Sequence[NewPlanStep]) -> tuple[NewPlanStep, ...]:
+    normalized: list[NewPlanStep] = []
+    for step in steps:
+        title = _normalize_required_text(
+            step.title,
+            field_name="title",
+            max_length=_MAX_TITLE_LENGTH,
+        )
+        details = _normalize_optional_text(
+            step.details,
+            field_name="details",
+            max_length=_MAX_DETAIL_LENGTH,
+        )
+        normalized.append(NewPlanStep(title=title, details=details))
+    return tuple(normalized)
+
+
+def _normalize_required_text(
+    value: str,
+    *,
+    field_name: str,
+    max_length: int,
+) -> str:
+    stripped = value.strip()
+    if not stripped:
+        raise ToolValidationError(f"{field_name.title()} must not be empty.")
+    if len(stripped) > max_length:
+        raise ToolValidationError(
+            f"{field_name.title()} must be <= {max_length} characters."
+        )
+    _ensure_ascii(stripped, field_name)
+    return stripped
+
+
+def _normalize_optional_text(
+    value: str | None,
+    *,
+    field_name: str,
+    max_length: int,
+    require_content: bool = False,
+) -> str | None:
+    if value is None:
+        return None
+    stripped = value.strip()
+    if not stripped:
+        if require_content:
+            raise ToolValidationError(
+                f"{field_name.title()} must not be empty when provided."
+            )
+        return None
+    if len(stripped) > max_length:
+        raise ToolValidationError(
+            f"{field_name.title()} must be <= {max_length} characters."
+        )
+    _ensure_ascii(stripped, field_name)
+    return stripped
+
+
+def _ensure_ascii(value: str, field_name: str) -> None:
+    try:
+        value.encode(_ASCII)
+    except UnicodeEncodeError as error:  # pragma: no cover - defensive
+        raise ToolValidationError(f"{field_name.title()} must be ASCII.") from error
+
+
+def _require_plan(session: Session) -> Plan:
+    plan = select_latest(session, Plan)
+    if plan is None:
+        raise ToolValidationError("No plan is currently initialised.")
+    return plan
+
+
+def _ensure_active(plan: Plan, action: str) -> None:
+    if plan.status != "active":
+        raise ToolValidationError(
+            f"Plan must be active to {action}. Current status: {plan.status}."
+        )
+
+
+def _ensure_step_exists(plan: Plan, step_id: str) -> None:
+    if not any(step.step_id == step_id for step in plan.steps):
+        raise ToolValidationError(f"Step {step_id} does not exist.")
+
+
+def _summarize_steps(steps: Sequence[PlanStep]) -> str:
+    if not steps:
+        return "no steps recorded"
+    summaries = [f"{step.step_id}:{step.status}" for step in steps]
+    return ", ".join(summaries)
+
+
+__all__ = [
+    "Plan",
+    "PlanStatus",
+    "PlanStep",
+    "StepStatus",
+    "NewPlanStep",
+    "SetupPlan",
+    "AddStep",
+    "UpdateStep",
+    "MarkStep",
+    "ClearPlan",
+    "ReadPlan",
+    "PlanningToolsSection",
+]