File size: 5,348 Bytes
83136ac | 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 | """OpenEnv client + FastAPI server entry points.
The :class:`ChaosOpsClient` is the thing TRL-style training code imports:
from chaosops.env.openenv_wrapper import ChaosOpsClient
with ChaosOpsClient(base_url="http://localhost:8000").sync() as env:
obs = env.reset()
obs = env.step(action)
The server side is created with :func:`build_fastapi_app`, which returns a
ready-to-serve FastAPI instance. ``server/app.py`` simply re-exports it so
``uvicorn chaosops.env.openenv_wrapper:app`` works out of the box.
If ``openenv-core`` is not installed the client import will raise at the
call site — consistent with the rest of the package: the simulator itself
has no hard dependency on OpenEnv so unit tests stay lightweight.
"""
from __future__ import annotations
from typing import Any
from chaosops.env.environment import ChaosOpsEnvironment
from chaosops.env.models import (
ChaosOpsAction,
ChaosOpsObservation,
ChaosOpsState,
RoleView,
)
try: # pragma: no cover — optional dependency
from openenv.core.env_client import EnvClient
from openenv.core.env_server import create_fastapi_app
from openenv.core.client_types import StepResult
_HAS_OPENENV = True
except ImportError: # pragma: no cover
_HAS_OPENENV = False
EnvClient = object # type: ignore[assignment, misc]
StepResult = None # type: ignore[assignment]
create_fastapi_app = None # type: ignore[assignment]
if _HAS_OPENENV:
class ChaosOpsClient(EnvClient[ChaosOpsAction, ChaosOpsObservation, ChaosOpsState]): # type: ignore[misc]
"""Typed OpenEnv client for ChaosOps AI."""
def _step_payload(self, action: ChaosOpsAction) -> dict[str, Any]:
return action.model_dump(mode="json")
def _parse_result(self, payload: dict[str, Any]) -> "StepResult": # type: ignore[name-defined]
obs_data = payload.get("observation", {})
view_data = obs_data.get("view", {})
view = RoleView.model_validate(view_data) if view_data else None
observation = ChaosOpsObservation(
done=payload.get("done", False),
reward=payload.get("reward"),
view=view, # type: ignore[arg-type]
step=obs_data.get("step", 0),
turn_role=obs_data.get("turn_role"),
message=obs_data.get("message", ""),
)
return StepResult( # type: ignore[operator]
observation=observation,
reward=payload.get("reward"),
done=payload.get("done", False),
)
def _parse_state(self, payload: dict[str, Any]) -> ChaosOpsState:
return ChaosOpsState.model_validate(payload)
def build_fastapi_app(): # type: ignore[no-untyped-def]
"""Return a ready-to-serve FastAPI app exposing ChaosOps AI.
``create_fastapi_app`` expects an environment *factory*, plus the
action and observation Pydantic classes so it can validate incoming
payloads before handing them to the simulator.
"""
assert create_fastapi_app is not None
return create_fastapi_app(
env=ChaosOpsEnvironment,
action_cls=ChaosOpsAction,
observation_cls=ChaosOpsObservation,
)
# Build the app at import time so ``uvicorn chaosops.env.openenv_wrapper:app``
# works. Any error here is a config/install problem — fail loudly rather
# than serve a half-broken app.
try:
app = build_fastapi_app()
except Exception: # pragma: no cover — surfaced via chaosops-serve
app = None
else:
class ChaosOpsClient: # type: ignore[no-redef]
"""Placeholder — install ``openenv-core`` to use the HTTP client."""
def __init__(self, *args: Any, **kwargs: Any) -> None:
raise ImportError(
"openenv-core is not installed. Run `pip install openenv-core` "
"to use the ChaosOps HTTP client."
)
def build_fastapi_app(): # type: ignore[no-untyped-def]
raise ImportError("openenv-core is required to build the FastAPI server app.")
app = None
def serve_cli() -> None:
"""``chaosops-serve`` console entry point.
Boots the FastAPI app with uvicorn. Raises a clear error if the
optional ``openenv`` extras are not installed.
"""
try:
import uvicorn # type: ignore[import-not-found]
except ImportError as exc:
raise SystemExit(
"uvicorn is required for chaosops-serve. "
"Install with: pip install 'chaosops[openenv]'"
) from exc
if not _HAS_OPENENV:
raise SystemExit(
"openenv-core is required for chaosops-serve. "
"Install with: pip install 'chaosops[openenv]'"
)
import argparse
parser = argparse.ArgumentParser(description="ChaosOps AI HTTP server")
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--reload", action="store_true")
args = parser.parse_args()
uvicorn.run(
"chaosops.env.openenv_wrapper:app",
host=args.host,
port=args.port,
reload=args.reload,
)
__all__ = ["ChaosOpsClient", "build_fastapi_app", "app", "serve_cli"]
|