Spaces:
Running
Running
File size: 2,612 Bytes
b85de73 | 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 | from __future__ import annotations
import time
from typing import Any
from ..registry import get_agent
from ..run_logger import append_stage_trace
from ..schemas import StageTrace
def now_ms() -> int:
return int(time.time() * 1000)
def append_timeline(state: dict[str, Any], event: str) -> dict[str, Any]:
memory = dict(state.get("working_memory") or {"shared": {}, "agent_private": {}, "timeline": []})
memory.setdefault("shared", {})
memory.setdefault("agent_private", {})
timeline = list(memory.get("timeline", []))
timeline.append({"timestamp_ms": now_ms(), "event": event})
memory["timeline"] = timeline
return memory
def stage_index(state: dict[str, Any]) -> int:
return len(state.get("stage_traces", [])) + 1
def record_trace(state: dict[str, Any], trace: StageTrace) -> list[dict[str, Any]]:
data = trace.model_dump()
run_id = state.get("run_metadata", {}).get("run_id")
if run_id:
append_stage_trace(str(run_id), data)
return list(state.get("stage_traces", [])) + [data]
def disabled_update(state: dict[str, Any], agent_key: str) -> dict[str, Any]:
agent = get_agent(agent_key)
trace = StageTrace(
iteration=int(state.get("iterations", 1) or 1),
stage_index=stage_index(state),
agent_key=agent_key,
agent_name=agent.name,
status="skipped",
output_summary=["Agent execution mode is Disabled."],
agent_explanation="The agent was skipped by configuration.",
)
return {
"stage_traces": record_trace(state, trace),
"working_memory": append_timeline(state, f"{agent.name} skipped."),
}
def failure_update(
state: dict[str, Any],
agent_key: str,
started_ms: int,
error: Exception,
) -> dict[str, Any]:
agent = get_agent(agent_key)
error_data = {"agent_key": agent_key, "message": str(error), "timestamp_ms": now_ms()}
trace = StageTrace(
iteration=int(state.get("iterations", 1) or 1),
stage_index=stage_index(state),
agent_key=agent_key,
agent_name=agent.name,
status="failed",
output_summary=[str(error)],
decision_basis=["The node raised an exception or returned invalid structured output."],
agent_explanation="The agent failed before producing a valid schema-conformant result.",
duration_ms=now_ms() - started_ms,
)
return {
"errors": list(state.get("errors", [])) + [error_data],
"stage_traces": record_trace(state, trace),
"working_memory": append_timeline(state, f"{agent.name} failed: {error}"),
}
|