File size: 10,869 Bytes
445125c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 | """Application-owned bounded execution with real, streamable lifecycle events."""
from typing import Literal
from uuid import uuid4
from pydantic import ValidationError
from service_analysis.schemas import Contract
from service_analysis.analysis import (
CalculationMismatch, SummaryResult, ComparisonResult, _dataset_id,
summarize_resolution, compare_periods,
)
from service_analysis.decomposition import DecompositionResult, decompose_case_mix
from service_analysis.workload import WorkloadResult, analyze_workload
from service_analysis.planning import (
MAX_STEPS, ToolCall, SummaryCall, ComparisonCall, DecompositionCall,
PlanRejected, compile_plan,
)
ToolResult = SummaryResult | ComparisonResult | DecompositionResult | WorkloadResult
StepStatus = Literal["completed", "partial", "blocked", "failed", "skipped"]
class StepOutcome(Contract):
step_id: str
call: ToolCall
depends_on: tuple[str, ...]
status: StepStatus
result: ToolResult | None = None
code: str | None = None
class RunEvent(Contract):
run_id: str
sequence: int
kind: Literal["run_started", "plan_accepted", "plan_rejected", "correction_requested",
"followup_requested", "followup_declined", "planner_failed",
"step_started", "step_completed", "step_partial", "step_blocked",
"step_failed", "step_skipped", "budget_exhausted", "run_finished"]
round_number: int | None = None
step_id: str | None = None
result_id: str | None = None
code: str | None = None
class PlanningContext(Contract):
round_number: int
remaining_steps: int
outcomes: tuple[StepOutcome, ...]
rejection_code: str | None = None
class RunReport(Contract):
run_id: str
dataset_id: str
status: Literal["completed", "partial", "blocked", "failed", "rejected"]
accepted_rounds: int
correction_used: bool
tool_calls: int
reserved_steps: int
outcomes: tuple[StepOutcome, ...]
events: tuple[RunEvent, ...]
# Registry is application-owned. Plans can name only the closed ToolCall union.
def _summary(dataset, call):
return summarize_resolution(dataset, call.period, call.grouping)
def _comparison(dataset, call):
return compare_periods(dataset, call.baseline, call.comparison, call.grouping)
def _decomposition(dataset, call):
return decompose_case_mix(dataset, call.baseline, call.comparison)
def _workload(dataset, call):
return analyze_workload(dataset, call.opening_period)
_HANDLERS = {
"summarize": (_summary, SummaryResult),
"compare": (_comparison, ComparisonResult),
"decompose": (_decomposition, DecompositionResult),
"workload": (_workload, WorkloadResult),
}
def _scope(call):
if isinstance(call, (SummaryCall, ComparisonCall)):
return "resolution" if call.grouping == "overall" else call.grouping
return "resolution" if isinstance(call, DecompositionCall) else "workload"
def _validated_result(value, result_type, call, dataset_id):
if not isinstance(value, result_type):
raise ValueError("Unexpected result type")
# Revalidate even model instances; do not trust a constructed/bypassed object.
result = result_type.model_validate(value.model_dump(mode="json"))
if result.dataset_id != dataset_id:
raise ValueError("Result belongs to a different dataset")
if isinstance(call, SummaryCall):
valid = result.period == call.period and result.grouping == call.grouping
elif isinstance(call, (ComparisonCall, DecompositionCall)):
comparison = result if isinstance(call, ComparisonCall) else result.overall
valid = comparison.baseline.period == call.baseline and comparison.comparison.period == call.comparison
if isinstance(call, ComparisonCall):
valid = valid and comparison.baseline.grouping == call.grouping and comparison.comparison.grouping == call.grouping
else:
valid = result.opening_period == call.opening_period
if not valid:
raise ValueError("Result does not match requested analysis")
return result
def run_analysis_iter(dataset, initial_plan, *, correction=None, followup=None):
"""Yield actual RunEvent objects, followed by one final RunReport.
correction(context) may supply one replacement plan across the entire run.
followup(context) may supply one additional plan or None. These optional
adapters are application code, not tool names or executable model output.
No inference provider is invoked by this module. Exhaust the iterator to
finish; abandoning it performs no further work and has no resume guarantee.
"""
run_id = "run-" + uuid4().hex
dataset_id = _dataset_id(dataset)
events, outcomes, compiled_history = [], [], []
tool_calls = accepted_rounds = 0
correction_used = False
terminal_issue = None
def event(kind, **values):
value = RunEvent(run_id=run_id, sequence=len(events) + 1, kind=kind, **values)
events.append(value)
return value
def context(round_number, rejection=None):
return PlanningContext(round_number=round_number,
remaining_steps=MAX_STEPS - len(compiled_history),
outcomes=tuple(outcomes), rejection_code=rejection)
yield event("run_started")
pending = initial_plan
for round_number in (0, 1):
if round_number:
if followup is None:
break
if len(compiled_history) == MAX_STEPS:
terminal_issue = "STEP_BUDGET_EXHAUSTED"
yield event("budget_exhausted", round_number=round_number, code=terminal_issue)
break
yield event("followup_requested", round_number=round_number)
try:
pending = followup(context(round_number))
except Exception:
terminal_issue = "FOLLOWUP_ADAPTER_FAILED"
yield event("planner_failed", round_number=round_number, code=terminal_issue)
break
if pending is None:
yield event("followup_declined", round_number=round_number)
break
try:
steps = compile_plan(pending, dataset.coverage, tuple(compiled_history))
except PlanRejected as rejection:
yield event("plan_rejected", round_number=round_number, code=rejection.code)
if correction is None or correction_used:
terminal_issue = rejection.code
break
correction_used = True
yield event("correction_requested", round_number=round_number, code=rejection.code)
try:
replacement = correction(context(round_number, rejection.code))
except Exception:
terminal_issue = "CORRECTION_ADAPTER_FAILED"
yield event("planner_failed", round_number=round_number, code=terminal_issue)
break
try:
steps = compile_plan(replacement, dataset.coverage, tuple(compiled_history))
except PlanRejected as second:
terminal_issue = second.code
yield event("plan_rejected", round_number=round_number, code=second.code)
break
compiled_history.extend(steps)
accepted_rounds += 1
yield event("plan_accepted", round_number=round_number)
for step in steps:
previous = {value.step_id: value for value in outcomes}
if any(previous[dependency].status != "completed" for dependency in step.depends_on):
outcome = StepOutcome(**step.model_dump(), status="skipped", code="DEPENDENCY_NOT_COMPLETE")
else:
try:
dataset.require(_scope(step.call))
except ValueError:
outcome = StepOutcome(**step.model_dump(), status="blocked", code="DATA_QUALITY_BLOCKED")
else:
yield event("step_started", round_number=round_number, step_id=step.step_id)
tool_calls += 1
handler, result_type = _HANDLERS[step.call.operation]
try:
value = handler(dataset, step.call)
except CalculationMismatch:
outcome = StepOutcome(**step.model_dump(), status="failed", code="CALCULATION_MISMATCH")
except ValidationError:
outcome = StepOutcome(**step.model_dump(), status="failed", code="TOOL_OUTPUT_INVALID")
except ValueError:
outcome = StepOutcome(**step.model_dump(), status="blocked", code="TOOL_INPUT_REJECTED")
except Exception:
outcome = StepOutcome(**step.model_dump(), status="failed", code="TOOL_EXECUTION_FAILED")
else:
try:
result = _validated_result(value, result_type, step.call, dataset_id)
except (ValueError, TypeError, ValidationError):
outcome = StepOutcome(**step.model_dump(), status="failed", code="TOOL_OUTPUT_INVALID")
else:
status = "completed" if result.status == "complete" else "partial"
outcome = StepOutcome(**step.model_dump(), status=status, result=result)
outcomes.append(outcome)
yield event("step_" + outcome.status, round_number=round_number,
step_id=step.step_id, code=outcome.code,
result_id=outcome.result.result_id if outcome.result is not None else None)
if outcomes and all(value.status == "completed" for value in outcomes) and terminal_issue is None:
status = "completed"
elif any(value.result is not None for value in outcomes):
status = "partial"
elif any(value.status == "failed" for value in outcomes) or terminal_issue in (
"FOLLOWUP_ADAPTER_FAILED", "CORRECTION_ADAPTER_FAILED"
):
status = "failed"
elif outcomes:
status = "blocked"
else:
status = "rejected"
yield event("run_finished", code=status)
yield RunReport(run_id=run_id, dataset_id=dataset_id, status=status,
accepted_rounds=accepted_rounds, correction_used=correction_used,
tool_calls=tool_calls, reserved_steps=len(compiled_history),
outcomes=tuple(outcomes), events=tuple(events))
def run_analysis(dataset, initial_plan, *, correction=None, followup=None):
"""Convenience wrapper returning the final report from the same event loop."""
final = None
for item in run_analysis_iter(dataset, initial_plan, correction=correction, followup=followup):
if isinstance(item, RunReport):
final = item
return final
|