Spaces:
Running
Running
File size: 4,860 Bytes
33dd3ee | 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 | """Pydantic models for OpenEnv-Sentinel.
Typed discriminated union for actions — each tool has its own action + params class.
"""
from typing import Annotated, List, Literal, Union
from openenv.core.env_server.types import Action, Observation, State
from pydantic import BaseModel, Field, RootModel
# ── Parameter models ────────────────────────────────────────────────
class QueryLogsParams(BaseModel):
service: str
query: str = "all"
severity: str = "all"
class QueryMetricsParams(BaseModel):
service: str
metric: str
class GetServiceStatusParams(BaseModel):
service: str
class GetDependencyMapParams(BaseModel):
service: str = ""
class ConsultRunbookParams(BaseModel):
topic: str
class CheckRecentChangesParams(BaseModel):
service: str = ""
class SubmitResolutionParams(BaseModel):
root_cause: str
affected_service: str
recommendation: str
# ── Action models ───────────────────────────────────────────────────
class QueryLogsAction(Action):
tool_name: Literal["query_logs"] = "query_logs"
parameters: QueryLogsParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
class QueryMetricsAction(Action):
tool_name: Literal["query_metrics"] = "query_metrics"
parameters: QueryMetricsParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
class GetServiceStatusAction(Action):
tool_name: Literal["get_service_status"] = "get_service_status"
parameters: GetServiceStatusParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
class GetDependencyMapAction(Action):
tool_name: Literal["get_dependency_map"] = "get_dependency_map"
parameters: GetDependencyMapParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
class ConsultRunbookAction(Action):
tool_name: Literal["consult_runbook"] = "consult_runbook"
parameters: ConsultRunbookParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
class CheckRecentChangesAction(Action):
tool_name: Literal["check_recent_changes"] = "check_recent_changes"
parameters: CheckRecentChangesParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
class SubmitResolutionAction(Action):
tool_name: Literal["submit_resolution"] = "submit_resolution"
parameters: SubmitResolutionParams
def param_dict(self) -> dict:
return self.parameters.model_dump()
# ── Discriminated union ─────────────────────────────────────────────
_ActionUnion = Annotated[
Union[
QueryLogsAction,
QueryMetricsAction,
GetServiceStatusAction,
GetDependencyMapAction,
ConsultRunbookAction,
CheckRecentChangesAction,
SubmitResolutionAction,
],
Field(discriminator="tool_name"),
]
class SentinelAction(RootModel[_ActionUnion]):
"""Discriminated union action — delegates to the matched concrete action."""
@property
def tool_name(self) -> str:
return self.root.tool_name
def param_dict(self) -> dict:
return self.root.param_dict()
# ── Observation & State ─────────────────────────────────────────────
class SentinelObservation(Observation):
"""What the agent sees after each step."""
incident_summary: str = Field(default="", description="Initial alert / ongoing context")
tool_output: str = Field(default="", description="Result from the last tool call")
available_tools: List[str] = Field(default_factory=list, description="Tools the agent can use")
step_number: int = Field(default=0, description="Current step number")
max_steps: int = Field(default=20, description="Maximum steps per episode")
cumulative_reward: float = Field(default=0.0, description="Running total of per-step rewards")
last_action_error: str = Field(default="", description="Error from last invalid action")
tool_descriptions: dict = Field(default_factory=dict, description="Parameter metadata (populated on reset only)")
class SentinelState(State):
"""Internal environment state."""
task_id: int = 1
task_name: str = ""
tools_called: List[str] = Field(default_factory=list)
relevant_tools_called: List[str] = Field(default_factory=list)
resolution_submitted: bool = False
root_cause_correct: bool = False
recommendation_correct: bool = False
final_score: float = 0.0
|