566 lines
17 KiB
Python
566 lines
17 KiB
Python
"""
|
|
Stream Runtime - Thread-safe runtime for concurrent executions.
|
|
|
|
Unlike the original Runtime which has a single _current_run,
|
|
StreamRuntime tracks runs by execution_id, allowing concurrent
|
|
executions within the same stream without collision.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from framework.observability import set_trace_context
|
|
from framework.schemas.decision import Decision, DecisionType, Option, Outcome
|
|
from framework.schemas.run import Run, RunStatus
|
|
from framework.storage.concurrent import ConcurrentStorage
|
|
|
|
if TYPE_CHECKING:
|
|
from framework.host.outcome_aggregator import OutcomeAggregator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamDecisionTracker:
|
|
"""
|
|
Thread-safe runtime for a single execution stream.
|
|
|
|
Key differences from Runtime:
|
|
- Tracks multiple runs concurrently via execution_id
|
|
- Uses ConcurrentStorage for thread-safe persistence
|
|
- Reports decisions to OutcomeAggregator for cross-stream evaluation
|
|
|
|
Example:
|
|
runtime = StreamRuntime(
|
|
stream_id="webhook",
|
|
storage=concurrent_storage,
|
|
outcome_aggregator=aggregator,
|
|
)
|
|
|
|
# Start a run for a specific execution
|
|
run_id = runtime.start_run(
|
|
execution_id="exec_123",
|
|
goal_id="support-goal",
|
|
goal_description="Handle support tickets",
|
|
)
|
|
|
|
# Record decisions (thread-safe)
|
|
decision_id = runtime.decide(
|
|
execution_id="exec_123",
|
|
intent="Classify ticket",
|
|
options=[...],
|
|
chosen="howto",
|
|
reasoning="Question matches how-to pattern",
|
|
)
|
|
|
|
# Record outcome
|
|
runtime.record_outcome(
|
|
execution_id="exec_123",
|
|
decision_id=decision_id,
|
|
success=True,
|
|
result={"category": "howto"},
|
|
)
|
|
|
|
# End run
|
|
runtime.end_run(
|
|
execution_id="exec_123",
|
|
success=True,
|
|
narrative="Ticket resolved",
|
|
)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stream_id: str,
|
|
storage: ConcurrentStorage,
|
|
outcome_aggregator: "OutcomeAggregator | None" = None,
|
|
):
|
|
"""
|
|
Initialize stream runtime.
|
|
|
|
Args:
|
|
stream_id: Unique identifier for this stream
|
|
storage: Concurrent storage backend
|
|
outcome_aggregator: Optional aggregator for cross-stream evaluation
|
|
"""
|
|
self.stream_id = stream_id
|
|
self._storage = storage
|
|
self._outcome_aggregator = outcome_aggregator
|
|
|
|
# Track runs by execution_id (thread-safe via lock)
|
|
self._runs: dict[str, Run] = {}
|
|
self._run_locks: dict[str, asyncio.Lock] = {}
|
|
self._global_lock = asyncio.Lock()
|
|
|
|
# Track current node per execution (for decision context)
|
|
self._current_nodes: dict[str, str] = {}
|
|
|
|
# === RUN LIFECYCLE ===
|
|
|
|
def start_run(
|
|
self,
|
|
execution_id: str,
|
|
goal_id: str,
|
|
goal_description: str = "",
|
|
input_data: dict[str, Any] | None = None,
|
|
) -> str:
|
|
"""
|
|
Start a new run for an execution.
|
|
|
|
Args:
|
|
execution_id: Unique execution identifier
|
|
goal_id: The ID of the goal being pursued
|
|
goal_description: Human-readable description of the goal
|
|
input_data: Initial input to the run
|
|
|
|
Returns:
|
|
The run ID
|
|
"""
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
run_id = f"run_{self.stream_id}_{timestamp}_{uuid.uuid4().hex[:8]}"
|
|
trace_id = uuid.uuid4().hex
|
|
otel_execution_id = uuid.uuid4().hex # 32 hex, OTel/W3C-aligned for logs
|
|
|
|
set_trace_context(
|
|
trace_id=trace_id,
|
|
execution_id=otel_execution_id,
|
|
run_id=run_id,
|
|
goal_id=goal_id,
|
|
stream_id=self.stream_id,
|
|
)
|
|
|
|
run = Run(
|
|
id=run_id,
|
|
goal_id=goal_id,
|
|
goal_description=goal_description,
|
|
input_data=input_data or {},
|
|
)
|
|
|
|
self._runs[execution_id] = run
|
|
self._run_locks[execution_id] = asyncio.Lock()
|
|
self._current_nodes[execution_id] = "unknown"
|
|
|
|
logger.debug(
|
|
f"Started run {run_id} for execution {execution_id} in stream {self.stream_id}"
|
|
)
|
|
return run_id
|
|
|
|
def end_run(
|
|
self,
|
|
execution_id: str,
|
|
success: bool,
|
|
narrative: str = "",
|
|
output_data: dict[str, Any] | None = None,
|
|
) -> None:
|
|
"""
|
|
End a run for an execution.
|
|
|
|
Args:
|
|
execution_id: Execution identifier
|
|
success: Whether the run achieved its goal
|
|
narrative: Human-readable summary of what happened
|
|
output_data: Final output of the run
|
|
"""
|
|
run = self._runs.get(execution_id)
|
|
if run is None:
|
|
logger.warning(f"end_run called but no run for execution {execution_id}")
|
|
return
|
|
|
|
status = RunStatus.COMPLETED if success else RunStatus.FAILED
|
|
run.output_data = output_data or {}
|
|
run.complete(status, narrative)
|
|
|
|
# Save to storage asynchronously
|
|
asyncio.create_task(self._save_run(execution_id, run))
|
|
|
|
logger.debug(f"Ended run {run.id} for execution {execution_id}: {status.value}")
|
|
|
|
async def _save_run(self, execution_id: str, run: Run) -> None:
|
|
"""Save run to storage and clean up."""
|
|
try:
|
|
await self._storage.save_run(run)
|
|
except Exception as e:
|
|
logger.error(f"Failed to save run {run.id}: {e}")
|
|
finally:
|
|
# Clean up
|
|
self._runs.pop(execution_id, None)
|
|
self._run_locks.pop(execution_id, None)
|
|
self._current_nodes.pop(execution_id, None)
|
|
|
|
def set_node(self, execution_id: str, node_id: str) -> None:
|
|
"""Set the current node context for an execution."""
|
|
self._current_nodes[execution_id] = node_id
|
|
|
|
def get_run(self, execution_id: str) -> Run | None:
|
|
"""Get the current run for an execution."""
|
|
return self._runs.get(execution_id)
|
|
|
|
# === DECISION RECORDING ===
|
|
|
|
def decide(
|
|
self,
|
|
execution_id: str,
|
|
intent: str,
|
|
options: list[dict[str, Any]],
|
|
chosen: str,
|
|
reasoning: str,
|
|
node_id: str | None = None,
|
|
decision_type: DecisionType = DecisionType.CUSTOM,
|
|
constraints: list[str] | None = None,
|
|
context: dict[str, Any] | None = None,
|
|
) -> str:
|
|
"""
|
|
Record a decision for a specific execution.
|
|
|
|
Thread-safe: Multiple executions can record decisions concurrently.
|
|
|
|
Args:
|
|
execution_id: Which execution is making this decision
|
|
intent: What the agent was trying to accomplish
|
|
options: List of options considered
|
|
chosen: ID of the chosen option
|
|
reasoning: Why the agent chose this option
|
|
node_id: Which node made this decision
|
|
decision_type: Type of decision
|
|
constraints: Active constraints that influenced the decision
|
|
context: Additional context available when deciding
|
|
|
|
Returns:
|
|
The decision ID, or empty string if no run in progress
|
|
"""
|
|
run = self._runs.get(execution_id)
|
|
if run is None:
|
|
logger.warning(f"decide called but no run for execution {execution_id}: {intent}")
|
|
return ""
|
|
|
|
# Build Option objects
|
|
option_objects = []
|
|
for opt in options:
|
|
option_objects.append(
|
|
Option(
|
|
id=opt["id"],
|
|
description=opt.get("description", ""),
|
|
action_type=opt.get("action_type", "unknown"),
|
|
action_params=opt.get("action_params", {}),
|
|
pros=opt.get("pros", []),
|
|
cons=opt.get("cons", []),
|
|
confidence=opt.get("confidence", 0.5),
|
|
)
|
|
)
|
|
|
|
# Create decision
|
|
decision_id = f"dec_{len(run.decisions)}"
|
|
current_node = node_id or self._current_nodes.get(execution_id, "unknown")
|
|
|
|
decision = Decision(
|
|
id=decision_id,
|
|
node_id=current_node,
|
|
intent=intent,
|
|
decision_type=decision_type,
|
|
options=option_objects,
|
|
chosen_option_id=chosen,
|
|
reasoning=reasoning,
|
|
active_constraints=constraints or [],
|
|
input_context=context or {},
|
|
)
|
|
|
|
run.add_decision(decision)
|
|
|
|
# Report to outcome aggregator if available
|
|
if self._outcome_aggregator:
|
|
self._outcome_aggregator.record_decision(
|
|
stream_id=self.stream_id,
|
|
execution_id=execution_id,
|
|
decision=decision,
|
|
)
|
|
|
|
return decision_id
|
|
|
|
def record_outcome(
|
|
self,
|
|
execution_id: str,
|
|
decision_id: str,
|
|
success: bool,
|
|
result: Any = None,
|
|
error: str | None = None,
|
|
summary: str = "",
|
|
state_changes: dict[str, Any] | None = None,
|
|
tokens_used: int = 0,
|
|
latency_ms: int = 0,
|
|
) -> None:
|
|
"""
|
|
Record the outcome of a decision.
|
|
|
|
Args:
|
|
execution_id: Which execution
|
|
decision_id: ID returned from decide()
|
|
success: Whether the action succeeded
|
|
result: The actual result/output
|
|
error: Error message if failed
|
|
summary: Human-readable summary of what happened
|
|
state_changes: What state changed as a result
|
|
tokens_used: LLM tokens consumed
|
|
latency_ms: Time taken in milliseconds
|
|
"""
|
|
run = self._runs.get(execution_id)
|
|
if run is None:
|
|
logger.warning(f"record_outcome called but no run for execution {execution_id}")
|
|
return
|
|
|
|
outcome = Outcome(
|
|
success=success,
|
|
result=result,
|
|
error=error,
|
|
summary=summary,
|
|
state_changes=state_changes or {},
|
|
tokens_used=tokens_used,
|
|
latency_ms=latency_ms,
|
|
)
|
|
|
|
run.record_outcome(decision_id, outcome)
|
|
|
|
# Report to outcome aggregator if available
|
|
if self._outcome_aggregator:
|
|
self._outcome_aggregator.record_outcome(
|
|
stream_id=self.stream_id,
|
|
execution_id=execution_id,
|
|
decision_id=decision_id,
|
|
outcome=outcome,
|
|
)
|
|
|
|
# === PROBLEM RECORDING ===
|
|
|
|
def report_problem(
|
|
self,
|
|
execution_id: str,
|
|
severity: str,
|
|
description: str,
|
|
decision_id: str | None = None,
|
|
root_cause: str | None = None,
|
|
suggested_fix: str | None = None,
|
|
) -> str:
|
|
"""
|
|
Report a problem that occurred during an execution.
|
|
|
|
Args:
|
|
execution_id: Which execution
|
|
severity: "critical", "warning", or "minor"
|
|
description: What went wrong
|
|
decision_id: Which decision caused this (if known)
|
|
root_cause: Why it went wrong (if known)
|
|
suggested_fix: What might fix it (if known)
|
|
|
|
Returns:
|
|
The problem ID, or empty string if no run in progress
|
|
"""
|
|
run = self._runs.get(execution_id)
|
|
if run is None:
|
|
logger.warning(
|
|
f"report_problem called but no run for execution {execution_id}: "
|
|
f"[{severity}] {description}"
|
|
)
|
|
return ""
|
|
|
|
return run.add_problem(
|
|
severity=severity,
|
|
description=description,
|
|
decision_id=decision_id,
|
|
root_cause=root_cause,
|
|
suggested_fix=suggested_fix,
|
|
)
|
|
|
|
# === CONVENIENCE METHODS ===
|
|
|
|
def quick_decision(
|
|
self,
|
|
execution_id: str,
|
|
intent: str,
|
|
action: str,
|
|
reasoning: str,
|
|
node_id: str | None = None,
|
|
) -> str:
|
|
"""
|
|
Record a simple decision with a single action.
|
|
|
|
Args:
|
|
execution_id: Which execution
|
|
intent: What the agent is trying to do
|
|
action: What it's doing
|
|
reasoning: Why
|
|
|
|
Returns:
|
|
The decision ID
|
|
"""
|
|
return self.decide(
|
|
execution_id=execution_id,
|
|
intent=intent,
|
|
options=[
|
|
{
|
|
"id": "action",
|
|
"description": action,
|
|
"action_type": "execute",
|
|
}
|
|
],
|
|
chosen="action",
|
|
reasoning=reasoning,
|
|
node_id=node_id,
|
|
)
|
|
|
|
# === STATS AND MONITORING ===
|
|
|
|
def get_active_executions(self) -> list[str]:
|
|
"""Get list of active execution IDs."""
|
|
return list(self._runs.keys())
|
|
|
|
def get_stats(self) -> dict:
|
|
"""Get runtime statistics."""
|
|
return {
|
|
"stream_id": self.stream_id,
|
|
"active_executions": len(self._runs),
|
|
"execution_ids": list(self._runs.keys()),
|
|
}
|
|
|
|
|
|
class StreamRuntimeAdapter:
|
|
"""
|
|
Adapter to make StreamRuntime compatible with existing Runtime interface.
|
|
|
|
This allows StreamRuntime to be used with existing GraphExecutor code
|
|
by providing the same API as Runtime but routing to a specific execution.
|
|
"""
|
|
|
|
def __init__(self, stream_runtime: StreamDecisionTracker, execution_id: str):
|
|
"""
|
|
Create adapter for a specific execution.
|
|
|
|
Args:
|
|
stream_runtime: The underlying stream runtime
|
|
execution_id: Which execution this adapter is for
|
|
"""
|
|
self._runtime = stream_runtime
|
|
self._execution_id = execution_id
|
|
self._current_node = "unknown"
|
|
|
|
# Expose storage for compatibility
|
|
@property
|
|
def storage(self):
|
|
return self._runtime._storage
|
|
|
|
@property
|
|
def execution_id(self) -> str:
|
|
return self._execution_id
|
|
|
|
@property
|
|
def current_run(self) -> Run | None:
|
|
return self._runtime.get_run(self._execution_id)
|
|
|
|
def start_run(
|
|
self,
|
|
goal_id: str,
|
|
goal_description: str = "",
|
|
input_data: dict[str, Any] | None = None,
|
|
) -> str:
|
|
return self._runtime.start_run(
|
|
execution_id=self._execution_id,
|
|
goal_id=goal_id,
|
|
goal_description=goal_description,
|
|
input_data=input_data,
|
|
)
|
|
|
|
def end_run(
|
|
self,
|
|
success: bool,
|
|
narrative: str = "",
|
|
output_data: dict[str, Any] | None = None,
|
|
) -> None:
|
|
self._runtime.end_run(
|
|
execution_id=self._execution_id,
|
|
success=success,
|
|
narrative=narrative,
|
|
output_data=output_data,
|
|
)
|
|
|
|
def set_node(self, node_id: str) -> None:
|
|
self._current_node = node_id
|
|
self._runtime.set_node(self._execution_id, node_id)
|
|
|
|
def decide(
|
|
self,
|
|
intent: str,
|
|
options: list[dict[str, Any]],
|
|
chosen: str,
|
|
reasoning: str,
|
|
node_id: str | None = None,
|
|
decision_type: DecisionType = DecisionType.CUSTOM,
|
|
constraints: list[str] | None = None,
|
|
context: dict[str, Any] | None = None,
|
|
) -> str:
|
|
return self._runtime.decide(
|
|
execution_id=self._execution_id,
|
|
intent=intent,
|
|
options=options,
|
|
chosen=chosen,
|
|
reasoning=reasoning,
|
|
node_id=node_id or self._current_node,
|
|
decision_type=decision_type,
|
|
constraints=constraints,
|
|
context=context,
|
|
)
|
|
|
|
def record_outcome(
|
|
self,
|
|
decision_id: str,
|
|
success: bool,
|
|
result: Any = None,
|
|
error: str | None = None,
|
|
summary: str = "",
|
|
state_changes: dict[str, Any] | None = None,
|
|
tokens_used: int = 0,
|
|
latency_ms: int = 0,
|
|
) -> None:
|
|
self._runtime.record_outcome(
|
|
execution_id=self._execution_id,
|
|
decision_id=decision_id,
|
|
success=success,
|
|
result=result,
|
|
error=error,
|
|
summary=summary,
|
|
state_changes=state_changes,
|
|
tokens_used=tokens_used,
|
|
latency_ms=latency_ms,
|
|
)
|
|
|
|
def report_problem(
|
|
self,
|
|
severity: str,
|
|
description: str,
|
|
decision_id: str | None = None,
|
|
root_cause: str | None = None,
|
|
suggested_fix: str | None = None,
|
|
) -> str:
|
|
return self._runtime.report_problem(
|
|
execution_id=self._execution_id,
|
|
severity=severity,
|
|
description=description,
|
|
decision_id=decision_id,
|
|
root_cause=root_cause,
|
|
suggested_fix=suggested_fix,
|
|
)
|
|
|
|
def quick_decision(
|
|
self,
|
|
intent: str,
|
|
action: str,
|
|
reasoning: str,
|
|
node_id: str | None = None,
|
|
) -> str:
|
|
return self._runtime.quick_decision(
|
|
execution_id=self._execution_id,
|
|
intent=intent,
|
|
action=action,
|
|
reasoning=reasoning,
|
|
node_id=node_id or self._current_node,
|
|
)
|