"""AgentState TypedDict and its Pydantic validator model. ``AgentStateModel`` is the single source of truth for field defaults and validation. ``AgentState`` composes investigation slices from :mod:`core.state.runtime_slices` and chat slice from :mod:`core.state.slices`; the runtime dict remains flat. Whenever you add or remove a field, update ``AgentStateModel`` and the appropriate slice in ``runtime_slices.py`` or ``slices.py``. ``tests/core/state/test_agent_state_sync.py`` asserts slice keys and Pydantic fields stay aligned with ``AgentState``. """ from __future__ import annotations from typing import Any, cast from pydantic import ConfigDict, Field from config.constants.investigation import MAX_INVESTIGATION_LOOPS from config.strict_config import StrictConfigModel from core.domain.types.retrieval import RetrievalControlsMap from core.state.runtime_slices import ( AlertInputSlice, CallerMetadataSlice, DeliveryContextSlice, DeliveryOutputSlice, DiagnosisSlice, EvalHarnessSlice, InvestigationPlanSlice, InvestigationRuntimeSlice, MaskingSlice, ) from core.state.slices import ChatStateSlice from core.state.types import AgentMode, ChatMessage, ChatMessageModel class AgentState( CallerMetadataSlice, ChatStateSlice, AlertInputSlice, InvestigationPlanSlice, InvestigationRuntimeSlice, DiagnosisSlice, MaskingSlice, DeliveryContextSlice, DeliveryOutputSlice, EvalHarnessSlice, total=False, ): """Unified flat state for chat and investigation modes. Chat mode primarily uses ``ChatStateSlice`` + ``CallerMetadataSlice``. Investigation mode uses alert, plan, runtime, diagnosis, and delivery slices. See :mod:`core.state.runtime_slices` for investigation field groupings. """ InvestigationState = AgentState class AgentStateModel(StrictConfigModel): """Runtime-validated state envelope used by state constructors.""" model_config = ConfigDict(extra="forbid", protected_namespaces=(), populate_by_name=True) mode: AgentMode = "chat" route: str = "" org_id: str = "" user_id: str = "" user_email: str = "" user_name: str = "" organization_slug: str = "" messages: list[ChatMessageModel] = Field(default_factory=list) is_noise: bool = False alert_name: str = "" pipeline_name: str = "" severity: str = "" alert_source: str = "" raw_alert: str | dict[str, Any] = Field(default_factory=lambda: {}) alert_json: dict[str, Any] = Field(default_factory=dict) planned_actions: list[str] = Field(default_factory=list) plan_rationale: str = "" retrieval_controls: RetrievalControlsMap | None = None available_sources: dict[str, dict[str, Any]] = Field(default_factory=dict) available_action_names: list[str] = Field(default_factory=list) tool_budget: int = Field( default=10, ge=1, le=50, description="Maximum tools to select per step" ) plan_audit: dict[str, Any] = Field( default_factory=dict, description="Audit trail for planning step" ) resolved_integrations: dict[str, Any] = Field(default_factory=dict) context: dict[str, Any] = Field( default_factory=dict, description=( "Legacy investigation evidence envelope. Not for REPL session state, " "prompt grounding, or generic runtime request metadata." ), ) evidence: dict[str, Any] = Field(default_factory=dict) correlation: dict[str, Any] = Field(default_factory=dict) root_cause: str = "" root_cause_category: str = "" validated_claims: list[dict[str, Any]] = Field(default_factory=list) non_validated_claims: list[dict[str, Any]] = Field(default_factory=list) validity_score: float = 0.0 investigation_recommendations: list[str] = Field(default_factory=list) remediation_steps: list[str] = Field(default_factory=list) triage_summary: str = "" incident_status: str = "" investigation_hypotheses: list[str] = Field(default_factory=list) verification_summary: list[str] = Field(default_factory=list) follow_up_questions: list[str] = Field(default_factory=list) remediation_tradeoffs: str = "" investigation_loop_count: int = 0 investigation_iteration_cap: int = MAX_INVESTIGATION_LOOPS hypotheses: list[str] = Field(default_factory=list) executed_hypotheses: list[dict[str, Any]] = Field(default_factory=list) evidence_entries: list[dict[str, Any]] = Field(default_factory=list) hypothesis_results: list[dict[str, Any]] = Field(default_factory=list) action_to_run: str = "" investigation_started_at: float = 0.0 incident_window: dict[str, Any] | None = None incident_window_history: list[dict[str, Any]] | None = None masking_map: dict[str, str] = Field(default_factory=dict) slack_context: dict[str, Any] = Field(default_factory=dict) discord_context: dict[str, Any] = Field(default_factory=dict) telegram_context: dict[str, Any] = Field(default_factory=dict) whatsapp_context: dict[str, Any] = Field(default_factory=dict) twilio_sms_context: dict[str, Any] = Field(default_factory=dict) openclaw_context: dict[str, Any] = Field(default_factory=dict) thread_id: str = "" run_id: str = "" auth_token: str = Field(default="", alias="_auth_token", exclude=True) slack_message: str = "" problem_md: str = "" summary: str = "" problem_report: dict[str, Any] = Field(default_factory=dict) report: str = "" opensre_evaluate: bool = False opensre_eval_rubric: str = "" opensre_llm_eval: dict[str, Any] = Field(default_factory=dict) def model_default_payload(*exclude: str) -> dict[str, Any]: """Return default field values from ``AgentStateModel``, omitting ``exclude`` keys.""" skip = frozenset(exclude) model = AgentStateModel() dumped = model.model_dump(mode="python", by_alias=True, exclude_none=True) return {key: value for key, value in dumped.items() if key not in skip} def make_chat_state( org_id: str = "", user_id: str = "", user_email: str = "", user_name: str = "", organization_slug: str = "", messages: list[ChatMessage] | None = None, ) -> AgentState: """Create initial state for chat mode.""" state = AgentStateModel.model_validate( { "mode": "chat", "org_id": org_id, "user_id": user_id, "user_email": user_email, "user_name": user_name, "organization_slug": organization_slug, "messages": messages or [], "context": {}, } ) return cast(AgentState, state.model_dump(mode="python", by_alias=True, exclude_none=True)) __all__ = [ "AgentState", "AgentStateModel", "InvestigationState", "make_chat_state", "model_default_payload", ]