Files
nvi/api/plan/types.py

51 lines
1.4 KiB
Python

"""Public data shapes for the Plan layer."""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Optional
@dataclass
class PlanStep:
analysis: str
args: dict[str, Any] = field(default_factory=dict)
why: str = ""
# Optional sibling step to run when the primary errors or yields no
# finding. Lets the planner hedge between a speculative Analysis (e.g.
# `drill_down`) and a safer one (e.g. `direct_answer`).
fallback: Optional["PlanStep"] = None
@classmethod
def from_dict(cls, raw: dict[str, Any]) -> "PlanStep":
fb_raw = raw.get("fallback")
return cls(
analysis=raw["analysis"],
args=raw.get("args") or {},
why=raw.get("why", ""),
fallback=cls.from_dict(fb_raw) if fb_raw else None,
)
def to_dict(self) -> dict[str, Any]:
out: dict[str, Any] = {
"analysis": self.analysis,
"args": self.args,
"why": self.why,
}
if self.fallback is not None:
out["fallback"] = self.fallback.to_dict()
return out
@dataclass
class Plan:
rationale: str
steps: list[PlanStep]
@classmethod
def from_dict(cls, raw: dict[str, Any]) -> "Plan":
return cls(
rationale=raw.get("rationale", ""),
steps=[PlanStep.from_dict(s) for s in raw.get("steps", [])],
)