Angshuman28 commited on
Commit
5c4e77e
·
verified ·
1 Parent(s): 0d99aee

Upload folder using huggingface_hub

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. CLAUDE.md +4 -4
  2. README.md +255 -255
  3. __init__.py +16 -16
  4. cortex/brains/__init__.py +26 -0
  5. cortex/brains/_base.py +201 -0
  6. cortex/brains/_executive.py +144 -0
  7. cortex/brains/epidemiology.py +25 -0
  8. cortex/brains/governance.py +22 -0
  9. cortex/brains/logistics.py +23 -0
  10. cortex/lenses.py +186 -0
  11. cortex/schemas.py +46 -1
  12. cortex/subagents/__init__.py +27 -0
  13. cortex/subagents/_base.py +240 -0
  14. cortex/subagents/critic.py +68 -0
  15. cortex/subagents/perception.py +192 -0
  16. cortex/subagents/planner.py +61 -0
  17. cortex/subagents/prompts/critic.txt +19 -0
  18. cortex/subagents/prompts/planner.txt +26 -0
  19. cortex/subagents/prompts/world_modeler.txt +32 -0
  20. cortex/subagents/world_modeler.py +61 -0
  21. demo/CLAUDE.md +45 -45
  22. docs/CORTEX_ARCHITECTURE.md +1 -1
  23. inference.py +9 -4
  24. notebooks/train_b1_grpo.ipynb +496 -0
  25. notebooks/train_cortex_router_grpo.ipynb +567 -0
  26. openenv.yaml +7 -7
  27. openenv_CrisisWorldCortex.egg-info/PKG-INFO +13 -0
  28. openenv_CrisisWorldCortex.egg-info/SOURCES.txt +48 -0
  29. openenv_CrisisWorldCortex.egg-info/dependency_links.txt +1 -0
  30. openenv_CrisisWorldCortex.egg-info/entry_points.txt +2 -0
  31. openenv_CrisisWorldCortex.egg-info/requires.txt +9 -0
  32. openenv_CrisisWorldCortex.egg-info/top_level.txt +1 -0
  33. pyproject.toml +1 -0
  34. server/CLAUDE.md +9 -18
  35. server/CrisisWorldCortex_environment.py +19 -3
  36. server/__init__.py +11 -11
  37. server/app.py +2 -2
  38. server/graders/outer_reward.py +96 -34
  39. server/requirements.txt +6 -6
  40. tests/CLAUDE.md +48 -48
  41. tests/_helpers/__init__.py +6 -0
  42. tests/_helpers/llm_stub.py +82 -0
  43. tests/test_baseline_b1.py +28 -30
  44. tests/test_cortex_brain_executive.py +206 -0
  45. tests/test_cortex_brain_smoke.py +142 -0
  46. tests/test_cortex_lenses.py +238 -0
  47. tests/test_cortex_perception.py +116 -0
  48. tests/test_cortex_subagents.py +387 -0
  49. tests/test_outer_reward_in_range.py +12 -10
  50. tests/test_outer_reward_non_constancy.py +14 -5
CLAUDE.md CHANGED
@@ -28,7 +28,7 @@ Wire-protocol classes are template-locked: `CrisisworldcortexAction`, `Crisiswor
28
 
29
  - The four wire-protocol class names above. Do not rename.
30
  - Flat package layout: `pyproject.toml:package-dir` maps `CrisisWorldCortex` to repo root. Do not move `models.py`, `client.py`, or root `__init__.py` off root.
31
- - Canonical wire-type imports in every `server/` module: `from CrisisWorldCortex.models import ...`.
32
  - `pyproject.toml`: deps, entry points, `packages`, `package-dir`.
33
  - `openenv.yaml` keys: `spec_version`, `runtime`, `app`, `port`.
34
  - `server/app.py:create_app(...)` call signature.
@@ -47,7 +47,7 @@ Wire-protocol classes are template-locked: `CrisisworldcortexAction`, `Crisiswor
47
  ## Import-graph rule (enforced)
48
 
49
  - `cortex/` imports `models` and `cortex/*` only.
50
- - `server/` imports `models`, `openenv.core.*`, and package-relative `server` internals only. No `cortex/*`, no `training/*`, no `baselines/*`, no `demo/*`.
51
  - `baselines/` imports `models`, `client`, `cortex/*`. Must not import `server/*` — baselines hit the env over HTTP like production.
52
  - `training/` imports `models`, `client`, `cortex/*`, `server.graders` (reward-name constants only). Must not import `server.simulator/*`.
53
  - `demo/` imports `cortex.schemas` (types only) and stdlib. Must not import `server/*`, `training/*`, `baselines/*`, `cortex.council`, `cortex.routing_policy`.
@@ -70,9 +70,9 @@ Do not restate subsystem APIs in this file or in other subsystem files.
70
 
71
  - Wire package: `from CrisisWorldCortex import CrisisworldcortexAction, ...`.
72
  - Dev / research directories (sibling-of-repo-root): bare-name — `import cortex`, `import baselines`, `import training`, `import demo`, `import scripts`.
73
- - Server-internal: package-relative imports such as `from .simulator import ...` or `from ..simulator import ...`; do not use `from CrisisWorldCortex.server...`.
74
  - **Cross-boundary into wire types**: when a dev / research directory (`cortex/`, `baselines/`, `training/`, `demo/`) crosses into the wire-protocol package, use `from CrisisWorldCortex.models import …` — **never** bare `from models import …`. Bare-name only applies to dev-directory siblings and to same-subpackage imports inside `server/`. The dual-path import creates distinct `sys.modules` entries and breaks Pydantic discriminator checks across import boundaries (verified by session 4's class-identity bug; `cortex/schemas.py:21` documents the canonicalised import).
75
- - **Cross-boundary into wire types from server modules**: every `server/` file imports wire types with `from CrisisWorldCortex.models import `. Do not use `from ..models` or bare `from models`; both can create a second class identity under different launch modes and break Pydantic discriminated-union validation.
76
 
77
  ## Commands (Git Bash; quote Windows paths)
78
 
 
28
 
29
  - The four wire-protocol class names above. Do not rename.
30
  - Flat package layout: `pyproject.toml:package-dir` maps `CrisisWorldCortex` to repo root. Do not move `models.py`, `client.py`, or root `__init__.py` off root.
31
+ - Dual-import fallback pattern (`try: from ..models / except: from models`) in every `server/` module that imports `models`.
32
  - `pyproject.toml`: deps, entry points, `packages`, `package-dir`.
33
  - `openenv.yaml` keys: `spec_version`, `runtime`, `app`, `port`.
34
  - `server/app.py:create_app(...)` call signature.
 
47
  ## Import-graph rule (enforced)
48
 
49
  - `cortex/` imports `models` and `cortex/*` only.
50
+ - `server/` imports `models`, `openenv.core.*`, `server/*` only. No `cortex/*`, no `training/*`, no `baselines/*`, no `demo/*`.
51
  - `baselines/` imports `models`, `client`, `cortex/*`. Must not import `server/*` — baselines hit the env over HTTP like production.
52
  - `training/` imports `models`, `client`, `cortex/*`, `server.graders` (reward-name constants only). Must not import `server.simulator/*`.
53
  - `demo/` imports `cortex.schemas` (types only) and stdlib. Must not import `server/*`, `training/*`, `baselines/*`, `cortex.council`, `cortex.routing_policy`.
 
70
 
71
  - Wire package: `from CrisisWorldCortex import CrisisworldcortexAction, ...`.
72
  - Dev / research directories (sibling-of-repo-root): bare-name — `import cortex`, `import baselines`, `import training`, `import demo`, `import scripts`.
73
+ - Server-internal: `from server.simulator import ...`, `from server.graders import ...`.
74
  - **Cross-boundary into wire types**: when a dev / research directory (`cortex/`, `baselines/`, `training/`, `demo/`) crosses into the wire-protocol package, use `from CrisisWorldCortex.models import …` — **never** bare `from models import …`. Bare-name only applies to dev-directory siblings and to same-subpackage imports inside `server/`. The dual-path import creates distinct `sys.modules` entries and breaks Pydantic discriminator checks across import boundaries (verified by session 4's class-identity bug; `cortex/schemas.py:21` documents the canonicalised import).
75
+ - **Cross-boundary into wire types from deep server modules**: files inside `server/` more than one level deep (e.g., `server/simulator/seir_model.py`, `server/graders/outer_reward.py`) cannot use the `try: from ..models / except: from models` fallback — `..models` from a two-level-deep module resolves to a non-existent `CrisisWorldCortex.server.models`, the fallback fires, and bare `models` loads as a separate `sys.modules` entry. Use the absolute path: `from CrisisWorldCortex.models import …`. Session 4 (`cortex/schemas.py:21`) and Session 5a (`server/simulator/seir_model.py`, `server/simulator/tasks.py`) document this with inline comments. The dual-import fallback in `server/CrisisWorldCortex_environment.py` and `server/app.py` works only because they are one level deep (`..models` → `CrisisWorldCortex.models` directly).
76
 
77
  ## Commands (Git Bash; quote Windows paths)
78
 
README.md CHANGED
@@ -1,255 +1,255 @@
1
- ---
2
- title: Crisisworldcortex Environment Server
3
- emoji: 🌟
4
- colorFrom: yellow
5
- colorTo: pink
6
- sdk: docker
7
- pinned: false
8
- app_port: 8000
9
- base_path: /web
10
- tags:
11
- - openenv
12
- ---
13
-
14
- # Crisisworldcortex Environment
15
-
16
- A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
17
-
18
- ## Quick Start
19
-
20
- The simplest way to use the Crisisworldcortex environment is through the `CrisisworldcortexEnv` class:
21
-
22
- ```python
23
- from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
24
-
25
- try:
26
- # Create environment from Docker image
27
- CrisisWorldCortexenv = CrisisworldcortexEnv.from_docker_image("CrisisWorldCortex-env:latest")
28
-
29
- # Reset
30
- result = CrisisWorldCortexenv.reset()
31
- print(f"Reset: {result.observation.echoed_message}")
32
-
33
- # Send multiple messages
34
- messages = ["Hello, World!", "Testing echo", "Final message"]
35
-
36
- for msg in messages:
37
- result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message=msg))
38
- print(f"Sent: '{msg}'")
39
- print(f" → Echoed: '{result.observation.echoed_message}'")
40
- print(f" → Length: {result.observation.message_length}")
41
- print(f" → Reward: {result.reward}")
42
-
43
- finally:
44
- # Always clean up
45
- CrisisWorldCortexenv.close()
46
- ```
47
-
48
- That's it! The `CrisisworldcortexEnv.from_docker_image()` method handles:
49
- - Starting the Docker container
50
- - Waiting for the server to be ready
51
- - Connecting to the environment
52
- - Container cleanup when you call `close()`
53
-
54
- ## Building the Docker Image
55
-
56
- Before using the environment, you need to build the Docker image:
57
-
58
- ```bash
59
- # From project root
60
- docker build -t CrisisWorldCortex-env:latest -f server/Dockerfile .
61
- ```
62
-
63
- ## Deploying to Hugging Face Spaces
64
-
65
- You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
66
-
67
- ```bash
68
- # From the environment directory (where openenv.yaml is located)
69
- openenv push
70
-
71
- # Or specify options
72
- openenv push --namespace my-org --private
73
- ```
74
-
75
- The `openenv push` command will:
76
- 1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
77
- 2. Prepare a custom build for Hugging Face Docker space (enables web interface)
78
- 3. Upload to Hugging Face (ensuring you're logged in)
79
-
80
- ### Prerequisites
81
-
82
- - Authenticate with Hugging Face: The command will prompt for login if not already authenticated
83
-
84
- ### Options
85
-
86
- - `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
87
- - `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
88
- - `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
89
- - `--private`: Deploy the space as private (default: public)
90
-
91
- ### Examples
92
-
93
- ```bash
94
- # Push to your personal namespace (defaults to username/env-name from openenv.yaml)
95
- openenv push
96
-
97
- # Push to a specific repository
98
- openenv push --repo-id my-org/my-env
99
-
100
- # Push with a custom base image
101
- openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
102
-
103
- # Push as a private space
104
- openenv push --private
105
-
106
- # Combine options
107
- openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
108
- ```
109
-
110
- After deployment, your space will be available at:
111
- `https://huggingface.co/spaces/<repo-id>`
112
-
113
- The deployed space includes:
114
- - **Web Interface** at `/web` - Interactive UI for exploring the environment
115
- - **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
116
- - **Health Check** at `/health` - Container health monitoring
117
- - **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
118
-
119
- ## Environment Details
120
-
121
- ### Action
122
- **CrisisworldcortexAction**: Contains a single field
123
- - `message` (str) - The message to echo back
124
-
125
- ### Observation
126
- **CrisisworldcortexObservation**: Contains the echo response and metadata
127
- - `echoed_message` (str) - The message echoed back
128
- - `message_length` (int) - Length of the message
129
- - `reward` (float) - Reward based on message length (length × 0.1)
130
- - `done` (bool) - Always False for echo environment
131
- - `metadata` (dict) - Additional info like step count
132
-
133
- ### Reward
134
- The reward is calculated as: `message_length × 0.1`
135
- - "Hi" → reward: 0.2
136
- - "Hello, World!" → reward: 1.3
137
- - Empty message → reward: 0.0
138
-
139
- ## Advanced Usage
140
-
141
- ### Connecting to an Existing Server
142
-
143
- If you already have a Crisisworldcortex environment server running, you can connect directly:
144
-
145
- ```python
146
- from CrisisWorldCortex import CrisisworldcortexEnv
147
-
148
- # Connect to existing server
149
- CrisisWorldCortexenv = CrisisworldcortexEnv(base_url="<ENV_HTTP_URL_HERE>")
150
-
151
- # Use as normal
152
- result = CrisisWorldCortexenv.reset()
153
- result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message="Hello!"))
154
- ```
155
-
156
- Note: When connecting to an existing server, `CrisisWorldCortexenv.close()` will NOT stop the server.
157
-
158
- ### Using the Context Manager
159
-
160
- The client supports context manager usage for automatic connection management:
161
-
162
- ```python
163
- from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
164
-
165
- # Connect with context manager (auto-connects and closes)
166
- with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
167
- result = env.reset()
168
- print(f"Reset: {result.observation.echoed_message}")
169
- # Multiple steps with low latency
170
- for msg in ["Hello", "World", "!"]:
171
- result = env.step(CrisisworldcortexAction(message=msg))
172
- print(f"Echoed: {result.observation.echoed_message}")
173
- ```
174
-
175
- The client uses WebSocket connections for:
176
- - **Lower latency**: No HTTP connection overhead per request
177
- - **Persistent session**: Server maintains your environment state
178
- - **Efficient for episodes**: Better for many sequential steps
179
-
180
- ### Concurrent WebSocket Sessions
181
-
182
- The server supports multiple concurrent WebSocket connections. To enable this,
183
- modify `server/app.py` to use factory mode:
184
-
185
- ```python
186
- # In server/app.py - use factory mode for concurrent sessions
187
- app = create_app(
188
- CrisisworldcortexEnvironment, # Pass class, not instance
189
- CrisisworldcortexAction,
190
- CrisisworldcortexObservation,
191
- max_concurrent_envs=4, # Allow 4 concurrent sessions
192
- )
193
- ```
194
-
195
- Then multiple clients can connect simultaneously:
196
-
197
- ```python
198
- from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
199
- from concurrent.futures import ThreadPoolExecutor
200
-
201
- def run_episode(client_id: int):
202
- with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
203
- result = env.reset()
204
- for i in range(10):
205
- result = env.step(CrisisworldcortexAction(message=f"Client {client_id}, step {i}"))
206
- return client_id, result.observation.message_length
207
-
208
- # Run 4 episodes concurrently
209
- with ThreadPoolExecutor(max_workers=4) as executor:
210
- results = list(executor.map(run_episode, range(4)))
211
- ```
212
-
213
- ## Development & Testing
214
-
215
- ### Direct Environment Testing
216
-
217
- Test the environment logic directly without starting the HTTP server:
218
-
219
- ```bash
220
- # From the server directory
221
- python3 server/CrisisWorldCortex_environment.py
222
- ```
223
-
224
- This verifies that:
225
- - Environment resets correctly
226
- - Step executes actions properly
227
- - State tracking works
228
- - Rewards are calculated correctly
229
-
230
- ### Running Locally
231
-
232
- Run the server locally for development:
233
-
234
- ```bash
235
- uvicorn server.app:app --reload
236
- ```
237
-
238
- ## Project Structure
239
-
240
- ```
241
- CrisisWorldCortex/
242
- ├── .dockerignore # Docker build exclusions
243
- ├── __init__.py # Module exports
244
- ├── README.md # This file
245
- ├── openenv.yaml # OpenEnv manifest
246
- ├── pyproject.toml # Project metadata and dependencies
247
- ├── uv.lock # Locked dependencies (generated)
248
- ├── client.py # CrisisworldcortexEnv client
249
- ├── models.py # Action and Observation models
250
- └── server/
251
- ├── __init__.py # Server module exports
252
- ├── CrisisWorldCortex_environment.py # Core environment logic
253
- ├── app.py # FastAPI application (HTTP + WebSocket endpoints)
254
- └── Dockerfile # Container image definition
255
- ```
 
1
+ ---
2
+ title: Crisisworldcortex Environment Server
3
+ emoji: 🌟
4
+ colorFrom: yellow
5
+ colorTo: pink
6
+ sdk: docker
7
+ pinned: false
8
+ app_port: 8000
9
+ base_path: /web
10
+ tags:
11
+ - openenv
12
+ ---
13
+
14
+ # Crisisworldcortex Environment
15
+
16
+ A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
17
+
18
+ ## Quick Start
19
+
20
+ The simplest way to use the Crisisworldcortex environment is through the `CrisisworldcortexEnv` class:
21
+
22
+ ```python
23
+ from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
24
+
25
+ try:
26
+ # Create environment from Docker image
27
+ CrisisWorldCortexenv = CrisisworldcortexEnv.from_docker_image("CrisisWorldCortex-env:latest")
28
+
29
+ # Reset
30
+ result = CrisisWorldCortexenv.reset()
31
+ print(f"Reset: {result.observation.echoed_message}")
32
+
33
+ # Send multiple messages
34
+ messages = ["Hello, World!", "Testing echo", "Final message"]
35
+
36
+ for msg in messages:
37
+ result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message=msg))
38
+ print(f"Sent: '{msg}'")
39
+ print(f" → Echoed: '{result.observation.echoed_message}'")
40
+ print(f" → Length: {result.observation.message_length}")
41
+ print(f" → Reward: {result.reward}")
42
+
43
+ finally:
44
+ # Always clean up
45
+ CrisisWorldCortexenv.close()
46
+ ```
47
+
48
+ That's it! The `CrisisworldcortexEnv.from_docker_image()` method handles:
49
+ - Starting the Docker container
50
+ - Waiting for the server to be ready
51
+ - Connecting to the environment
52
+ - Container cleanup when you call `close()`
53
+
54
+ ## Building the Docker Image
55
+
56
+ Before using the environment, you need to build the Docker image:
57
+
58
+ ```bash
59
+ # From project root
60
+ docker build -t CrisisWorldCortex-env:latest -f server/Dockerfile .
61
+ ```
62
+
63
+ ## Deploying to Hugging Face Spaces
64
+
65
+ You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
66
+
67
+ ```bash
68
+ # From the environment directory (where openenv.yaml is located)
69
+ openenv push
70
+
71
+ # Or specify options
72
+ openenv push --namespace my-org --private
73
+ ```
74
+
75
+ The `openenv push` command will:
76
+ 1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
77
+ 2. Prepare a custom build for Hugging Face Docker space (enables web interface)
78
+ 3. Upload to Hugging Face (ensuring you're logged in)
79
+
80
+ ### Prerequisites
81
+
82
+ - Authenticate with Hugging Face: The command will prompt for login if not already authenticated
83
+
84
+ ### Options
85
+
86
+ - `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
87
+ - `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
88
+ - `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
89
+ - `--private`: Deploy the space as private (default: public)
90
+
91
+ ### Examples
92
+
93
+ ```bash
94
+ # Push to your personal namespace (defaults to username/env-name from openenv.yaml)
95
+ openenv push
96
+
97
+ # Push to a specific repository
98
+ openenv push --repo-id my-org/my-env
99
+
100
+ # Push with a custom base image
101
+ openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
102
+
103
+ # Push as a private space
104
+ openenv push --private
105
+
106
+ # Combine options
107
+ openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
108
+ ```
109
+
110
+ After deployment, your space will be available at:
111
+ `https://huggingface.co/spaces/<repo-id>`
112
+
113
+ The deployed space includes:
114
+ - **Web Interface** at `/web` - Interactive UI for exploring the environment
115
+ - **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
116
+ - **Health Check** at `/health` - Container health monitoring
117
+ - **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
118
+
119
+ ## Environment Details
120
+
121
+ ### Action
122
+ **CrisisworldcortexAction**: Contains a single field
123
+ - `message` (str) - The message to echo back
124
+
125
+ ### Observation
126
+ **CrisisworldcortexObservation**: Contains the echo response and metadata
127
+ - `echoed_message` (str) - The message echoed back
128
+ - `message_length` (int) - Length of the message
129
+ - `reward` (float) - Reward based on message length (length × 0.1)
130
+ - `done` (bool) - Always False for echo environment
131
+ - `metadata` (dict) - Additional info like step count
132
+
133
+ ### Reward
134
+ The reward is calculated as: `message_length × 0.1`
135
+ - "Hi" → reward: 0.2
136
+ - "Hello, World!" → reward: 1.3
137
+ - Empty message → reward: 0.0
138
+
139
+ ## Advanced Usage
140
+
141
+ ### Connecting to an Existing Server
142
+
143
+ If you already have a Crisisworldcortex environment server running, you can connect directly:
144
+
145
+ ```python
146
+ from CrisisWorldCortex import CrisisworldcortexEnv
147
+
148
+ # Connect to existing server
149
+ CrisisWorldCortexenv = CrisisworldcortexEnv(base_url="<ENV_HTTP_URL_HERE>")
150
+
151
+ # Use as normal
152
+ result = CrisisWorldCortexenv.reset()
153
+ result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message="Hello!"))
154
+ ```
155
+
156
+ Note: When connecting to an existing server, `CrisisWorldCortexenv.close()` will NOT stop the server.
157
+
158
+ ### Using the Context Manager
159
+
160
+ The client supports context manager usage for automatic connection management:
161
+
162
+ ```python
163
+ from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
164
+
165
+ # Connect with context manager (auto-connects and closes)
166
+ with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
167
+ result = env.reset()
168
+ print(f"Reset: {result.observation.echoed_message}")
169
+ # Multiple steps with low latency
170
+ for msg in ["Hello", "World", "!"]:
171
+ result = env.step(CrisisworldcortexAction(message=msg))
172
+ print(f"Echoed: {result.observation.echoed_message}")
173
+ ```
174
+
175
+ The client uses WebSocket connections for:
176
+ - **Lower latency**: No HTTP connection overhead per request
177
+ - **Persistent session**: Server maintains your environment state
178
+ - **Efficient for episodes**: Better for many sequential steps
179
+
180
+ ### Concurrent WebSocket Sessions
181
+
182
+ The server supports multiple concurrent WebSocket connections. To enable this,
183
+ modify `server/app.py` to use factory mode:
184
+
185
+ ```python
186
+ # In server/app.py - use factory mode for concurrent sessions
187
+ app = create_app(
188
+ CrisisworldcortexEnvironment, # Pass class, not instance
189
+ CrisisworldcortexAction,
190
+ CrisisworldcortexObservation,
191
+ max_concurrent_envs=4, # Allow 4 concurrent sessions
192
+ )
193
+ ```
194
+
195
+ Then multiple clients can connect simultaneously:
196
+
197
+ ```python
198
+ from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
199
+ from concurrent.futures import ThreadPoolExecutor
200
+
201
+ def run_episode(client_id: int):
202
+ with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
203
+ result = env.reset()
204
+ for i in range(10):
205
+ result = env.step(CrisisworldcortexAction(message=f"Client {client_id}, step {i}"))
206
+ return client_id, result.observation.message_length
207
+
208
+ # Run 4 episodes concurrently
209
+ with ThreadPoolExecutor(max_workers=4) as executor:
210
+ results = list(executor.map(run_episode, range(4)))
211
+ ```
212
+
213
+ ## Development & Testing
214
+
215
+ ### Direct Environment Testing
216
+
217
+ Test the environment logic directly without starting the HTTP server:
218
+
219
+ ```bash
220
+ # From the server directory
221
+ python3 server/CrisisWorldCortex_environment.py
222
+ ```
223
+
224
+ This verifies that:
225
+ - Environment resets correctly
226
+ - Step executes actions properly
227
+ - State tracking works
228
+ - Rewards are calculated correctly
229
+
230
+ ### Running Locally
231
+
232
+ Run the server locally for development:
233
+
234
+ ```bash
235
+ uvicorn server.app:app --reload
236
+ ```
237
+
238
+ ## Project Structure
239
+
240
+ ```
241
+ CrisisWorldCortex/
242
+ ├── .dockerignore # Docker build exclusions
243
+ ├── __init__.py # Module exports
244
+ ├── README.md # This file
245
+ ├── openenv.yaml # OpenEnv manifest
246
+ ├── pyproject.toml # Project metadata and dependencies
247
+ ├── uv.lock # Locked dependencies (generated)
248
+ ├── client.py # CrisisworldcortexEnv client
249
+ ├── models.py # Action and Observation models
250
+ └── server/
251
+ ├── __init__.py # Server module exports
252
+ ├── CrisisWorldCortex_environment.py # Core environment logic
253
+ ├── app.py # FastAPI application (HTTP + WebSocket endpoints)
254
+ └── Dockerfile # Container image definition
255
+ ```
__init__.py CHANGED
@@ -1,16 +1,16 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
- #
4
- # This source code is licensed under the BSD-style license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- """Crisisworldcortex Environment."""
8
-
9
- from .client import CrisisworldcortexEnv
10
- from .models import CrisisworldcortexAction, CrisisworldcortexObservation
11
-
12
- __all__ = [
13
- "CrisisworldcortexAction",
14
- "CrisisworldcortexObservation",
15
- "CrisisworldcortexEnv",
16
- ]
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Crisisworldcortex Environment."""
8
+
9
+ from .client import CrisisworldcortexEnv
10
+ from .models import CrisisworldcortexAction, CrisisworldcortexObservation
11
+
12
+ __all__ = [
13
+ "CrisisworldcortexAction",
14
+ "CrisisworldcortexObservation",
15
+ "CrisisworldcortexEnv",
16
+ ]
cortex/brains/__init__.py CHANGED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Cortex brains (Session 11+).
2
+
3
+ Public surface:
4
+ - Brain: per-brain class wiring Perception + Lens + 3 LLM subagents + Brain Executive.
5
+ - EpiBrain, LogisticsBrain, GovernanceBrain: factory functions.
6
+ - aggregate_brain_outputs: Brain Executive aggregation function.
7
+
8
+ Each Brain holds its own LLMClient instance; the orchestration layer
9
+ (Council Executive in Session 12, Workstream B trainers) constructs one
10
+ Brain per brain id, optionally with different LLMClients pointing to
11
+ different models. NO module-level state.
12
+ """
13
+
14
+ from ._base import Brain
15
+ from ._executive import aggregate_brain_outputs
16
+ from .epidemiology import EpiBrain
17
+ from .governance import GovernanceBrain
18
+ from .logistics import LogisticsBrain
19
+
20
+ __all__ = [
21
+ "Brain",
22
+ "EpiBrain",
23
+ "GovernanceBrain",
24
+ "LogisticsBrain",
25
+ "aggregate_brain_outputs",
26
+ ]
cortex/brains/_base.py ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Brain class - composes Perception + Lens + 3 Subagents + Brain Executive.
2
+
3
+ Per cortex/CLAUDE.md: each brain has a deterministic Python Perception
4
+ + Lens, three LLM subagents (router-callable), and a deterministic
5
+ Python Brain Executive. The Brain class wires these together.
6
+
7
+ Multi-model deployment: each Brain holds a SINGLE LLMClient instance
8
+ passed at construction. Different brains can use different models by
9
+ constructing each Brain with a different LLMClient (e.g., Qwen for epi,
10
+ Llama for logistics). NO module-level state, NO shared singletons.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ from typing import List, Literal
16
+
17
+ from cortex.lenses import lens_for
18
+ from cortex.schemas import (
19
+ BeliefState,
20
+ BrainLensedObservation,
21
+ BrainRecommendation,
22
+ CandidatePlan,
23
+ CriticReport,
24
+ PerceptionReport,
25
+ SubagentInput,
26
+ )
27
+ from cortex.subagents import (
28
+ CriticSubagent,
29
+ PlannerSubagent,
30
+ WorldModelerSubagent,
31
+ perception_for,
32
+ )
33
+ from cortex.subagents._base import _LLMClientLike
34
+ from CrisisWorldCortex.models import CrisisworldcortexObservation
35
+
36
+ from ._executive import aggregate_brain_outputs
37
+
38
+ _BrainId = Literal["epidemiology", "logistics", "governance"]
39
+
40
+
41
+ class Brain:
42
+ """Per-brain pipeline holder.
43
+
44
+ Each Brain instance owns its own LLMClient. The orchestration layer
45
+ (Session 12 Council, Workstream B trainers) instantiates one Brain
46
+ per brain id, optionally with different LLMClients pointing to
47
+ different models. The Brain class itself has NO module-level state
48
+ and NO forced singleton.
49
+
50
+ Args:
51
+ brain_id: One of "epidemiology", "logistics", "governance".
52
+ llm_client: This brain's LLM client. Subagents are constructed
53
+ with the SAME client so token billing aggregates correctly.
54
+ wm: WorldModeler subagent.
55
+ planner: Planner subagent.
56
+ critic: Critic subagent.
57
+ """
58
+
59
+ def __init__(
60
+ self,
61
+ brain_id: _BrainId,
62
+ llm_client: _LLMClientLike,
63
+ wm: WorldModelerSubagent,
64
+ planner: PlannerSubagent,
65
+ critic: CriticSubagent,
66
+ ) -> None:
67
+ self.brain_id = brain_id
68
+ self.llm_client = llm_client
69
+ self.wm = wm
70
+ self.planner = planner
71
+ self.critic = critic
72
+
73
+ # ------------------------------------------------------------------
74
+ # Deterministic Python pieces (no LLM)
75
+ # ------------------------------------------------------------------
76
+
77
+ def compute_perception(self, obs: CrisisworldcortexObservation) -> PerceptionReport:
78
+ """Run this brain's Perception. Pure Python; no LLM."""
79
+ return perception_for(self.brain_id, obs)
80
+
81
+ def compute_lens(
82
+ self, obs: CrisisworldcortexObservation, last_reward: float
83
+ ) -> BrainLensedObservation:
84
+ """Run this brain's Lens. Pure Python; no LLM."""
85
+ return lens_for(self.brain_id, obs, last_reward)
86
+
87
+ def aggregate(
88
+ self,
89
+ perception: PerceptionReport,
90
+ beliefs: List[BeliefState],
91
+ plans: List[CandidatePlan],
92
+ critics: List[CriticReport],
93
+ tokens_used: int = 0,
94
+ ) -> BrainRecommendation:
95
+ """Run this brain's Brain Executive. Pure Python; no LLM."""
96
+ return aggregate_brain_outputs(
97
+ brain_id=self.brain_id,
98
+ perception=perception,
99
+ beliefs=beliefs,
100
+ plans=plans,
101
+ critics=critics,
102
+ tokens_used=tokens_used,
103
+ )
104
+
105
+ # ------------------------------------------------------------------
106
+ # High-level convenience: round-1 single tick
107
+ # ------------------------------------------------------------------
108
+
109
+ def run_tick(
110
+ self,
111
+ obs: CrisisworldcortexObservation,
112
+ last_reward: float,
113
+ tick: int,
114
+ round_: int = 1,
115
+ ) -> BrainRecommendation:
116
+ """Round-1 single-tick pipeline (Session 11 smoke).
117
+
118
+ Round 2 is orchestrated by the Council Executive (Session 12)
119
+ via the fine-grained methods (compute_perception, compute_lens,
120
+ wm.run / planner.run / critic.run, aggregate). Calling this
121
+ convenience method with ``round_!=1`` raises NotImplementedError
122
+ to prevent accidental misuse before the Council exists.
123
+ """
124
+ if round_ != 1:
125
+ raise NotImplementedError(
126
+ f"Round {round_} orchestration is the Council Executive's "
127
+ f"responsibility (Session 12). Use Brain.compute_perception/"
128
+ f"compute_lens + WorldModelerSubagent.run/PlannerSubagent.run/"
129
+ f"CriticSubagent.run + Brain.aggregate directly."
130
+ )
131
+
132
+ perception = self.compute_perception(obs)
133
+ # Lens is computed for completeness; Session 11 doesn't yet plumb
134
+ # it into SubagentInput (M-FR-4 step indices fixed). Session 12
135
+ # Council will extend the SubagentInput contract to carry lens
136
+ # output if subagents need it.
137
+ _ = self.compute_lens(obs, last_reward)
138
+
139
+ # WorldModeler (step_idx=0)
140
+ wm_input = SubagentInput(
141
+ brain=self.brain_id,
142
+ role="world_modeler",
143
+ tick=tick,
144
+ round=round_,
145
+ perception=perception,
146
+ prior_belief=None,
147
+ prior_plans=[],
148
+ target_plan_id=None,
149
+ last_reward=last_reward,
150
+ recent_action_log_excerpt=list(obs.recent_action_log),
151
+ )
152
+ belief = self.wm.run(wm_input, step_idx=0)
153
+
154
+ # Planner (step_idx=1)
155
+ planner_input = SubagentInput(
156
+ brain=self.brain_id,
157
+ role="planner",
158
+ tick=tick,
159
+ round=round_,
160
+ perception=perception,
161
+ prior_belief=belief,
162
+ prior_plans=[],
163
+ target_plan_id=None,
164
+ last_reward=last_reward,
165
+ recent_action_log_excerpt=list(obs.recent_action_log),
166
+ )
167
+ plan = self.planner.run(planner_input, step_idx=1)
168
+
169
+ # Critic (step_idx=2)
170
+ critic_input = SubagentInput(
171
+ brain=self.brain_id,
172
+ role="critic",
173
+ tick=tick,
174
+ round=round_,
175
+ perception=perception,
176
+ prior_belief=belief,
177
+ prior_plans=[plan],
178
+ target_plan_id="plan-0",
179
+ last_reward=last_reward,
180
+ recent_action_log_excerpt=list(obs.recent_action_log),
181
+ )
182
+ critic = self.critic.run(critic_input, step_idx=2)
183
+
184
+ # Tally tokens billed to this brain's caller_ids.
185
+ caller_id_base = f"cortex:{self.brain_id}"
186
+ tokens_used = sum(
187
+ self.llm_client.tokens_used_for(f"{caller_id_base}:{role}:t{tick}:r{round_}:s{idx}")
188
+ for role, idx in (
189
+ ("world_modeler", 0),
190
+ ("planner", 1),
191
+ ("critic", 2),
192
+ )
193
+ )
194
+
195
+ return self.aggregate(
196
+ perception=perception,
197
+ beliefs=[belief],
198
+ plans=[plan],
199
+ critics=[critic],
200
+ tokens_used=tokens_used,
201
+ )
cortex/brains/_executive.py ADDED
@@ -0,0 +1,144 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Brain Executive - deterministic Python aggregation.
2
+
3
+ Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 15-21 + M-FR-3 partial
4
+ evidence union (perception + beliefs only; CandidatePlan and CriticReport
5
+ schemas have no evidence fields).
6
+
7
+ Brain Executive runs ONCE per brain at round end. NOT router-callable
8
+ per cortex/CLAUDE.md.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from typing import List
14
+
15
+ from cortex.schemas import (
16
+ BeliefState,
17
+ BrainRecommendation,
18
+ CandidatePlan,
19
+ CriticReport,
20
+ EvidenceCitation,
21
+ PerceptionReport,
22
+ )
23
+ from CrisisWorldCortex.models import NoOp
24
+
25
+ _REASONING_SUMMARY_MAX_CHARS = 400 # matches BrainRecommendation.reasoning_summary cap
26
+ _FALSIFIERS_TO_JOIN = 3
27
+ _FALSIFIER_FALLBACK = "(no falsifier provided)"
28
+ _EMPTY_REASONING = "(empty: no subagent produced a parseable plan)"
29
+
30
+
31
+ def aggregate_brain_outputs(
32
+ brain_id: str,
33
+ perception: PerceptionReport,
34
+ beliefs: List[BeliefState],
35
+ plans: List[CandidatePlan],
36
+ critics: List[CriticReport],
37
+ tokens_used: int = 0,
38
+ ) -> BrainRecommendation:
39
+ """Aggregate one brain's per-round subagent outputs into a recommendation.
40
+
41
+ Decisions:
42
+ D15: argmax over expected_value * confidence.
43
+ D16: top_confidence = chosen.confidence * (1 - chosen_belief.uncertainty).
44
+ D17: minority_actions = all expected_outer_actions except chosen.
45
+ D19: reasoning_summary = chosen.action_sketch[:400].
46
+ D20 (M-FR-3): evidence = perception.evidence + flat-union of beliefs[*].evidence.
47
+ CandidatePlan/CriticReport carry no evidence fields.
48
+ D21: brain_id is lowercase per Pydantic Literal in BrainRecommendation.
49
+
50
+ Empty fallback (M-FR-7): no plans, or chosen plan has confidence==0
51
+ -> top_action=NoOp, top_confidence=0, uncertainty=1.0,
52
+ reasoning_summary=_EMPTY_REASONING.
53
+
54
+ Args:
55
+ brain_id: lowercase brain id ("epidemiology" / "logistics" / "governance").
56
+ perception: This brain's PerceptionReport for the tick.
57
+ beliefs: Per-round BeliefStates. Index aligned with ``plans``.
58
+ plans: Per-round CandidatePlans.
59
+ critics: Per-round CriticReports (currently unused in aggregation but
60
+ kept on the signature so the trajectory log captures the full
61
+ chain).
62
+ tokens_used: Total tokens billed across this brain's subagents.
63
+ """
64
+ if not plans:
65
+ return _empty_recommendation(brain_id, perception, beliefs, tokens_used)
66
+
67
+ # D15: argmax over expected_value * confidence
68
+ chosen_idx = max(
69
+ range(len(plans)),
70
+ key=lambda i: plans[i].expected_value * plans[i].confidence,
71
+ )
72
+ chosen_plan = plans[chosen_idx]
73
+
74
+ if chosen_plan.confidence == 0.0:
75
+ # All plans are empty fallbacks (or the only plan is empty).
76
+ # Brain Executive treats this as no-signal.
77
+ return _empty_recommendation(brain_id, perception, beliefs, tokens_used)
78
+
79
+ # D16: top_confidence = chosen.confidence * (1 - belief.uncertainty)
80
+ if chosen_idx < len(beliefs):
81
+ chosen_belief = beliefs[chosen_idx]
82
+ uncertainty = chosen_belief.uncertainty
83
+ else:
84
+ # Defensive: parallel arrays should match. If not, treat as max uncertainty.
85
+ uncertainty = 1.0
86
+ top_confidence = chosen_plan.confidence * (1.0 - uncertainty)
87
+
88
+ # D17: minority_actions = all plans except chosen
89
+ minority_actions = [
90
+ plans[i].expected_outer_action for i in range(len(plans)) if i != chosen_idx
91
+ ]
92
+
93
+ # D19: reasoning_summary
94
+ reasoning_summary = chosen_plan.action_sketch[:_REASONING_SUMMARY_MAX_CHARS]
95
+
96
+ # D20 (M-FR-3): evidence union from perception + beliefs only.
97
+ # CandidatePlan and CriticReport schemas (Session 9) have no evidence fields;
98
+ # the perception+beliefs union captures the actionable evidence chain since
99
+ # plans/critics derive from beliefs.
100
+ evidence: List[EvidenceCitation] = list(perception.evidence)
101
+ for b in beliefs:
102
+ evidence.extend(b.evidence)
103
+
104
+ # falsifier (M-FR-6): join up to 3 falsifiers; fallback if empty.
105
+ if chosen_plan.falsifiers:
106
+ falsifier = "; ".join(chosen_plan.falsifiers[:_FALSIFIERS_TO_JOIN])
107
+ else:
108
+ falsifier = _FALSIFIER_FALLBACK
109
+
110
+ return BrainRecommendation(
111
+ brain=brain_id,
112
+ top_action=chosen_plan.expected_outer_action,
113
+ top_confidence=top_confidence,
114
+ minority_actions=minority_actions,
115
+ reasoning_summary=reasoning_summary,
116
+ evidence=evidence,
117
+ falsifier=falsifier,
118
+ uncertainty=uncertainty,
119
+ tokens_used=tokens_used,
120
+ )
121
+
122
+
123
+ def _empty_recommendation(
124
+ brain_id: str,
125
+ perception: PerceptionReport,
126
+ beliefs: List[BeliefState],
127
+ tokens_used: int,
128
+ ) -> BrainRecommendation:
129
+ """M-FR-7 empty fallback: NoOp + confidence=0 + uncertainty=1."""
130
+ evidence: List[EvidenceCitation] = list(perception.evidence)
131
+ for b in beliefs:
132
+ evidence.extend(b.evidence)
133
+
134
+ return BrainRecommendation(
135
+ brain=brain_id,
136
+ top_action=NoOp(),
137
+ top_confidence=0.0,
138
+ minority_actions=[],
139
+ reasoning_summary=_EMPTY_REASONING,
140
+ evidence=evidence,
141
+ falsifier=_FALSIFIER_FALLBACK,
142
+ uncertainty=1.0,
143
+ tokens_used=tokens_used,
144
+ )
cortex/brains/epidemiology.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Epidemiology brain factory."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from cortex.subagents import CriticSubagent, PlannerSubagent, WorldModelerSubagent
6
+ from cortex.subagents._base import _LLMClientLike
7
+
8
+ from ._base import Brain
9
+
10
+
11
+ def EpiBrain(llm_client: _LLMClientLike) -> Brain:
12
+ """Construct an Epidemiology Brain bound to ``llm_client``.
13
+
14
+ Multi-model deployment: pass a different ``llm_client`` per brain
15
+ instance to use different models per brain (e.g., Qwen for epi,
16
+ Llama for logistics). The 3 LLM subagents are constructed with the
17
+ SAME client so token billing aggregates correctly.
18
+ """
19
+ return Brain(
20
+ brain_id="epidemiology",
21
+ llm_client=llm_client,
22
+ wm=WorldModelerSubagent(llm_client),
23
+ planner=PlannerSubagent(llm_client),
24
+ critic=CriticSubagent(llm_client),
25
+ )
cortex/brains/governance.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Governance brain factory."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from cortex.subagents import CriticSubagent, PlannerSubagent, WorldModelerSubagent
6
+ from cortex.subagents._base import _LLMClientLike
7
+
8
+ from ._base import Brain
9
+
10
+
11
+ def GovernanceBrain(llm_client: _LLMClientLike) -> Brain:
12
+ """Construct a Governance Brain bound to ``llm_client``.
13
+
14
+ Multi-model deployment: see EpiBrain.
15
+ """
16
+ return Brain(
17
+ brain_id="governance",
18
+ llm_client=llm_client,
19
+ wm=WorldModelerSubagent(llm_client),
20
+ planner=PlannerSubagent(llm_client),
21
+ critic=CriticSubagent(llm_client),
22
+ )
cortex/brains/logistics.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Logistics brain factory."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from cortex.subagents import CriticSubagent, PlannerSubagent, WorldModelerSubagent
6
+ from cortex.subagents._base import _LLMClientLike
7
+
8
+ from ._base import Brain
9
+
10
+
11
+ def LogisticsBrain(llm_client: _LLMClientLike) -> Brain:
12
+ """Construct a Logistics Brain bound to ``llm_client``.
13
+
14
+ Multi-model deployment: see EpiBrain. Pass a Llama-bound client
15
+ here while the other brains use Qwen, etc.
16
+ """
17
+ return Brain(
18
+ brain_id="logistics",
19
+ llm_client=llm_client,
20
+ wm=WorldModelerSubagent(llm_client),
21
+ planner=PlannerSubagent(llm_client),
22
+ critic=CriticSubagent(llm_client),
23
+ )
cortex/lenses.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Brain-specific observation lenses (Session 10).
2
+
3
+ Phase A docs/CORTEX_ARCHITECTURE.md Decisions 9-14 + §2 A1.
4
+
5
+ Lenses are pure-Python: no LLM, no I/O, no state. ``lens_for(brain, obs,
6
+ last_reward)`` dispatches to one of 3 brain-specific helpers and returns
7
+ a ``BrainLensedObservation``. V2 brain ids raise ``KeyError`` per
8
+ Decision 9 (post-review) -- no MVP stub functions.
9
+
10
+ The lens does NOT strip the raw observation (Decision 13); subagents
11
+ may need fields the lens didn't emphasise. ``salient_field_ids`` is a
12
+ salience map alongside ``raw_obs``, not a replacement for it.
13
+
14
+ ``transmission_rate_trend`` is fixed at 0.0 in MVP (M-FR-2): the lens
15
+ sees one observation per call. Session 11 plumbs prior-tick obs into
16
+ the lens to enable real trend computation.
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ from typing import Callable, Dict
22
+
23
+ from cortex.schemas import BrainLensedObservation
24
+ from CrisisWorldCortex.models import CrisisworldcortexObservation
25
+
26
+ _V2_BRAINS = frozenset({"communications", "equity"})
27
+
28
+
29
+ def lens_for(
30
+ brain: str,
31
+ obs: CrisisworldcortexObservation,
32
+ last_reward: float,
33
+ ) -> BrainLensedObservation:
34
+ """Return the per-brain lensed observation.
35
+
36
+ Args:
37
+ brain: One of {"epidemiology", "logistics", "governance"}.
38
+ obs: The current tick's observation.
39
+ last_reward: Previous tick's reward (plumbed from B1's pattern,
40
+ included on the lensed object so all 3 subagents in this
41
+ brain see the same recency signal).
42
+
43
+ Raises:
44
+ KeyError: If ``brain`` is a V2-deferred brain (communications,
45
+ equity) or unknown.
46
+ """
47
+ helper = _LENS_REGISTRY.get(brain)
48
+ if helper is not None:
49
+ return helper(obs, last_reward)
50
+ if brain in _V2_BRAINS:
51
+ raise KeyError(
52
+ f"V2 brain {brain!r} deferred per Phase A Decision 9; "
53
+ f"no MVP stub lens. See docs/CORTEX_ARCHITECTURE.md."
54
+ )
55
+ raise KeyError(f"unknown brain: {brain!r}")
56
+
57
+
58
+ # ============================================================================
59
+ # Per-brain lens helpers
60
+ # ============================================================================
61
+
62
+
63
+ def _epi_lens(obs: CrisisworldcortexObservation, last_reward: float) -> BrainLensedObservation:
64
+ """Epidemiology lens (Decision 10, M-FR-4 rename: epi_pressure)."""
65
+ n_regions = max(1, len(obs.regions))
66
+ mean_hospital_load = sum(r.hospital_load for r in obs.regions) / n_regions
67
+ # M-FR-4: pressure scalar correlated with R_eff but not a true R_eff
68
+ # estimate. WorldModeler subagent computes proper R_eff during reasoning.
69
+ epi_pressure = max(0.0, min(3.0, mean_hospital_load * 2.0))
70
+
71
+ max_cases = max((r.reported_cases_d_ago for r in obs.regions), default=0)
72
+ # /1000 normaliser matches the design-doc "~30 cases / 1000 pop" spec
73
+ worst_region_infection = max(0.0, min(1.0, max_cases / 1000.0))
74
+
75
+ return BrainLensedObservation(
76
+ brain="epidemiology",
77
+ raw_obs=obs,
78
+ salient_field_ids=[
79
+ "regions[*].reported_cases_d_ago",
80
+ "regions[*].hospital_load",
81
+ "regions[*].compliance_proxy",
82
+ ],
83
+ derived_features={
84
+ "epi_pressure": float(epi_pressure),
85
+ "worst_region_infection": float(worst_region_infection),
86
+ # M-FR-2: needs history; Session 11 plumbs prior-tick obs.
87
+ "transmission_rate_trend": 0.0,
88
+ },
89
+ last_reward=last_reward,
90
+ )
91
+
92
+
93
+ def _logistics_lens(
94
+ obs: CrisisworldcortexObservation, last_reward: float
95
+ ) -> BrainLensedObservation:
96
+ """Logistics lens (Decision 11, M-FR-3 floor 0.5, M-FR-6 flat keys)."""
97
+ res = obs.resources
98
+ total_inventory = float(
99
+ res.test_kits + res.hospital_beds_free + res.mobile_units + res.vaccine_doses
100
+ )
101
+
102
+ hospital_load_max = (
103
+ max((r.hospital_load for r in obs.regions), default=0.0) if obs.regions else 0.0
104
+ )
105
+
106
+ # Per-region feasibility flat keys (D14 + M-FR-6).
107
+ strict_regions = {r.region for r in obs.active_restrictions if r.severity == "strict"}
108
+ feasibility: Dict[str, float] = {}
109
+ for r in obs.regions:
110
+ key = f"deployment_feasibility_{r.region}"
111
+ if total_inventory <= 0.0:
112
+ feasibility[key] = 0.0
113
+ elif r.region in strict_regions:
114
+ # M-FR-3: 0.5 floor when strict restriction is in place but
115
+ # units could still be helicoptered in; Planner does the
116
+ # legal-check.
117
+ feasibility[key] = 0.5
118
+ else:
119
+ feasibility[key] = 1.0
120
+
121
+ derived_features: Dict[str, float] = {
122
+ "total_inventory": total_inventory,
123
+ "hospital_load_max": float(hospital_load_max),
124
+ **feasibility,
125
+ }
126
+
127
+ return BrainLensedObservation(
128
+ brain="logistics",
129
+ raw_obs=obs,
130
+ salient_field_ids=[
131
+ "resources.test_kits",
132
+ "resources.hospital_beds_free",
133
+ "resources.mobile_units",
134
+ "resources.vaccine_doses",
135
+ "regions[*].hospital_load",
136
+ "active_restrictions[*]",
137
+ ],
138
+ derived_features=derived_features,
139
+ last_reward=last_reward,
140
+ )
141
+
142
+
143
+ def _governance_lens(
144
+ obs: CrisisworldcortexObservation, last_reward: float
145
+ ) -> BrainLensedObservation:
146
+ """Governance lens (Decision 12)."""
147
+ # escalation_unlocked_strict: any accepted escalate(national) in the log
148
+ escalation_unlocked = 0.0
149
+ for ea in obs.recent_action_log:
150
+ if (
151
+ ea.accepted
152
+ and ea.action.kind == "escalate"
153
+ and getattr(ea.action, "to_authority", None) == "national"
154
+ ):
155
+ escalation_unlocked = 1.0
156
+ break
157
+
158
+ return BrainLensedObservation(
159
+ brain="governance",
160
+ raw_obs=obs,
161
+ salient_field_ids=[
162
+ "active_restrictions[*]",
163
+ "legal_constraints[*]",
164
+ "recent_action_log[*]",
165
+ ],
166
+ derived_features={
167
+ "escalation_unlocked_strict": escalation_unlocked,
168
+ "legal_constraints_count": float(len(obs.legal_constraints)),
169
+ "restrictions_active_count": float(len(obs.active_restrictions)),
170
+ },
171
+ last_reward=last_reward,
172
+ )
173
+
174
+
175
+ # ============================================================================
176
+ # Registry (defined after helpers so closures resolve cleanly)
177
+ # ============================================================================
178
+
179
+
180
+ _LENS_REGISTRY: Dict[
181
+ str, Callable[[CrisisworldcortexObservation, float], BrainLensedObservation]
182
+ ] = {
183
+ "epidemiology": _epi_lens,
184
+ "logistics": _logistics_lens,
185
+ "governance": _governance_lens,
186
+ }
cortex/schemas.py CHANGED
@@ -35,7 +35,12 @@ from pydantic import BaseModel, Field
35
  # OWN internal types (cortex.subagents, cortex.brains, etc.) continue
36
  # to use bare-name sibling imports per Phase 1 C1 — only the cross-package
37
  # wire boundary is canonicalised.
38
- from CrisisWorldCortex.models import OuterActionPayload, RegionId
 
 
 
 
 
39
 
40
  EpistemicPhase = Literal["Divergence", "Challenge", "Narrowing", "Convergence"]
41
 
@@ -137,6 +142,46 @@ SubagentReport = Union[BeliefState, CandidatePlan, CriticReport]
137
  trajectory buffers that need to carry 'any subagent output' generically."""
138
 
139
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
140
  # ============================================================================
141
  # Brain output + Council decision
142
  # ============================================================================
 
35
  # OWN internal types (cortex.subagents, cortex.brains, etc.) continue
36
  # to use bare-name sibling imports per Phase 1 C1 — only the cross-package
37
  # wire boundary is canonicalised.
38
+ from CrisisWorldCortex.models import (
39
+ CrisisworldcortexObservation,
40
+ ExecutedAction,
41
+ OuterActionPayload,
42
+ RegionId,
43
+ )
44
 
45
  EpistemicPhase = Literal["Divergence", "Challenge", "Narrowing", "Convergence"]
46
 
 
142
  trajectory buffers that need to carry 'any subagent output' generically."""
143
 
144
 
145
+ class SubagentInput(BaseModel):
146
+ """Typed input handed to one of the 3 LLM subagents per call.
147
+
148
+ Per Phase A §2 A2: each subagent call receives a fully-typed input
149
+ so prompts are deterministic and testable. ``prior_belief`` is
150
+ ``None`` on round 1 (nothing to revise yet); on round 2 it carries
151
+ the previous round's BeliefState (or an empty BeliefState if round 1
152
+ failed, per Phase A Decision 62). ``prior_plans`` is empty for
153
+ WorldModeler / Planner; populated for Critic so it can attack a
154
+ specific plan. ``target_plan_id`` is required when ``role='critic'``.
155
+ """
156
+
157
+ brain: Literal["epidemiology", "logistics", "governance"]
158
+ role: Literal["world_modeler", "planner", "critic"]
159
+ tick: int = Field(ge=0)
160
+ round: int = Field(ge=1, le=2, description="MVP cap: 1 or 2 only")
161
+ perception: PerceptionReport
162
+ prior_belief: Optional[BeliefState] = None
163
+ prior_plans: List[CandidatePlan] = Field(default_factory=list)
164
+ target_plan_id: Optional[str] = None
165
+ last_reward: float
166
+ recent_action_log_excerpt: List[ExecutedAction] = Field(default_factory=list)
167
+
168
+
169
+ class BrainLensedObservation(BaseModel):
170
+ """Per-brain salience-mapped observation per Phase A §2 A1.
171
+
172
+ Lenses do not strip fields from the raw observation (Decision 13);
173
+ they project a salience map alongside it. ``derived_features`` lets
174
+ each brain pre-compute domain-specific scalars once and pass them
175
+ to all three of its LLM subagents without re-reading ``raw_obs``.
176
+ """
177
+
178
+ brain: Literal["epidemiology", "logistics", "governance"]
179
+ raw_obs: CrisisworldcortexObservation
180
+ salient_field_ids: List[str] = Field(default_factory=list)
181
+ derived_features: Dict[str, float] = Field(default_factory=dict)
182
+ last_reward: float
183
+
184
+
185
  # ============================================================================
186
  # Brain output + Council decision
187
  # ============================================================================
cortex/subagents/__init__.py CHANGED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Cortex per-brain subagents (Session 9+).
2
+
3
+ Public surface:
4
+ - WorldModelerSubagent: emits BeliefState (LLM, router-callable).
5
+ - PlannerSubagent: emits CandidatePlan (LLM, router-callable).
6
+ - CriticSubagent: emits CriticReport (LLM, router-callable).
7
+ - perception_for: deterministic Python Perception function (Session 11+;
8
+ NOT router-callable per cortex/CLAUDE.md role-split binding).
9
+ - PROMPTS_DIR: directory holding the per-role SYS prompt templates.
10
+
11
+ Brain Executive (Python-only, NOT router-callable) lives in
12
+ ``cortex/brains/_executive.py``.
13
+ """
14
+
15
+ from ._base import PROMPTS_DIR
16
+ from .critic import CriticSubagent
17
+ from .perception import perception_for
18
+ from .planner import PlannerSubagent
19
+ from .world_modeler import WorldModelerSubagent
20
+
21
+ __all__ = [
22
+ "CriticSubagent",
23
+ "PROMPTS_DIR",
24
+ "PlannerSubagent",
25
+ "WorldModelerSubagent",
26
+ "perception_for",
27
+ ]
cortex/subagents/_base.py ADDED
@@ -0,0 +1,240 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Abstract base for the 3 LLM subagents (WorldModeler, Planner, Critic).
2
+
3
+ Phase A docs/CORTEX_ARCHITECTURE.md Decisions 1-8 + 62 lock the role split,
4
+ prompt-loading mechanism, retry-with-history semantics, empty fallback
5
+ shape, caller-id format, and TypeAdapter validation pattern. This base
6
+ class implements the shared mechanics; concrete subclasses pin the
7
+ role name, output type, prompt path, TypeAdapter, and USR builder.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from abc import ABC, abstractmethod
13
+ from pathlib import Path
14
+ from typing import ClassVar, List, Optional, Protocol
15
+
16
+ from pydantic import BaseModel, TypeAdapter, ValidationError
17
+
18
+ from cortex.llm_client import ChatMessage, ChatResponse
19
+ from cortex.schemas import SubagentInput
20
+ from CrisisWorldCortex.models import ExecutedAction
21
+
22
+ # ============================================================================
23
+ # Module-level constants (loaded at import time)
24
+ # ============================================================================
25
+
26
+ PROMPTS_DIR: Path = Path(__file__).parent / "prompts"
27
+ """Directory holding the per-role SYS prompt template files (Decision 4)."""
28
+
29
+ _RETRY_SNIPPET_MAX_CHARS: int = 200
30
+ """Cap on the failed-response snippet included in the retry message
31
+ (Decision 8). 200 chars ~= 50 tokens; keeps retry overhead bounded."""
32
+
33
+ _BASE_RETRY_USER_TEMPLATE: str = (
34
+ "Your previous response failed to parse as JSON. The response was:\n"
35
+ "{snippet}\n\n"
36
+ "Emit ONLY valid JSON matching the schema specified in the system prompt. "
37
+ "No prose, no code fences."
38
+ )
39
+
40
+ _RECENT_ACTION_LOG_TAIL: int = 8
41
+ """How many entries from ``recent_action_log_excerpt`` to render into the
42
+ USR summary. Matches the design-doc 8-deep history (M-FR-3)."""
43
+
44
+
45
+ # ============================================================================
46
+ # Duck-typed LLM client protocol (so tests can pass StubLLMClient)
47
+ # ============================================================================
48
+
49
+
50
+ class _LLMClientLike(Protocol):
51
+ """Subset of ``cortex.llm_client.LLMClient`` that subagents call.
52
+
53
+ Production: ``LLMClient``. Tests: ``tests._helpers.llm_stub.StubLLMClient``.
54
+ """
55
+
56
+ def chat(
57
+ self,
58
+ caller_id: str,
59
+ messages: List[ChatMessage],
60
+ max_tokens: Optional[int] = ...,
61
+ temperature: Optional[float] = ...,
62
+ ) -> ChatResponse: ...
63
+
64
+
65
+ # ============================================================================
66
+ # Abstract base
67
+ # ============================================================================
68
+
69
+
70
+ class _LLMSubagent(ABC):
71
+ """Shared run/retry/parse/empty-fallback skeleton for the 3 subagents.
72
+
73
+ Subclasses override the class-level vars below and implement
74
+ ``_build_user_message`` + ``empty_fallback``.
75
+ """
76
+
77
+ # --- Subclass class-level overrides -------------------------------------
78
+ _role_name: ClassVar[str] # one of: "world_modeler", "planner", "critic"
79
+ _output_type: ClassVar[type] # BeliefState / CandidatePlan / CriticReport
80
+ _system_prompt_filename: ClassVar[str] # e.g. "world_modeler.txt"
81
+ _SYSTEM_PROMPT_TEMPLATE: ClassVar[str] # populated by load_prompt() at module load
82
+ _ADAPTER: ClassVar[TypeAdapter] # populated at module load
83
+
84
+ # --- Construction --------------------------------------------------------
85
+
86
+ def __init__(self, llm_client: _LLMClientLike) -> None:
87
+ self._llm = llm_client
88
+
89
+ # --- Public surface ------------------------------------------------------
90
+
91
+ def run(self, input: SubagentInput, step_idx: int) -> BaseModel:
92
+ """Call the LLM (with 1 retry), parse, return typed output or empty fallback.
93
+
94
+ Always returns a typed object - never ``None``. Decision 6: on
95
+ any failure (parse, retry-parse, LLM call exception) returns the
96
+ role-specific empty fallback.
97
+ """
98
+ # Defensive: subclass enforces role-input alignment so harnesses
99
+ # don't accidentally route a Planner input through a Critic class.
100
+ assert input.role == self._role_name, (
101
+ f"SubagentInput.role={input.role!r} does not match "
102
+ f"{type(self).__name__}._role_name={self._role_name!r}"
103
+ )
104
+
105
+ sys_content = self._SYSTEM_PROMPT_TEMPLATE.format(
106
+ brain=input.brain,
107
+ target_plan_id=input.target_plan_id or "",
108
+ )
109
+ usr_content = self._build_user_message(input)
110
+ messages: List[ChatMessage] = [
111
+ ChatMessage(role="system", content=sys_content),
112
+ ChatMessage(role="user", content=usr_content),
113
+ ]
114
+ caller_id = self._caller_id(input, step_idx)
115
+
116
+ # ---- Attempt 1 -----------------------------------------------------
117
+ first_response = self._safe_chat(caller_id, messages)
118
+ if first_response is None:
119
+ return self._empty_fallback_for(input)
120
+ parsed = self._try_parse(first_response.content)
121
+ if parsed is not None:
122
+ return parsed
123
+
124
+ # ---- Attempt 2 (retry with chat-history continuation) --------------
125
+ snippet = self._truncate_snippet(first_response.content)
126
+ retry_messages: List[ChatMessage] = [
127
+ *messages,
128
+ ChatMessage(role="assistant", content=first_response.content),
129
+ ChatMessage(
130
+ role="user",
131
+ content=_BASE_RETRY_USER_TEMPLATE.format(snippet=snippet),
132
+ ),
133
+ ]
134
+ retry_response = self._safe_chat(caller_id, retry_messages)
135
+ if retry_response is None:
136
+ return self._empty_fallback_for(input)
137
+ parsed_retry = self._try_parse(retry_response.content)
138
+ if parsed_retry is not None:
139
+ return parsed_retry
140
+
141
+ # ---- Both attempts failed - empty fallback -------------------------
142
+ return self._empty_fallback_for(input)
143
+
144
+ # --- Subclass extension points ------------------------------------------
145
+
146
+ @abstractmethod
147
+ def _build_user_message(self, input: SubagentInput) -> str:
148
+ """Render the role-specific USR message body."""
149
+
150
+ @classmethod
151
+ @abstractmethod
152
+ def empty_fallback(cls, brain: str, target_plan_id: str = "") -> BaseModel:
153
+ """Return the empty / no-signal output for this role.
154
+
155
+ Phase A Decision 6: confidence/severity = 0 and empty evidence/attacks
156
+ signal "no useful input from this subagent" to the Brain Executive.
157
+ """
158
+
159
+ # --- Internal helpers ---------------------------------------------------
160
+
161
+ def _caller_id(self, input: SubagentInput, step_idx: int) -> str:
162
+ # Phase A Decision 7: cortex:<brain>:<role>:t<tick>:r<round>:s<step_idx>
163
+ return f"cortex:{input.brain}:{self._role_name}:t{input.tick}:r{input.round}:s{step_idx}"
164
+
165
+ def _safe_chat(self, caller_id: str, messages: List[ChatMessage]) -> Optional[ChatResponse]:
166
+ """Call LLM; on exception, return None so caller can empty-fallback."""
167
+ try:
168
+ return self._llm.chat(caller_id=caller_id, messages=messages)
169
+ except Exception:
170
+ # Decision 6: LLM call failure folds into the same empty-fallback path
171
+ # as parse failure. Brain Executive sees a no-signal subagent.
172
+ return None
173
+
174
+ def _try_parse(self, content: str) -> Optional[BaseModel]:
175
+ """Validate ``content`` as JSON via this role's TypeAdapter.
176
+
177
+ Strips common markdown code fences before validating since some
178
+ models wrap JSON in ```json ... ```.
179
+ """
180
+ cleaned = _strip_code_fences(content.strip())
181
+ if not cleaned:
182
+ return None
183
+ try:
184
+ return self._ADAPTER.validate_json(cleaned)
185
+ except (ValidationError, ValueError):
186
+ return None
187
+
188
+ def _empty_fallback_for(self, input: SubagentInput) -> BaseModel:
189
+ return type(self).empty_fallback(
190
+ brain=input.brain,
191
+ target_plan_id=input.target_plan_id or "",
192
+ )
193
+
194
+ @staticmethod
195
+ def _truncate_snippet(content: str) -> str:
196
+ if len(content) <= _RETRY_SNIPPET_MAX_CHARS:
197
+ return content
198
+ return content[:_RETRY_SNIPPET_MAX_CHARS] + "..."
199
+
200
+ @staticmethod
201
+ def _format_action_log(log: List[ExecutedAction]) -> str:
202
+ """M-FR-3 - render recent_action_log_excerpt as a compact text summary.
203
+
204
+ Format: ``"tick 4: deploy_resource accepted; tick 5: restrict_movement.strict rejected"``.
205
+ Capped at the most recent 8 entries.
206
+ """
207
+ if not log:
208
+ return "(empty)"
209
+ items: List[str] = []
210
+ for ea in log[-_RECENT_ACTION_LOG_TAIL:]:
211
+ status = "accepted" if ea.accepted else "rejected"
212
+ kind = ea.action.kind
213
+ extra = ""
214
+ if kind == "restrict_movement":
215
+ extra = f".{getattr(ea.action, 'severity', '?')}"
216
+ items.append(f"tick {ea.tick}: {kind}{extra} {status}")
217
+ return "; ".join(items)
218
+
219
+
220
+ # ============================================================================
221
+ # Helpers (module-level)
222
+ # ============================================================================
223
+
224
+
225
+ def load_prompt(filename: str) -> str:
226
+ """Load a SYS prompt template at module-load time (Decision 4)."""
227
+ return (PROMPTS_DIR / filename).read_text(encoding="utf-8")
228
+
229
+
230
+ def _strip_code_fences(s: str) -> str:
231
+ """Remove leading ``` / ```json fence and trailing ``` if present."""
232
+ s = s.strip()
233
+ if not s.startswith("```"):
234
+ return s
235
+ lines = s.split("\n")
236
+ if lines[0].startswith("```"):
237
+ lines = lines[1:]
238
+ if lines and lines[-1].strip() == "```":
239
+ lines = lines[:-1]
240
+ return "\n".join(lines).strip()
cortex/subagents/critic.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Critic LLM subagent.
2
+
3
+ Phase A docs/CORTEX_ARCHITECTURE.md §9 Decision 3: SYS = critic role;
4
+ USR = perception + target plan + WM belief (M-FR-5). Critic emits prose
5
+ ``CriticReport`` only; never proposes alternative actions.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import ClassVar, List
11
+
12
+ from pydantic import TypeAdapter
13
+
14
+ from cortex.schemas import CriticReport, SubagentInput
15
+
16
+ from ._base import _LLMSubagent, load_prompt
17
+
18
+ _CRITIC_ADAPTER: TypeAdapter[CriticReport] = TypeAdapter(CriticReport)
19
+
20
+
21
+ class CriticSubagent(_LLMSubagent):
22
+ """LLM subagent that emits ``CriticReport`` for one brain per call."""
23
+
24
+ _role_name: ClassVar[str] = "critic"
25
+ _output_type: ClassVar[type] = CriticReport
26
+ _system_prompt_filename: ClassVar[str] = "critic.txt"
27
+ _SYSTEM_PROMPT_TEMPLATE: ClassVar[str] = load_prompt("critic.txt")
28
+ _ADAPTER: ClassVar[TypeAdapter] = _CRITIC_ADAPTER
29
+
30
+ def _build_user_message(self, input: SubagentInput) -> str:
31
+ sections: List[str] = []
32
+ sections.append(f"# Perception\n{input.perception.model_dump_json(indent=2)}")
33
+ # M-FR-5: target plan + WM belief, both as full JSON.
34
+ target_json = self._select_target_plan(input)
35
+ sections.append(f"# Target plan (id={input.target_plan_id})\n{target_json}")
36
+ if input.prior_belief is not None:
37
+ sections.append(f"# WM BeliefState\n{input.prior_belief.model_dump_json(indent=2)}")
38
+ return "\n\n".join(sections)
39
+
40
+ @staticmethod
41
+ def _select_target_plan(input: SubagentInput) -> str:
42
+ """Render the target plan's JSON body for the USR.
43
+
44
+ Session 11's Brain Executive populates ``prior_plans`` from the
45
+ Planner's outputs and sets ``target_plan_id`` to identify which
46
+ plan the Critic should attack. Here we render the first plan
47
+ (or a placeholder if none) — Session 11 wires up id-based
48
+ lookup once plans carry ids.
49
+ """
50
+ if not input.prior_plans:
51
+ return "(no target plan provided)"
52
+ return input.prior_plans[0].model_dump_json(indent=2)
53
+
54
+ @classmethod
55
+ def empty_fallback(cls, brain: str, target_plan_id: str = "") -> CriticReport:
56
+ # Phase A Decision 6: severity=0 + empty attacks signal "no
57
+ # critique". Brain Executive ignores this Critic's vote weight.
58
+ return CriticReport(
59
+ brain=brain,
60
+ target_plan_id=target_plan_id,
61
+ attacks=[],
62
+ missing_considerations=[],
63
+ would_change_mind_if=[],
64
+ severity=0.0,
65
+ )
66
+
67
+ def run(self, input: SubagentInput, step_idx: int) -> CriticReport: # type: ignore[override]
68
+ return super().run(input, step_idx) # type: ignore[return-value]
cortex/subagents/perception.py ADDED
@@ -0,0 +1,192 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Perception subagent - deterministic Python; not router-callable.
2
+
3
+ Per cortex/CLAUDE.md binding: Perception is pure Python; no LLM calls.
4
+ Phase A Decisions 9 (V2 KeyError) + 63 (salient_signals cap at 5) +
5
+ M-FR-1 (pinned confidence per brain).
6
+
7
+ Perception runs ONCE per brain at tick start (not router-callable).
8
+ The Council Executive (Session 12) calls ``perception_for`` once per
9
+ brain at the start of each tick; the resulting ``PerceptionReport`` is
10
+ plumbed into all subsequent SubagentInputs for that brain in that tick.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ from typing import Callable, Dict, List
16
+
17
+ from cortex.schemas import EvidenceCitation, PerceptionReport
18
+ from CrisisWorldCortex.models import CrisisworldcortexObservation
19
+
20
+ _V2_BRAINS = frozenset({"communications", "equity"})
21
+
22
+ _SALIENT_SIGNALS_CAP = 5 # Phase A Decision 63 / OQ-2
23
+
24
+ _LOGISTICS_THRESHOLDS = {
25
+ "test_kits": 300,
26
+ "hospital_beds_free": 100,
27
+ "mobile_units": 5,
28
+ "vaccine_doses": 500,
29
+ }
30
+
31
+ _HOSPITAL_LOAD_ANOMALY_THRESHOLD = 0.6
32
+
33
+
34
+ def perception_for(brain: str, obs: CrisisworldcortexObservation) -> PerceptionReport:
35
+ """Compute the per-brain Perception report.
36
+
37
+ Args:
38
+ brain: One of {"epidemiology", "logistics", "governance"}.
39
+ obs: The current tick's observation.
40
+
41
+ Raises:
42
+ KeyError: If ``brain`` is V2-deferred or unknown.
43
+ """
44
+ helper = _PERCEPTION_REGISTRY.get(brain)
45
+ if helper is not None:
46
+ return helper(obs)
47
+ if brain in _V2_BRAINS:
48
+ raise KeyError(
49
+ f"V2 brain {brain!r} deferred per Phase A Decision 9; no MVP stub perception."
50
+ )
51
+ raise KeyError(f"unknown brain: {brain!r}")
52
+
53
+
54
+ def _epi_perception(obs: CrisisworldcortexObservation) -> PerceptionReport:
55
+ """Epidemiology perception: top-cases regions + high-hospital-load anomalies."""
56
+ sorted_regions = sorted(obs.regions, key=lambda r: r.reported_cases_d_ago, reverse=True)
57
+ salient_signals: List[str] = []
58
+ evidence: List[EvidenceCitation] = []
59
+
60
+ for r in sorted_regions[:3]:
61
+ if r.reported_cases_d_ago > 0:
62
+ salient_signals.append(f"{r.region}: cases={r.reported_cases_d_ago}")
63
+ evidence.append(
64
+ EvidenceCitation(
65
+ source="telemetry",
66
+ ref=f"{r.region}.reported_cases_d_ago",
67
+ excerpt=str(r.reported_cases_d_ago),
68
+ )
69
+ )
70
+
71
+ if not salient_signals and obs.regions:
72
+ # Fallback: cite the first region so we have at least one signal
73
+ r = obs.regions[0]
74
+ salient_signals.append(f"{r.region}: cases={r.reported_cases_d_ago}")
75
+ evidence.append(
76
+ EvidenceCitation(
77
+ source="telemetry",
78
+ ref=f"{r.region}.reported_cases_d_ago",
79
+ excerpt=str(r.reported_cases_d_ago),
80
+ )
81
+ )
82
+
83
+ salient_signals = salient_signals[:_SALIENT_SIGNALS_CAP]
84
+
85
+ anomalies = [
86
+ f"{r.region}: hospital_load={r.hospital_load:.2f}"
87
+ for r in obs.regions
88
+ if r.hospital_load > _HOSPITAL_LOAD_ANOMALY_THRESHOLD
89
+ ]
90
+
91
+ return PerceptionReport(
92
+ brain="epidemiology",
93
+ salient_signals=salient_signals,
94
+ anomalies=anomalies,
95
+ # M-FR-1: telemetry is delayed and noisy per mm.md; pinned proxy
96
+ confidence=0.7,
97
+ evidence=evidence,
98
+ )
99
+
100
+
101
+ def _logistics_perception(obs: CrisisworldcortexObservation) -> PerceptionReport:
102
+ """Logistics perception: low-resource flags + depleted-resource anomalies."""
103
+ res = obs.resources
104
+ salient_signals: List[str] = []
105
+ evidence: List[EvidenceCitation] = []
106
+
107
+ for resource_name, threshold in _LOGISTICS_THRESHOLDS.items():
108
+ value = getattr(res, resource_name)
109
+ if value < threshold:
110
+ salient_signals.append(f"{resource_name} low: {value}")
111
+ evidence.append(
112
+ EvidenceCitation(
113
+ source="resource",
114
+ ref=f"resources.{resource_name}",
115
+ excerpt=str(value),
116
+ )
117
+ )
118
+
119
+ salient_signals = salient_signals[:_SALIENT_SIGNALS_CAP]
120
+
121
+ anomalies = []
122
+ for resource_name in _LOGISTICS_THRESHOLDS:
123
+ if getattr(res, resource_name) == 0:
124
+ anomalies.append(f"{resource_name}: depleted")
125
+
126
+ return PerceptionReport(
127
+ brain="logistics",
128
+ salient_signals=salient_signals,
129
+ anomalies=anomalies,
130
+ # M-FR-1: resource counts are deterministic, no telemetry noise
131
+ confidence=1.0,
132
+ evidence=evidence,
133
+ )
134
+
135
+
136
+ def _governance_perception(obs: CrisisworldcortexObservation) -> PerceptionReport:
137
+ """Governance perception: active restrictions + legal constraints + about-to-expire anomalies."""
138
+ salient_signals: List[str] = []
139
+ evidence: List[EvidenceCitation] = []
140
+
141
+ for restr in obs.active_restrictions:
142
+ salient_signals.append(f"{restr.region}: {restr.severity} ({restr.ticks_remaining}t)")
143
+ evidence.append(
144
+ EvidenceCitation(
145
+ source="policy",
146
+ ref=f"active_restrictions.{restr.region}",
147
+ excerpt=f"{restr.severity}@{restr.ticks_remaining}",
148
+ )
149
+ )
150
+
151
+ for lc in obs.legal_constraints:
152
+ salient_signals.append(f"legal: {lc.rule_id} blocks {lc.blocked_action}")
153
+ evidence.append(
154
+ EvidenceCitation(
155
+ source="policy",
156
+ ref=f"legal_constraints.{lc.rule_id}",
157
+ excerpt=lc.blocked_action,
158
+ )
159
+ )
160
+
161
+ salient_signals = salient_signals[:_SALIENT_SIGNALS_CAP]
162
+
163
+ has_recent_escalate_national = any(
164
+ ea.accepted
165
+ and ea.action.kind == "escalate"
166
+ and getattr(ea.action, "to_authority", None) == "national"
167
+ for ea in obs.recent_action_log
168
+ )
169
+ anomalies = []
170
+ for restr in obs.active_restrictions:
171
+ if (
172
+ restr.severity == "strict"
173
+ and restr.ticks_remaining <= 1
174
+ and not has_recent_escalate_national
175
+ ):
176
+ anomalies.append(f"{restr.region}: strict expiring without escalation")
177
+
178
+ return PerceptionReport(
179
+ brain="governance",
180
+ salient_signals=salient_signals,
181
+ anomalies=anomalies,
182
+ # M-FR-1: policy state is deterministic
183
+ confidence=1.0,
184
+ evidence=evidence,
185
+ )
186
+
187
+
188
+ _PERCEPTION_REGISTRY: Dict[str, Callable[[CrisisworldcortexObservation], PerceptionReport]] = {
189
+ "epidemiology": _epi_perception,
190
+ "logistics": _logistics_perception,
191
+ "governance": _governance_perception,
192
+ }
cortex/subagents/planner.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Planner LLM subagent.
2
+
3
+ Phase A docs/CORTEX_ARCHITECTURE.md §9 Decision 2: SYS = role + action
4
+ schema (B1's shape); USR = perception + WM BeliefState (full JSON if
5
+ provided per M-FR-4) + last_reward.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import ClassVar, List
11
+
12
+ from pydantic import TypeAdapter
13
+
14
+ from cortex.schemas import CandidatePlan, SubagentInput
15
+ from CrisisWorldCortex.models import NoOp
16
+
17
+ from ._base import _LLMSubagent, load_prompt
18
+
19
+ _PLAN_ADAPTER: TypeAdapter[CandidatePlan] = TypeAdapter(CandidatePlan)
20
+
21
+
22
+ class PlannerSubagent(_LLMSubagent):
23
+ """LLM subagent that emits ``CandidatePlan`` for one brain per call."""
24
+
25
+ _role_name: ClassVar[str] = "planner"
26
+ _output_type: ClassVar[type] = CandidatePlan
27
+ _system_prompt_filename: ClassVar[str] = "planner.txt"
28
+ _SYSTEM_PROMPT_TEMPLATE: ClassVar[str] = load_prompt("planner.txt")
29
+ _ADAPTER: ClassVar[TypeAdapter] = _PLAN_ADAPTER
30
+
31
+ def _build_user_message(self, input: SubagentInput) -> str:
32
+ sections: List[str] = []
33
+ sections.append(f"# Perception\n{input.perception.model_dump_json(indent=2)}")
34
+ if input.prior_belief is not None:
35
+ sections.append(
36
+ "# BeliefState (from this brain's WorldModeler)\n"
37
+ f"{input.prior_belief.model_dump_json(indent=2)}"
38
+ )
39
+ sections.append(f"# Last tick reward: {input.last_reward}")
40
+ sections.append(
41
+ f"# Recent action log: {self._format_action_log(input.recent_action_log_excerpt)}"
42
+ )
43
+ return "\n\n".join(sections)
44
+
45
+ @classmethod
46
+ def empty_fallback(cls, brain: str, target_plan_id: str = "") -> CandidatePlan:
47
+ # Phase A Decision 6: NoOp + confidence=0 means "no signal". The
48
+ # Brain Executive's argmax(expected_value * confidence) picks any
49
+ # non-empty plan over this one.
50
+ return CandidatePlan(
51
+ action_sketch="(empty: planner failed to produce a parseable plan)",
52
+ expected_outer_action=NoOp(),
53
+ expected_value=0.0,
54
+ cost=0.0,
55
+ assumptions=[],
56
+ falsifiers=[],
57
+ confidence=0.0,
58
+ )
59
+
60
+ def run(self, input: SubagentInput, step_idx: int) -> CandidatePlan: # type: ignore[override]
61
+ return super().run(input, step_idx) # type: ignore[return-value]
cortex/subagents/prompts/critic.txt ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ You are the {brain} brain's Critic in the CrisisWorldCortex multi-brain agent.
2
+
3
+ Your job: read a CandidatePlan from this brain's Planner and identify what is wrong with it. Emit a CriticReport as JSON.
4
+
5
+ You write prose critique only. Do NOT propose alternative actions; that is the Planner's role.
6
+
7
+ Output strict JSON matching this schema:
8
+ {{
9
+ "brain": "{brain}",
10
+ "target_plan_id": "{target_plan_id}",
11
+ "attacks": ["<one sentence per concrete attack>"],
12
+ "missing_considerations": ["<one sentence per missed factor>"],
13
+ "would_change_mind_if": ["<observation or fact that would soften your critique>"],
14
+ "severity": <float in [0,1]; 0=plan is fine, 1=do not execute this plan>
15
+ }}
16
+
17
+ Use prose strings inside the lists. Do NOT emit JSON action variants.
18
+
19
+ Emit ONLY the JSON. No prose outside the JSON, no code fences.
cortex/subagents/prompts/planner.txt ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ You are the {brain} brain's Planner in the CrisisWorldCortex multi-brain agent.
2
+
3
+ Your job: read the perception summary and prior BeliefState, then emit ONE candidate plan as JSON.
4
+
5
+ Output strict JSON matching this schema:
6
+ {{
7
+ "action_sketch": "<one sentence describing the plan>",
8
+ "expected_outer_action": <one of the action variants below>,
9
+ "expected_value": <float; estimated reward delta for this tick>,
10
+ "cost": <float; relative cost of this plan>,
11
+ "assumptions": ["<assumption 1>"],
12
+ "falsifiers": ["<observation that would refute this plan>"],
13
+ "confidence": <float in [0,1]>
14
+ }}
15
+
16
+ Action variants for expected_outer_action (exactly one kind):
17
+ - {{"kind":"deploy_resource", "region":"R1|R2|R3|R4", "resource_type":"test_kits|hospital_beds|mobile_units|vaccine_doses", "quantity": <int >= 0>}}
18
+ - {{"kind":"request_data", "region":"<id>", "data_type":"case_survey|hospital_audit|compliance_check"}}
19
+ - {{"kind":"restrict_movement", "region":"<id>", "severity":"none|light|moderate|strict"}}
20
+ - {{"kind":"escalate", "to_authority":"regional|national"}}
21
+ - {{"kind":"reallocate_budget", "from_resource":"<resource_type>", "to_resource":"<resource_type>", "amount": <int >= 0>}}
22
+ - {{"kind":"no_op"}}
23
+
24
+ Strict severity may require a prior escalate(national) - check legal_constraints.
25
+
26
+ Emit ONLY the JSON. No prose, no code fences.
cortex/subagents/prompts/world_modeler.txt ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ You are the {brain} brain's World Modeler in the CrisisWorldCortex multi-brain agent.
2
+
3
+ Your job: read the perception summary plus prior belief (round 2 only), and emit a BeliefState as JSON.
4
+
5
+ A BeliefState describes what you think the latent epidemiological / logistical / governance state is for each region, citing evidence from observed telemetry, resources, policy state, or recent actions.
6
+
7
+ Output strict JSON matching this schema:
8
+ {{
9
+ "brain": "{brain}",
10
+ "latent_estimates": {{
11
+ "<region_id>": {{
12
+ "estimated_infection_rate": <float in [0,1]>,
13
+ "estimated_r_effective": <float >= 0>,
14
+ "estimated_compliance": <float in [0,1]>,
15
+ "confidence_intervals": {{}}
16
+ }}
17
+ }},
18
+ "hypotheses": [
19
+ {{"label": "<short>", "weight": <float in [0,1]>, "explanation": "<one sentence>"}}
20
+ ],
21
+ "uncertainty": <float in [0,1]>,
22
+ "reducible_by_more_thought": <float in [0,1]>,
23
+ "evidence": [
24
+ {{"source": "telemetry|resource|policy|action_log|belief|memory",
25
+ "ref": "<concise pointer like region=R2.hospital_load@tick=7>",
26
+ "excerpt": "<the value or text>"}}
27
+ ]
28
+ }}
29
+
30
+ Cite at least 2 EvidenceCitations. Uncited claims zero your protocol-integrity reward.
31
+
32
+ Emit ONLY the JSON. No prose, no code fences.
cortex/subagents/world_modeler.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """WorldModeler LLM subagent.
2
+
3
+ Phase A docs/CORTEX_ARCHITECTURE.md §9 Decision 1: SYS = role + schema;
4
+ USR = perception + last_reward + recent_action_log_excerpt (with prior
5
+ BeliefState in round 2 per Decision 62).
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import ClassVar, List
11
+
12
+ from pydantic import TypeAdapter
13
+
14
+ from cortex.schemas import BeliefState, SubagentInput
15
+
16
+ from ._base import _LLMSubagent, load_prompt
17
+
18
+ _BELIEF_ADAPTER: TypeAdapter[BeliefState] = TypeAdapter(BeliefState)
19
+ """Module-level constant per Phase A Decision 5 (encapsulation; avoid
20
+ circular imports through the package init)."""
21
+
22
+
23
+ class WorldModelerSubagent(_LLMSubagent):
24
+ """LLM subagent that emits ``BeliefState`` for one brain per call."""
25
+
26
+ _role_name: ClassVar[str] = "world_modeler"
27
+ _output_type: ClassVar[type] = BeliefState
28
+ _system_prompt_filename: ClassVar[str] = "world_modeler.txt"
29
+ _SYSTEM_PROMPT_TEMPLATE: ClassVar[str] = load_prompt("world_modeler.txt")
30
+ _ADAPTER: ClassVar[TypeAdapter] = _BELIEF_ADAPTER
31
+
32
+ def _build_user_message(self, input: SubagentInput) -> str:
33
+ sections: List[str] = []
34
+ sections.append(f"# Perception\n{input.perception.model_dump_json(indent=2)}")
35
+ if input.prior_belief is not None:
36
+ sections.append(
37
+ "# Prior BeliefState (round 1 result)\n"
38
+ f"{input.prior_belief.model_dump_json(indent=2)}"
39
+ )
40
+ sections.append(f"# Last tick reward: {input.last_reward}")
41
+ sections.append(
42
+ f"# Recent action log: {self._format_action_log(input.recent_action_log_excerpt)}"
43
+ )
44
+ return "\n\n".join(sections)
45
+
46
+ @classmethod
47
+ def empty_fallback(cls, brain: str, target_plan_id: str = "") -> BeliefState:
48
+ # Phase A Decision 6 + Decision 62: empty BeliefState as the
49
+ # honest "no signal" state. uncertainty=1.0, no evidence -> r_proto = 0.
50
+ return BeliefState(
51
+ brain=brain,
52
+ latent_estimates={},
53
+ hypotheses=[],
54
+ uncertainty=1.0,
55
+ reducible_by_more_thought=0.0,
56
+ evidence=[],
57
+ )
58
+
59
+ # Narrow run() return type for callers (refinement #1).
60
+ def run(self, input: SubagentInput, step_idx: int) -> BeliefState: # type: ignore[override]
61
+ return super().run(input, step_idx) # type: ignore[return-value]
demo/CLAUDE.md CHANGED
@@ -1,45 +1,45 @@
1
- # demo/CLAUDE.md
2
-
3
- Replay-only visualization. Live demos fail under judging pressure; ship canned scenarios.
4
-
5
- ## Belongs here
6
-
7
- - `visualizer/trace_renderer.py` — renders a JSON trace as a "council in action" view.
8
- - `visualizer/reward_curve_plot.py` — plots reward curves from training logs.
9
- - `demo_scenarios/*.json` — pre-recorded trajectories for the pitch (e.g. `scenario_flat_fails.json`, `scenario_cortex_holds_dissent.json`).
10
-
11
- ## Does not belong here
12
-
13
- Live agent execution (record offline, replay here). Training logic. Graders.
14
-
15
- ## Allowed imports
16
-
17
- - `cortex.schemas` — typed parse of trace JSON. Types only, no logic.
18
- - stdlib + plotting libs (matplotlib / plotly).
19
-
20
- ## Forbidden imports
21
-
22
- - `server/*`, `training/*`, `baselines/*`.
23
- - `cortex.council`, `cortex.routing_policy` — if you need to re-run the agent, do it offline and ship a new JSON.
24
-
25
- ## Binding contracts
26
-
27
- - Every JSON scenario conforms to `cortex.schemas.Trajectory`.
28
- - Rendering is deterministic: same JSON → same output, modulo timestamps.
29
- - The pitch-demo scenario must showcase B2 overcommit/misallocate vs Cortex dissent-preservation (design §27).
30
- - A pre-recorded demo video (MP4) lives alongside the JSON scenarios as the live-demo fallback.
31
-
32
- ## Public APIs (owned here)
33
-
34
- - `render_trace(json_path: str, out_path: str) -> None`
35
- - `plot_reward_curves(log_paths: list[str], out_path: str) -> None`
36
-
37
- ## Testing requirements
38
-
39
- - Each committed JSON scenario parses into a `Trajectory` without error.
40
- - `render_trace` produces a non-empty output file for each scenario.
41
-
42
- ## Common failure modes
43
-
44
- - Live re-run during the demo — network/Colab flakiness kills the pitch. Replay only.
45
- - Renderer depending on a `cortex.council` instance — import breaks when Cortex API shifts. Keep read-only on types.
 
1
+ # demo/CLAUDE.md
2
+
3
+ Replay-only visualization. Live demos fail under judging pressure; ship canned scenarios.
4
+
5
+ ## Belongs here
6
+
7
+ - `visualizer/trace_renderer.py` — renders a JSON trace as a "council in action" view.
8
+ - `visualizer/reward_curve_plot.py` — plots reward curves from training logs.
9
+ - `demo_scenarios/*.json` — pre-recorded trajectories for the pitch (e.g. `scenario_flat_fails.json`, `scenario_cortex_holds_dissent.json`).
10
+
11
+ ## Does not belong here
12
+
13
+ Live agent execution (record offline, replay here). Training logic. Graders.
14
+
15
+ ## Allowed imports
16
+
17
+ - `cortex.schemas` — typed parse of trace JSON. Types only, no logic.
18
+ - stdlib + plotting libs (matplotlib / plotly).
19
+
20
+ ## Forbidden imports
21
+
22
+ - `server/*`, `training/*`, `baselines/*`.
23
+ - `cortex.council`, `cortex.routing_policy` — if you need to re-run the agent, do it offline and ship a new JSON.
24
+
25
+ ## Binding contracts
26
+
27
+ - Every JSON scenario conforms to `cortex.schemas.Trajectory`.
28
+ - Rendering is deterministic: same JSON → same output, modulo timestamps.
29
+ - The pitch-demo scenario must showcase B2 overcommit/misallocate vs Cortex dissent-preservation (design §27).
30
+ - A pre-recorded demo video (MP4) lives alongside the JSON scenarios as the live-demo fallback.
31
+
32
+ ## Public APIs (owned here)
33
+
34
+ - `render_trace(json_path: str, out_path: str) -> None`
35
+ - `plot_reward_curves(log_paths: list[str], out_path: str) -> None`
36
+
37
+ ## Testing requirements
38
+
39
+ - Each committed JSON scenario parses into a `Trajectory` without error.
40
+ - `render_trace` produces a non-empty output file for each scenario.
41
+
42
+ ## Common failure modes
43
+
44
+ - Live re-run during the demo — network/Colab flakiness kills the pitch. Replay only.
45
+ - Renderer depending on a `cortex.council` instance — import breaks when Cortex API shifts. Keep read-only on types.
docs/CORTEX_ARCHITECTURE.md CHANGED
@@ -705,7 +705,7 @@ Decisions are grouped by layer. Each entry: **decision** / **rationale** /
705
 
706
  19. **Reasoning summary: a 1–2 sentence string from the Planner's `action_sketch`.** / Fits the 400-char `BrainRecommendation.reasoning_summary` cap. / Considered an LLM call to produce a summary; rejected — Brain Executive must be Python-only per cortex/CLAUDE.md.
707
 
708
- 20. **`evidence` field on `BrainRecommendation` = union of all `EvidenceCitation` lists from `BeliefState`, `CandidatePlan`, `CriticReport`.** / Ensures the council sees the brain's full evidence chain. / Considered Critic only; rejected — claims with no upstream evidence get zeroed `r_proto`.
709
 
710
  21. **Brain identifier strings: `"epidemiology"`, `"logistics"`, `"governance"` (lowercase, full word).** / Readable and grep-friendly. / Considered abbreviations (epi, log, gov); rejected — log-grep collisions.
711
 
 
705
 
706
  19. **Reasoning summary: a 1–2 sentence string from the Planner's `action_sketch`.** / Fits the 400-char `BrainRecommendation.reasoning_summary` cap. / Considered an LLM call to produce a summary; rejected — Brain Executive must be Python-only per cortex/CLAUDE.md.
707
 
708
+ 20. **`evidence` field on `BrainRecommendation` = union of all `EvidenceCitation` lists from `BeliefState`, `CandidatePlan`, `CriticReport`.** / Ensures the council sees the brain's full evidence chain. / Considered Critic only; rejected — claims with no upstream evidence get zeroed `r_proto`. **(Session 11 implementation note — M-FR-3)** Implementation reads evidence from `PerceptionReport.evidence` + `BeliefState.evidence` only, since `CandidatePlan` and `CriticReport` schemas (Session 9) carry no `evidence` field. Adding evidence fields to those schemas was rejected as schema-churn risk; the perception+beliefs union captures the actionable evidence chain since plans/critics derive from beliefs. See `cortex/brains/_executive.py:aggregate_brain_outputs`.
709
 
710
  21. **Brain identifier strings: `"epidemiology"`, `"logistics"`, `"governance"` (lowercase, full word).** / Readable and grep-friendly. / Considered abbreviations (epi, log, gov); rejected — log-grep collisions.
711
 
inference.py CHANGED
@@ -40,10 +40,15 @@ from __future__ import annotations
40
  import os
41
  import sys
42
  from dataclasses import dataclass
 
 
 
 
 
43
  from typing import Any, Dict, List, Optional
44
 
45
- from baselines.flat_agent import B1FlatAgent, B1StepEvent
46
- from cortex.llm_client import LLMClient
47
  from CrisisWorldCortex.models import OuterActionPayload
48
  from CrisisWorldCortex.server.graders import terminal_bonus
49
  from CrisisWorldCortex.server.simulator import WorldState
@@ -62,8 +67,8 @@ DEFAULT_MODEL = "Qwen/Qwen2.5-72B-Instruct"
62
  # distinct seeds per task for cross-episode reproducibility.
63
  TASK_CONFIGS: List[dict] = [
64
  {"task_name": "outbreak_easy", "seed": 0, "max_ticks": 12},
65
- {"task_name": "outbreak_medium", "seed": 1, "max_ticks": 12},
66
- {"task_name": "outbreak_hard", "seed": 2, "max_ticks": 12},
67
  ]
68
 
69
  # Score-clamp bounds keep .3f formatting strictly inside (0, 1) so the
 
40
  import os
41
  import sys
42
  from dataclasses import dataclass
43
+ try:
44
+ from dotenv import load_dotenv
45
+ load_dotenv()
46
+ except ImportError:
47
+ pass
48
  from typing import Any, Dict, List, Optional
49
 
50
+ from CrisisWorldCortex.baselines.flat_agent import B1FlatAgent, B1StepEvent
51
+ from CrisisWorldCortex.cortex.llm_client import LLMClient
52
  from CrisisWorldCortex.models import OuterActionPayload
53
  from CrisisWorldCortex.server.graders import terminal_bonus
54
  from CrisisWorldCortex.server.simulator import WorldState
 
67
  # distinct seeds per task for cross-episode reproducibility.
68
  TASK_CONFIGS: List[dict] = [
69
  {"task_name": "outbreak_easy", "seed": 0, "max_ticks": 12},
70
+ # {"task_name": "outbreak_medium", "seed": 1, "max_ticks": 12},
71
+ # {"task_name": "outbreak_hard", "seed": 2, "max_ticks": 12},
72
  ]
73
 
74
  # Score-clamp bounds keep .3f formatting strictly inside (0, 1) so the
notebooks/train_b1_grpo.ipynb ADDED
@@ -0,0 +1,496 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "metadata": {},
6
+ "source": [
7
+ "# B1 GRPO training on Qwen3-1.7B (Workstream B Phase 3)\n",
8
+ "\n",
9
+ "Trains the **B1 flat-agent baseline** with [Unsloth](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide) + [TRL GRPO](https://huggingface.co/docs/trl/main/en/grpo_trainer) against the deployed CrisisWorldCortex HF Space env.\n",
10
+ "\n",
11
+ "**One-shot run:** `Runtime → Run all` on a fresh Colab T4. 300 training steps, ~30 minutes wall-clock. Saves the trained LoRA adapter to your HF Hub at the end.\n",
12
+ "\n",
13
+ "**Reward source:** the Phase-1-fixed `outer_reward` (range `[-1.0, 1.0]`, signal-quality gates passed). Single-step GRPO — each rollout = one env reset + one env step.\n",
14
+ "\n",
15
+ "**Prereqs:**\n",
16
+ "1. Colab Secrets has `HF_TOKEN` set (Tools → Secrets, name = `HF_TOKEN`, value = your `hf_xxx` token with write access).\n",
17
+ "2. The HF Space `Angshuman28/CrisisWorldCortex` is running with the post-Phase-1 reward.\n"
18
+ ]
19
+ },
20
+ {
21
+ "cell_type": "markdown",
22
+ "metadata": {},
23
+ "source": [
24
+ "## 1. Install dependencies\n",
25
+ "\n",
26
+ "Unsloth pulls a custom torch + vllm + xformers stack tuned for free Colab T4. Pin trl to a version compatible with `GRPOTrainer` (>=0.14)."
27
+ ]
28
+ },
29
+ {
30
+ "cell_type": "code",
31
+ "execution_count": null,
32
+ "metadata": {},
33
+ "outputs": [],
34
+ "source": [
35
+ "%%capture\n",
36
+ "!pip install --upgrade pip\n",
37
+ "!pip install unsloth vllm\n",
38
+ "!pip install --upgrade --no-deps \"trl>=0.14\" peft accelerate bitsandbytes\n",
39
+ "!pip install pydantic openenv huggingface_hub matplotlib"
40
+ ]
41
+ },
42
+ {
43
+ "cell_type": "markdown",
44
+ "metadata": {},
45
+ "source": [
46
+ "## 2. Authenticate with Hugging Face\n",
47
+ "\n",
48
+ "Reads `HF_TOKEN` from Colab Secrets. Falls back to interactive login if not set."
49
+ ]
50
+ },
51
+ {
52
+ "cell_type": "code",
53
+ "execution_count": null,
54
+ "metadata": {},
55
+ "outputs": [],
56
+ "source": [
57
+ "import os\n",
58
+ "\n",
59
+ "try:\n",
60
+ " from google.colab import userdata\n",
61
+ " HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
62
+ " os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
63
+ "except Exception:\n",
64
+ " from huggingface_hub import login\n",
65
+ " login()\n",
66
+ " HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
67
+ "\n",
68
+ "assert HF_TOKEN, \"HF_TOKEN is required (set in Colab Secrets or via login())\"\n",
69
+ "print(\"HF auth OK\")"
70
+ ]
71
+ },
72
+ {
73
+ "cell_type": "markdown",
74
+ "metadata": {},
75
+ "source": [
76
+ "## 3. Clone CrisisWorldCortex and install\n",
77
+ "\n",
78
+ "Pulls the deployed HF Space's repo and installs locally. Provides the `CrisisworldcortexEnv` HTTP client + the `baselines.flat_agent` system prompt + parser used by B1."
79
+ ]
80
+ },
81
+ {
82
+ "cell_type": "code",
83
+ "execution_count": null,
84
+ "metadata": {},
85
+ "outputs": [],
86
+ "source": [
87
+ "%%capture\n",
88
+ "!rm -rf /content/CrisisWorldCortex\n",
89
+ "!git clone https://huggingface.co/spaces/Angshuman28/CrisisWorldCortex /content/CrisisWorldCortex\n",
90
+ "%cd /content/CrisisWorldCortex\n",
91
+ "!pip install -e ."
92
+ ]
93
+ },
94
+ {
95
+ "cell_type": "code",
96
+ "execution_count": null,
97
+ "metadata": {},
98
+ "outputs": [],
99
+ "source": [
100
+ "# Sanity: imports resolve, env client constructs.\n",
101
+ "import sys\n",
102
+ "sys.path.insert(0, \"/content/CrisisWorldCortex\")\n",
103
+ "\n",
104
+ "from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexObservation\n",
105
+ "from CrisisWorldCortex.client import CrisisworldcortexEnv\n",
106
+ "from baselines.flat_agent import (\n",
107
+ " build_system_prompt,\n",
108
+ " parse_action,\n",
109
+ " parse_failure_marker,\n",
110
+ " serialize_observation,\n",
111
+ ")\n",
112
+ "print(\"CrisisWorld imports OK\")"
113
+ ]
114
+ },
115
+ {
116
+ "cell_type": "markdown",
117
+ "metadata": {},
118
+ "source": [
119
+ "## 4. Load Qwen3-1.7B with LoRA via Unsloth\n",
120
+ "\n",
121
+ "Qwen3-1.7B fits comfortably on a T4 with 4-bit quantization. LoRA rank 32 — enough to learn the JSON-action format and modest policy improvements; cheap to merge."
122
+ ]
123
+ },
124
+ {
125
+ "cell_type": "code",
126
+ "execution_count": null,
127
+ "metadata": {},
128
+ "outputs": [],
129
+ "source": [
130
+ "from unsloth import FastLanguageModel\n",
131
+ "import torch\n",
132
+ "\n",
133
+ "MAX_SEQ_LEN = 4096\n",
134
+ "MODEL_NAME = \"unsloth/Qwen3-1.7B\"\n",
135
+ "\n",
136
+ "model, tokenizer = FastLanguageModel.from_pretrained(\n",
137
+ " model_name=MODEL_NAME,\n",
138
+ " max_seq_length=MAX_SEQ_LEN,\n",
139
+ " load_in_4bit=True,\n",
140
+ " fast_inference=True, # vLLM-backed generate, required by GRPOTrainer\n",
141
+ " max_lora_rank=32,\n",
142
+ " gpu_memory_utilization=0.6,\n",
143
+ ")\n",
144
+ "\n",
145
+ "model = FastLanguageModel.get_peft_model(\n",
146
+ " model,\n",
147
+ " r=32,\n",
148
+ " target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
149
+ " \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
150
+ " lora_alpha=64,\n",
151
+ " use_gradient_checkpointing=\"unsloth\",\n",
152
+ " random_state=42,\n",
153
+ ")\n",
154
+ "print(\"Model + LoRA loaded\")"
155
+ ]
156
+ },
157
+ {
158
+ "cell_type": "markdown",
159
+ "metadata": {},
160
+ "source": [
161
+ "## 5. Connect to the deployed CrisisWorld env\n",
162
+ "\n",
163
+ "Uses the public HF Space URL. Each rollout calls `env.reset()` then `env.step(action)` once."
164
+ ]
165
+ },
166
+ {
167
+ "cell_type": "code",
168
+ "execution_count": null,
169
+ "metadata": {},
170
+ "outputs": [],
171
+ "source": [
172
+ "ENV_URL = \"https://angshuman28-crisisworldcortex.hf.space\"\n",
173
+ "TASKS = (\"outbreak_easy\", \"outbreak_medium\", \"outbreak_hard\")\n",
174
+ "EPISODE_TICKS = 12\n",
175
+ "\n",
176
+ "def make_env() -> CrisisworldcortexEnv:\n",
177
+ " return CrisisworldcortexEnv(base_url=ENV_URL)\n",
178
+ "\n",
179
+ "_test_env = make_env()\n",
180
+ "_obs = _test_env.reset(task_name=\"outbreak_easy\", seed=0, max_ticks=EPISODE_TICKS)\n",
181
+ "print(f\"Env OK. Initial tick={_obs.tick}, regions={[r.region for r in _obs.regions]}\")"
182
+ ]
183
+ },
184
+ {
185
+ "cell_type": "markdown",
186
+ "metadata": {},
187
+ "source": [
188
+ "## 6. Build the prompt dataset and reward function\n",
189
+ "\n",
190
+ "Each example in the dataset is a `(task, seed)` pair. The reward function:\n",
191
+ "1. Resets the env to that `(task, seed)`.\n",
192
+ "2. Parses the model's completion as a `OuterActionPayload`.\n",
193
+ "3. Submits to the env, returns `obs.reward` (post-Phase-1 range `[-1, 1]`).\n",
194
+ "4. On parse failure, submits `parse_failure_marker()` so the §19 -1.0 + terminate contract fires."
195
+ ]
196
+ },
197
+ {
198
+ "cell_type": "code",
199
+ "execution_count": null,
200
+ "metadata": {},
201
+ "outputs": [],
202
+ "source": [
203
+ "from datasets import Dataset\n",
204
+ "import random\n",
205
+ "\n",
206
+ "SYSTEM_PROMPT = build_system_prompt()\n",
207
+ "\n",
208
+ "def build_user_prompt(obs: CrisisworldcortexObservation) -> str:\n",
209
+ " return serialize_observation(obs)\n",
210
+ "\n",
211
+ "def make_chat_prompt(obs: CrisisworldcortexObservation) -> str:\n",
212
+ " return tokenizer.apply_chat_template(\n",
213
+ " [\n",
214
+ " {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
215
+ " {\"role\": \"user\", \"content\": build_user_prompt(obs)},\n",
216
+ " ],\n",
217
+ " tokenize=False,\n",
218
+ " add_generation_prompt=True,\n",
219
+ " )\n",
220
+ "\n",
221
+ "rng = random.Random(0)\n",
222
+ "_seed_pool = []\n",
223
+ "for task in TASKS:\n",
224
+ " for seed in range(50):\n",
225
+ " _seed_pool.append({\"task\": task, \"seed\": seed})\n",
226
+ "rng.shuffle(_seed_pool)\n",
227
+ "\n",
228
+ "_prompts = []\n",
229
+ "_meta = []\n",
230
+ "for entry in _seed_pool:\n",
231
+ " env = make_env()\n",
232
+ " obs = env.reset(task_name=entry[\"task\"], seed=entry[\"seed\"], max_ticks=EPISODE_TICKS)\n",
233
+ " _prompts.append(make_chat_prompt(obs))\n",
234
+ " _meta.append(entry)\n",
235
+ "\n",
236
+ "train_dataset = Dataset.from_dict({\n",
237
+ " \"prompt\": _prompts,\n",
238
+ " \"task\": [m[\"task\"] for m in _meta],\n",
239
+ " \"seed\": [m[\"seed\"] for m in _meta],\n",
240
+ "})\n",
241
+ "print(f\"Dataset built: {len(train_dataset)} examples\")"
242
+ ]
243
+ },
244
+ {
245
+ "cell_type": "code",
246
+ "execution_count": null,
247
+ "metadata": {},
248
+ "outputs": [],
249
+ "source": [
250
+ "def crisisworld_reward(\n",
251
+ " prompts: list[str],\n",
252
+ " completions: list[str],\n",
253
+ " task: list[str],\n",
254
+ " seed: list[int],\n",
255
+ " **_kwargs,\n",
256
+ ") -> list[float]:\n",
257
+ " \"\"\"GRPO reward function: one env step per (prompt, completion) pair.\n",
258
+ "\n",
259
+ " Reward source: Phase-1-fixed env outer_reward in [-1, 1].\n",
260
+ " Parse failure → submits parse_failure_marker → r_policy = -1.0 + terminate.\n",
261
+ " \"\"\"\n",
262
+ " rewards: list[float] = []\n",
263
+ " for completion, t, s in zip(completions, task, seed):\n",
264
+ " env = make_env()\n",
265
+ " env.reset(task_name=t, seed=int(s), max_ticks=EPISODE_TICKS)\n",
266
+ " action_payload = parse_action(completion)\n",
267
+ " if action_payload is None:\n",
268
+ " action_payload = parse_failure_marker()\n",
269
+ " try:\n",
270
+ " result = env.step(CrisisworldcortexAction(action=action_payload))\n",
271
+ " reward = result.observation.reward if hasattr(result, \"observation\") else result.reward\n",
272
+ " rewards.append(float(reward) if reward is not None else 0.0)\n",
273
+ " except Exception as exc:\n",
274
+ " print(f\"[WARN] env.step failed task={t} seed={s}: {exc}\")\n",
275
+ " rewards.append(-1.0)\n",
276
+ " return rewards"
277
+ ]
278
+ },
279
+ {
280
+ "cell_type": "markdown",
281
+ "metadata": {},
282
+ "source": [
283
+ "## 7. GRPO training\n",
284
+ "\n",
285
+ "300 steps × group size 4 = 1200 rollouts. Each rollout is one HF Space round-trip (~1s) — total ~20–30 min on T4."
286
+ ]
287
+ },
288
+ {
289
+ "cell_type": "code",
290
+ "execution_count": null,
291
+ "metadata": {},
292
+ "outputs": [],
293
+ "source": [
294
+ "from trl import GRPOConfig, GRPOTrainer\n",
295
+ "\n",
296
+ "MAX_TRAIN_STEPS = 300\n",
297
+ "GROUP_SIZE = 4\n",
298
+ "MAX_PROMPT_LEN = 2048\n",
299
+ "MAX_COMPLETION_LEN = 512\n",
300
+ "\n",
301
+ "training_args = GRPOConfig(\n",
302
+ " output_dir=\"/content/b1_grpo_output\",\n",
303
+ " learning_rate=5e-6,\n",
304
+ " per_device_train_batch_size=GROUP_SIZE,\n",
305
+ " gradient_accumulation_steps=1,\n",
306
+ " num_generations=GROUP_SIZE,\n",
307
+ " max_prompt_length=MAX_PROMPT_LEN,\n",
308
+ " max_completion_length=MAX_COMPLETION_LEN,\n",
309
+ " max_steps=MAX_TRAIN_STEPS,\n",
310
+ " save_steps=100,\n",
311
+ " logging_steps=5,\n",
312
+ " report_to=\"none\",\n",
313
+ " bf16=True,\n",
314
+ " optim=\"adamw_8bit\",\n",
315
+ " temperature=0.8,\n",
316
+ " use_vllm=True,\n",
317
+ " vllm_mode=\"colocate\",\n",
318
+ " seed=42,\n",
319
+ ")\n",
320
+ "\n",
321
+ "trainer = GRPOTrainer(\n",
322
+ " model=model,\n",
323
+ " processing_class=tokenizer,\n",
324
+ " reward_funcs=[crisisworld_reward],\n",
325
+ " args=training_args,\n",
326
+ " train_dataset=train_dataset,\n",
327
+ ")\n",
328
+ "print(\"GRPOTrainer constructed; starting train()...\")"
329
+ ]
330
+ },
331
+ {
332
+ "cell_type": "code",
333
+ "execution_count": null,
334
+ "metadata": {},
335
+ "outputs": [],
336
+ "source": [
337
+ "trainer.train()"
338
+ ]
339
+ },
340
+ {
341
+ "cell_type": "markdown",
342
+ "metadata": {},
343
+ "source": [
344
+ "## 8. Save the trained LoRA adapter to HF Hub\n",
345
+ "\n",
346
+ "Pushes to `<your-username>/crisisworld-b1-grpo-qwen3-1p7b`. Change the repo name below if you want a different namespace."
347
+ ]
348
+ },
349
+ {
350
+ "cell_type": "code",
351
+ "execution_count": null,
352
+ "metadata": {},
353
+ "outputs": [],
354
+ "source": [
355
+ "from huggingface_hub import HfApi\n",
356
+ "\n",
357
+ "HUB_REPO = \"Angshuman28/crisisworld-b1-grpo-qwen3-1p7b\"\n",
358
+ "\n",
359
+ "model.save_pretrained(\"/content/b1_grpo_lora\")\n",
360
+ "tokenizer.save_pretrained(\"/content/b1_grpo_lora\")\n",
361
+ "\n",
362
+ "api = HfApi()\n",
363
+ "api.create_repo(HUB_REPO, exist_ok=True, repo_type=\"model\", private=False, token=HF_TOKEN)\n",
364
+ "api.upload_folder(\n",
365
+ " folder_path=\"/content/b1_grpo_lora\",\n",
366
+ " repo_id=HUB_REPO,\n",
367
+ " repo_type=\"model\",\n",
368
+ " token=HF_TOKEN,\n",
369
+ ")\n",
370
+ "print(f\"Saved to https://huggingface.co/{HUB_REPO}\")"
371
+ ]
372
+ },
373
+ {
374
+ "cell_type": "markdown",
375
+ "metadata": {},
376
+ "source": [
377
+ "## 9. Eval: trained adapter vs base model on 3 tasks\n",
378
+ "\n",
379
+ "Runs a single full episode (12 ticks) per task, per model. Reports cumulative reward."
380
+ ]
381
+ },
382
+ {
383
+ "cell_type": "code",
384
+ "execution_count": null,
385
+ "metadata": {},
386
+ "outputs": [],
387
+ "source": [
388
+ "def _hf_chat(model_inst, tokenizer_inst, system: str, user: str, max_new_tokens: int = 256) -> str:\n",
389
+ " prompt = tokenizer_inst.apply_chat_template(\n",
390
+ " [{\"role\": \"system\", \"content\": system}, {\"role\": \"user\", \"content\": user}],\n",
391
+ " tokenize=False,\n",
392
+ " add_generation_prompt=True,\n",
393
+ " )\n",
394
+ " inputs = tokenizer_inst(prompt, return_tensors=\"pt\").to(model_inst.device)\n",
395
+ " with torch.no_grad():\n",
396
+ " out = model_inst.generate(\n",
397
+ " **inputs, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.0,\n",
398
+ " )\n",
399
+ " return tokenizer_inst.decode(out[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n",
400
+ "\n",
401
+ "def run_one_episode(model_inst, tokenizer_inst, task: str, seed: int) -> float:\n",
402
+ " env = make_env()\n",
403
+ " obs = env.reset(task_name=task, seed=seed, max_ticks=EPISODE_TICKS)\n",
404
+ " cumulative = 0.0\n",
405
+ " for tick in range(EPISODE_TICKS):\n",
406
+ " completion = _hf_chat(model_inst, tokenizer_inst, SYSTEM_PROMPT, serialize_observation(obs))\n",
407
+ " action = parse_action(completion) or parse_failure_marker()\n",
408
+ " result = env.step(CrisisworldcortexAction(action=action))\n",
409
+ " obs = result.observation if hasattr(result, \"observation\") else result\n",
410
+ " reward = obs.reward if obs.reward is not None else 0.0\n",
411
+ " cumulative += reward\n",
412
+ " if obs.done:\n",
413
+ " break\n",
414
+ " return cumulative\n",
415
+ "\n",
416
+ "FastLanguageModel.for_inference(model)\n",
417
+ "trained_results = {t: run_one_episode(model, tokenizer, t, seed=0) for t in TASKS}\n",
418
+ "print(f\"Trained model cumulative reward per task: {trained_results}\")"
419
+ ]
420
+ },
421
+ {
422
+ "cell_type": "code",
423
+ "execution_count": null,
424
+ "metadata": {},
425
+ "outputs": [],
426
+ "source": [
427
+ "# Reload base Qwen3-1.7B (no LoRA) for the comparison.\n",
428
+ "base_model, base_tokenizer = FastLanguageModel.from_pretrained(\n",
429
+ " model_name=MODEL_NAME,\n",
430
+ " max_seq_length=MAX_SEQ_LEN,\n",
431
+ " load_in_4bit=True,\n",
432
+ " fast_inference=False,\n",
433
+ ")\n",
434
+ "FastLanguageModel.for_inference(base_model)\n",
435
+ "base_results = {t: run_one_episode(base_model, base_tokenizer, t, seed=0) for t in TASKS}\n",
436
+ "print(f\"Base model cumulative reward per task: {base_results}\")"
437
+ ]
438
+ },
439
+ {
440
+ "cell_type": "markdown",
441
+ "metadata": {},
442
+ "source": [
443
+ "## 10. Plot eval comparison\n",
444
+ "\n",
445
+ "Bar chart: trained vs base, cumulative episode reward by task."
446
+ ]
447
+ },
448
+ {
449
+ "cell_type": "code",
450
+ "execution_count": null,
451
+ "metadata": {},
452
+ "outputs": [],
453
+ "source": [
454
+ "import matplotlib.pyplot as plt\n",
455
+ "import numpy as np\n",
456
+ "\n",
457
+ "task_names = list(TASKS)\n",
458
+ "trained_vals = [trained_results[t] for t in task_names]\n",
459
+ "base_vals = [base_results[t] for t in task_names]\n",
460
+ "\n",
461
+ "x = np.arange(len(task_names))\n",
462
+ "width = 0.35\n",
463
+ "\n",
464
+ "fig, ax = plt.subplots(figsize=(9, 5))\n",
465
+ "ax.bar(x - width/2, base_vals, width, label=\"Base Qwen3-1.7B\")\n",
466
+ "ax.bar(x + width/2, trained_vals, width, label=\"GRPO-trained Qwen3-1.7B\")\n",
467
+ "ax.set_xticks(x)\n",
468
+ "ax.set_xticklabels(task_names)\n",
469
+ "ax.set_ylabel(\"Cumulative episode reward\")\n",
470
+ "ax.set_title(\"B1 GRPO: trained vs base\")\n",
471
+ "ax.legend()\n",
472
+ "ax.axhline(0.0, linestyle=\":\", color=\"grey\")\n",
473
+ "plt.tight_layout()\n",
474
+ "plt.show()"
475
+ ]
476
+ }
477
+ ],
478
+ "metadata": {
479
+ "kernelspec": {
480
+ "display_name": "Python 3",
481
+ "language": "python",
482
+ "name": "python3"
483
+ },
484
+ "language_info": {
485
+ "name": "python",
486
+ "version": "3.10"
487
+ },
488
+ "accelerator": "GPU",
489
+ "colab": {
490
+ "gpuType": "T4",
491
+ "provenance": []
492
+ }
493
+ },
494
+ "nbformat": 4,
495
+ "nbformat_minor": 5
496
+ }
notebooks/train_cortex_router_grpo.ipynb ADDED
@@ -0,0 +1,567 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "id": "7fb27b941602401d91542211134fc71a",
6
+ "metadata": {},
7
+ "source": [
8
+ "# Cortex Small-LLM Router GRPO training (Workstream B Phase 4 — SKELETON, post-pivot)\n",
9
+ "\n",
10
+ "**Status:** scaffold with explicit `TODO(conv-point)` markers at every Cortex-dependent integration site. **Do not run end-to-end yet** — the cells marked TODO will fail until Workstream A's Session 13 ships:\n",
11
+ "\n",
12
+ "- `cortex.metacognition.MetacognitionState` — featurization source.\n",
13
+ "- `cortex.routing_policy.RoutingPolicy` — trainable interface (Phase A §6).\n",
14
+ "- `cortex.council.Council` — deliberation orchestrator that drives rollouts.\n",
15
+ "- `baselines.cortex_fixed_router.B3CortexFixedRouter` — generates the deterministic-router corpus.\n",
16
+ "\n",
17
+ "**Architecture (post-pivot — replaces the MLP-head approach in commit `5489e55`):**\n",
18
+ "- Router is a small LLM: `unsloth/Qwen3-1.5B-Instruct` + LoRA rank 16. Roughly 3 GB on a100-large in 4-bit.\n",
19
+ "- Input: NL summary of `MetacognitionState` (~300 tokens).\n",
20
+ "- Output: structured JSON matching `cortex.schemas.RoutingAction` (~150 tokens).\n",
21
+ "- GRPO via the same Unsloth + TRL `GRPOTrainer` pipeline as the B1 notebook (one less code path).\n",
22
+ "- Reward: Phase-1-fixed `outer_reward` ∈ [-1, 1] composed with the token-budget penalty via `training.reward_shaping.shape_reward`.\n",
23
+ "\n",
24
+ "**Why a small LLM (not an MLP)?** With A100 compute available, a small LLM gives better reasoning over complex MetacognitionState, is interpretable in the pitch demo (you can read what the router thinks), and reuses the exact training infrastructure as the B1 baseline."
25
+ ]
26
+ },
27
+ {
28
+ "cell_type": "markdown",
29
+ "id": "acae54e37e7d407bbb7b55eff062a284",
30
+ "metadata": {},
31
+ "source": [
32
+ "## 1. Install dependencies\n",
33
+ "\n",
34
+ "Same Unsloth + vLLM + TRL stack as the B1 notebook; smaller GPU footprint because the trainable router is 1.5B."
35
+ ]
36
+ },
37
+ {
38
+ "cell_type": "code",
39
+ "execution_count": null,
40
+ "id": "9a63283cbaf04dbcab1f6479b197f3a8",
41
+ "metadata": {},
42
+ "outputs": [],
43
+ "source": [
44
+ "%%capture\n",
45
+ "!pip install --upgrade pip\n",
46
+ "!pip install unsloth vllm\n",
47
+ "!pip install --upgrade --no-deps \"trl>=0.14\" peft accelerate bitsandbytes\n",
48
+ "!pip install pydantic openenv huggingface_hub matplotlib"
49
+ ]
50
+ },
51
+ {
52
+ "cell_type": "markdown",
53
+ "id": "8dd0d8092fe74a7c96281538738b07e2",
54
+ "metadata": {},
55
+ "source": [
56
+ "## 2. Authenticate with Hugging Face"
57
+ ]
58
+ },
59
+ {
60
+ "cell_type": "code",
61
+ "execution_count": null,
62
+ "id": "72eea5119410473aa328ad9291626812",
63
+ "metadata": {},
64
+ "outputs": [],
65
+ "source": [
66
+ "import os\n",
67
+ "\n",
68
+ "try:\n",
69
+ " from google.colab import userdata\n",
70
+ "\n",
71
+ " HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
72
+ " os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
73
+ "except Exception:\n",
74
+ " from huggingface_hub import login\n",
75
+ "\n",
76
+ " login()\n",
77
+ " HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
78
+ "\n",
79
+ "assert HF_TOKEN, \"HF_TOKEN required\""
80
+ ]
81
+ },
82
+ {
83
+ "cell_type": "markdown",
84
+ "id": "8edb47106e1a46a883d545849b8ab81b",
85
+ "metadata": {},
86
+ "source": [
87
+ "## 3. Clone CrisisWorldCortex (post-Cortex-Session-13 deploy)\n",
88
+ "\n",
89
+ "**TODO(conv-point):** the target Space repo at convergence point will have `cortex/*` populated through Session 13 (subagents + lenses + brains + council + metacognition + routing_policy + B3). Until then, the Cortex-dependent imports below will fail."
90
+ ]
91
+ },
92
+ {
93
+ "cell_type": "code",
94
+ "execution_count": null,
95
+ "id": "10185d26023b46108eb7d9f57d49d2b3",
96
+ "metadata": {},
97
+ "outputs": [],
98
+ "source": [
99
+ "%%capture\n",
100
+ "!rm -rf /content/CrisisWorldCortex\n",
101
+ "!git clone https://huggingface.co/spaces/Angshuman28/CrisisWorldCortex /content/CrisisWorldCortex\n",
102
+ "%cd /content/CrisisWorldCortex\n",
103
+ "!pip install -e ."
104
+ ]
105
+ },
106
+ {
107
+ "cell_type": "code",
108
+ "execution_count": null,
109
+ "id": "8763a12b2bbd4a93a75aff182afb95dc",
110
+ "metadata": {},
111
+ "outputs": [],
112
+ "source": [
113
+ "import sys\n",
114
+ "\n",
115
+ "sys.path.insert(0, \"/content/CrisisWorldCortex\")\n",
116
+ "\n",
117
+ "\n",
118
+ "# TODO(conv-point): uncomment when Cortex Session 13 lands.\n",
119
+ "# from cortex.schemas import MetacognitionState, RoutingAction, RouterStep\n",
120
+ "# from cortex.routing_policy import RoutingPolicy\n",
121
+ "# from cortex.council import Council\n",
122
+ "# from baselines.cortex_fixed_router import B3CortexFixedRouter\n",
123
+ "print(\"Phase-2 training utilities OK; Cortex imports gated until Session 13\")"
124
+ ]
125
+ },
126
+ {
127
+ "cell_type": "markdown",
128
+ "id": "7623eae2785240b9bd12b16a66d81610",
129
+ "metadata": {},
130
+ "source": [
131
+ "## 4. Load Qwen3-1.5B-Instruct (router) with LoRA\n",
132
+ "\n",
133
+ "Small enough to fit alongside frozen 7B/8B brain LLMs on a100-large (80GB). LoRA rank 16 — tighter than the B1 notebook's 32 because the action space is structured JSON (small effective vocabulary)."
134
+ ]
135
+ },
136
+ {
137
+ "cell_type": "code",
138
+ "execution_count": null,
139
+ "id": "7cdc8c89c7104fffa095e18ddfef8986",
140
+ "metadata": {},
141
+ "outputs": [],
142
+ "source": [
143
+ "from unsloth import FastLanguageModel\n",
144
+ "\n",
145
+ "ROUTER_MODEL = \"unsloth/Qwen3-1.5B-Instruct-bnb-4bit\"\n",
146
+ "MAX_SEQ_LEN = 2048\n",
147
+ "\n",
148
+ "router_model, router_tokenizer = FastLanguageModel.from_pretrained(\n",
149
+ " model_name=ROUTER_MODEL,\n",
150
+ " max_seq_length=MAX_SEQ_LEN,\n",
151
+ " load_in_4bit=True,\n",
152
+ " fast_inference=True,\n",
153
+ " max_lora_rank=16,\n",
154
+ " gpu_memory_utilization=0.5, # share GPU with frozen brain LLMs at conv-point\n",
155
+ ")\n",
156
+ "\n",
157
+ "router_model = FastLanguageModel.get_peft_model(\n",
158
+ " router_model,\n",
159
+ " r=16,\n",
160
+ " target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
161
+ " lora_alpha=32,\n",
162
+ " use_gradient_checkpointing=\"unsloth\",\n",
163
+ " random_state=42,\n",
164
+ ")\n",
165
+ "print(\"Router (Qwen3-1.5B + LoRA r=16) loaded\")"
166
+ ]
167
+ },
168
+ {
169
+ "cell_type": "markdown",
170
+ "id": "b118ea5561624da68c537baed56e602f",
171
+ "metadata": {},
172
+ "source": [
173
+ "## 5. Featurization: MetacognitionState → NL prompt + RoutingAction schema\n",
174
+ "\n",
175
+ "Replaces the MLP version's `(24,) np.float32` featurization. The router consumes a natural-language summary; output is structured JSON validated against `cortex.schemas.RoutingAction`.\n",
176
+ "\n",
177
+ "**TODO(conv-point):** the schema below references `cortex.schemas.MetacognitionState` (Session 13).\n",
178
+ "Per Phase A `cortex/schemas.py` and `docs/CORTEX_ARCHITECTURE.md` §6, the NL summary covers the 11 documented fields + the phase string."
179
+ ]
180
+ },
181
+ {
182
+ "cell_type": "code",
183
+ "execution_count": null,
184
+ "id": "938c804e27f84196a10c8828c723f798",
185
+ "metadata": {},
186
+ "outputs": [],
187
+ "source": [
188
+ "ROUTER_SYSTEM_PROMPT = \"\"\"You are the Cortex router for CrisisWorldCortex. You receive a metacognition\n",
189
+ "state summary describing the current deliberation phase and emit ONE routing action as JSON.\n",
190
+ "\n",
191
+ "Allowed action kinds (per cortex/CLAUDE.md):\n",
192
+ "- call_subagent: invoke a brain's subagent. Required: brain (epidemiology|logistics|governance),\n",
193
+ " subagent (world_modeler|planner|critic).\n",
194
+ "- request_challenge: cross-brain critique. Required: challenger_brain, target_brain.\n",
195
+ "- switch_phase: advance the phase machine. Required: new_phase (divergence|challenge|narrowing|convergence).\n",
196
+ "- preserve_dissent: tag a minority recommendation. Required: tag (string, max 80 chars).\n",
197
+ "- emit_outer_action: close the tick with a final action.\n",
198
+ "- stop_and_no_op: close the tick with a no-op.\n",
199
+ "\n",
200
+ "Hard caps (binding): ≤2 deliberation rounds/tick, ≤1 cross-brain challenge/tick,\n",
201
+ "≤1 critic call per brain/tick, ≤6000 token budget per tick.\n",
202
+ "\n",
203
+ "Output exactly one JSON object — no markdown fences, no prose around it.\"\"\"\n",
204
+ "\n",
205
+ "PHASE_NAMES = (\"divergence\", \"challenge\", \"narrowing\", \"convergence\")\n",
206
+ "\n",
207
+ "\n",
208
+ "def metacog_state_to_prompt(state) -> str:\n",
209
+ " \"\"\"Convert MetacognitionState → NL summary for the router.\n",
210
+ "\n",
211
+ " TODO(conv-point): change ``state`` typing to MetacognitionState\n",
212
+ " (cortex.schemas) once Session 13 lands. Duck-typed for now.\n",
213
+ " \"\"\"\n",
214
+ " return (\n",
215
+ " f\"Tick {getattr(state, 'tick', 0)}, round {getattr(state, 'round', 1)}, \"\n",
216
+ " f\"phase={getattr(state, 'phase', 'divergence')}.\\n\"\n",
217
+ " f\"Inter-brain agreement: {getattr(state, 'inter_brain_agreement', 0.0):.2f}.\\n\"\n",
218
+ " f\"Average confidence: {getattr(state, 'average_confidence', 0.0):.2f}.\\n\"\n",
219
+ " f\"Average evidence support: {getattr(state, 'average_evidence_support', 0.0):.2f}.\\n\"\n",
220
+ " f\"Novelty yield (last round): {getattr(state, 'novelty_yield_last_round', 0.0):.2f}.\\n\"\n",
221
+ " f\"Collapse suspicion: {getattr(state, 'collapse_suspicion', 0.0):.2f}.\\n\"\n",
222
+ " f\"Budget remaining: {getattr(state, 'budget_remaining_frac', 1.0):.0%}.\\n\"\n",
223
+ " f\"Urgency: {getattr(state, 'urgency', 0.0):.2f}.\\n\"\n",
224
+ " f\"Preserved dissent count: {getattr(state, 'preserved_dissent_count', 0)}.\\n\"\n",
225
+ " f\"Cross-brain challenge used this tick: \"\n",
226
+ " f\"{bool(getattr(state, 'challenge_used_this_tick', 0))}.\\n\\n\"\n",
227
+ " f\"Choose the next routing action.\"\n",
228
+ " )\n",
229
+ "\n",
230
+ "\n",
231
+ "def make_router_chat_prompt(state) -> str:\n",
232
+ " return router_tokenizer.apply_chat_template(\n",
233
+ " [\n",
234
+ " {\"role\": \"system\", \"content\": ROUTER_SYSTEM_PROMPT},\n",
235
+ " {\"role\": \"user\", \"content\": metacog_state_to_prompt(state)},\n",
236
+ " ],\n",
237
+ " tokenize=False,\n",
238
+ " add_generation_prompt=True,\n",
239
+ " )"
240
+ ]
241
+ },
242
+ {
243
+ "cell_type": "markdown",
244
+ "id": "504fb2a444614c0babb325280ed9130a",
245
+ "metadata": {},
246
+ "source": [
247
+ "## 6. Connect to the deployed CrisisWorld env\n",
248
+ "\n",
249
+ "Same env client pattern as the B1 notebook."
250
+ ]
251
+ },
252
+ {
253
+ "cell_type": "code",
254
+ "execution_count": null,
255
+ "id": "59bbdb311c014d738909a11f9e486628",
256
+ "metadata": {},
257
+ "outputs": [],
258
+ "source": [
259
+ "from CrisisWorldCortex.client import CrisisworldcortexEnv\n",
260
+ "\n",
261
+ "ENV_URL = \"https://angshuman28-crisisworldcortex.hf.space\"\n",
262
+ "TASKS = (\"outbreak_easy\", \"outbreak_medium\", \"outbreak_hard\")\n",
263
+ "EPISODE_TICKS = 12\n",
264
+ "\n",
265
+ "\n",
266
+ "def make_env() -> CrisisworldcortexEnv:\n",
267
+ " return CrisisworldcortexEnv(base_url=ENV_URL)\n",
268
+ "\n",
269
+ "\n",
270
+ "_test_env = make_env()\n",
271
+ "_obs = _test_env.reset(task_name=\"outbreak_easy\", seed=0, max_ticks=EPISODE_TICKS)\n",
272
+ "print(f\"Env OK. Initial tick={_obs.tick}, regions={[r.region for r in _obs.regions]}\")"
273
+ ]
274
+ },
275
+ {
276
+ "cell_type": "markdown",
277
+ "id": "b43b363d81ae4b689946ece5c682cd59",
278
+ "metadata": {},
279
+ "source": [
280
+ "## 7. Build training-data prompt set from B3 deterministic-router trajectories\n",
281
+ "\n",
282
+ "**TODO(conv-point):** B3CortexFixedRouter runs ~50 episodes. Each `RouterStep`'s `MetacognitionState` becomes a router-prompt; the GRPO completion is the router's emitted JSON.\n",
283
+ "\n",
284
+ "Replaces the MLP version's action-vocab + B3-trajectory-to-tuple pipeline. Same RolloutBuffer, different content shape."
285
+ ]
286
+ },
287
+ {
288
+ "cell_type": "code",
289
+ "execution_count": null,
290
+ "id": "8a65eabff63a45729fe45fb5ade58bdc",
291
+ "metadata": {},
292
+ "outputs": [],
293
+ "source": [
294
+ "from datasets import Dataset\n",
295
+ "\n",
296
+ "TRAIN_EPISODES = 50\n",
297
+ "\n",
298
+ "\n",
299
+ "def collect_b3_router_prompts(num_episodes: int = TRAIN_EPISODES) -> Dataset:\n",
300
+ " \"\"\"Run B3 and convert each RouterStep into a (prompt, task, seed) row.\n",
301
+ "\n",
302
+ " TODO(conv-point): uncomment the body when B3CortexFixedRouter ships.\n",
303
+ " \"\"\"\n",
304
+ " rows = {\"prompt\": [], \"task\": [], \"seed\": []}\n",
305
+ " # TODO(conv-point):\n",
306
+ " # b3 = B3CortexFixedRouter(env=make_env())\n",
307
+ " # for ep in range(num_episodes):\n",
308
+ " # trajectory = b3.run_episode(task=\"outbreak_easy\", seed=ep)\n",
309
+ " # for router_step in trajectory.router_steps:\n",
310
+ " # rows[\"prompt\"].append(make_router_chat_prompt(router_step.metacognition_state))\n",
311
+ " # rows[\"task\"].append(\"outbreak_easy\")\n",
312
+ " # rows[\"seed\"].append(ep)\n",
313
+ " return (\n",
314
+ " Dataset.from_dict(rows)\n",
315
+ " if rows[\"prompt\"]\n",
316
+ " else Dataset.from_dict(\n",
317
+ " {\"prompt\": [\"placeholder until conv-point\"], \"task\": [\"outbreak_easy\"], \"seed\": [0]}\n",
318
+ " )\n",
319
+ " )\n",
320
+ "\n",
321
+ "\n",
322
+ "train_dataset = collect_b3_router_prompts()\n",
323
+ "print(f\"Dataset: {len(train_dataset)} rows (will be ~{TRAIN_EPISODES * 8} at conv-point)\")"
324
+ ]
325
+ },
326
+ {
327
+ "cell_type": "markdown",
328
+ "id": "c3933fab20d04ec698c2621248eb3be0",
329
+ "metadata": {},
330
+ "source": [
331
+ "## 8. Reward function: full-episode rollout per (prompt, completion)\n",
332
+ "\n",
333
+ "Uses Phase-1 `outer_reward` summed across the episode, with Phase-2 `shape_reward` token-budget penalty.\n",
334
+ "\n",
335
+ "**TODO(conv-point):** the rollout loop calls `Council.step` with the trainable router policy. Until Session 13 lands, the loop is stubbed and returns 0.0 for every (prompt, completion)."
336
+ ]
337
+ },
338
+ {
339
+ "cell_type": "code",
340
+ "execution_count": null,
341
+ "id": "4dd4641cc4064e0191573fe9c69df29b",
342
+ "metadata": {},
343
+ "outputs": [],
344
+ "source": [
345
+ "def cortex_router_reward(prompts, completions, task, seed, **_kwargs):\n",
346
+ " \"\"\"Reward = sum(per-tick obs.reward) over a full episode driven by the\n",
347
+ " trainable router's emitted RoutingAction JSON.\n",
348
+ "\n",
349
+ " Each (prompt, completion) pair represents ONE router decision; the\n",
350
+ " full episode reward is shared across all router decisions in that\n",
351
+ " episode (GRPO group-relative advantage handles credit assignment).\n",
352
+ "\n",
353
+ " TODO(conv-point): replace the stub body with the real Council-driven\n",
354
+ " rollout once Session 13 ships.\n",
355
+ " \"\"\"\n",
356
+ " rewards = []\n",
357
+ " for completion, t, s in zip(completions, task, seed):\n",
358
+ " # TODO(conv-point):\n",
359
+ " # try:\n",
360
+ " # routing_action = RoutingAction.model_validate_json(completion)\n",
361
+ " # except ValidationError:\n",
362
+ " # rewards.append(-1.0) # invalid JSON → terminal-equivalent penalty\n",
363
+ " # continue\n",
364
+ " # council = Council(routing_policy=trainable_router_from(routing_action),\n",
365
+ " # env=make_env(), brains=cortex_brains)\n",
366
+ " # episode_return = council.run_episode(task=t, seed=int(s)).total_reward\n",
367
+ " # rewards.append(episode_return)\n",
368
+ " rewards.append(0.0) # stub\n",
369
+ " return rewards"
370
+ ]
371
+ },
372
+ {
373
+ "cell_type": "markdown",
374
+ "id": "8309879909854d7188b41380fd92a7c3",
375
+ "metadata": {},
376
+ "source": [
377
+ "## 9. GRPO training\n",
378
+ "\n",
379
+ "Same TRL `GRPOTrainer` shape as the B1 notebook. ~300 steps × group size 4 = ~1200 router decisions × full-episode rollouts. Wall-clock ~1.5 hours on a100-large at convergence point (dominated by brain-LLM rollouts, not router fine-tuning)."
380
+ ]
381
+ },
382
+ {
383
+ "cell_type": "code",
384
+ "execution_count": null,
385
+ "id": "3ed186c9a28b402fb0bc4494df01f08d",
386
+ "metadata": {},
387
+ "outputs": [],
388
+ "source": [
389
+ "from trl import GRPOConfig, GRPOTrainer\n",
390
+ "\n",
391
+ "MAX_TRAIN_STEPS = 300\n",
392
+ "GROUP_SIZE = 4\n",
393
+ "MAX_PROMPT_LEN = 512\n",
394
+ "MAX_COMPLETION_LEN = 256 # router output is structured JSON (M-FR-11)\n",
395
+ "\n",
396
+ "training_args = GRPOConfig(\n",
397
+ " output_dir=\"/content/cortex_router_grpo_output\",\n",
398
+ " learning_rate=5e-6,\n",
399
+ " per_device_train_batch_size=GROUP_SIZE,\n",
400
+ " gradient_accumulation_steps=1,\n",
401
+ " num_generations=GROUP_SIZE,\n",
402
+ " max_prompt_length=MAX_PROMPT_LEN,\n",
403
+ " max_completion_length=MAX_COMPLETION_LEN,\n",
404
+ " max_steps=MAX_TRAIN_STEPS,\n",
405
+ " save_steps=100,\n",
406
+ " logging_steps=5,\n",
407
+ " report_to=\"none\",\n",
408
+ " bf16=True,\n",
409
+ " optim=\"adamw_8bit\",\n",
410
+ " temperature=0.8,\n",
411
+ " use_vllm=True,\n",
412
+ " vllm_mode=\"colocate\",\n",
413
+ " seed=42,\n",
414
+ ")\n",
415
+ "\n",
416
+ "trainer = GRPOTrainer(\n",
417
+ " model=router_model,\n",
418
+ " processing_class=router_tokenizer,\n",
419
+ " reward_funcs=[cortex_router_reward],\n",
420
+ " args=training_args,\n",
421
+ " train_dataset=train_dataset,\n",
422
+ ")\n",
423
+ "print(\"Router GRPOTrainer constructed.\")\n",
424
+ "print(\"# TODO(conv-point): uncomment trainer.train() once Cortex Session 13 lands.\")\n",
425
+ "# trainer.train()"
426
+ ]
427
+ },
428
+ {
429
+ "cell_type": "markdown",
430
+ "id": "cb1e1581032b452c9409d6c6813c49d1",
431
+ "metadata": {},
432
+ "source": [
433
+ "## 10. Save the trained router LoRA to HF Hub"
434
+ ]
435
+ },
436
+ {
437
+ "cell_type": "code",
438
+ "execution_count": null,
439
+ "id": "379cbbc1e968416e875cc15c1202d7eb",
440
+ "metadata": {},
441
+ "outputs": [],
442
+ "source": [
443
+ "from huggingface_hub import HfApi\n",
444
+ "\n",
445
+ "HUB_REPO = \"Angshuman28/crisisworld-cortex-router-llm\"\n",
446
+ "\n",
447
+ "router_model.save_pretrained(\"/content/cortex_router_lora\")\n",
448
+ "router_tokenizer.save_pretrained(\"/content/cortex_router_lora\")\n",
449
+ "\n",
450
+ "api = HfApi()\n",
451
+ "api.create_repo(HUB_REPO, exist_ok=True, repo_type=\"model\", private=False, token=HF_TOKEN)\n",
452
+ "api.upload_folder(\n",
453
+ " folder_path=\"/content/cortex_router_lora\",\n",
454
+ " repo_id=HUB_REPO,\n",
455
+ " repo_type=\"model\",\n",
456
+ " token=HF_TOKEN,\n",
457
+ ")\n",
458
+ "print(f\"Saved to https://huggingface.co/{HUB_REPO}\")"
459
+ ]
460
+ },
461
+ {
462
+ "cell_type": "markdown",
463
+ "id": "277c27b1587741f2af2001be3712ef0d",
464
+ "metadata": {},
465
+ "source": [
466
+ "## 11. Eval: B6 (trained LLM router) vs B3 (deterministic) on 3 tasks\n",
467
+ "\n",
468
+ "**TODO(conv-point):** the comparison loop below requires both B3CortexFixedRouter (deterministic) and a way to swap the trainable router into Council.routing_policy. The eval is the headline result for the convergence point — its sign decides whether B6 ships or B3 ships per the Phase 7 hard-exit gate.\n",
469
+ "\n",
470
+ "Decision rule (per spec): if reward over training steps is INCREASING, ship B6. Else ship B3."
471
+ ]
472
+ },
473
+ {
474
+ "cell_type": "code",
475
+ "execution_count": null,
476
+ "id": "db7b79bc585a40fcaf58bf750017e135",
477
+ "metadata": {},
478
+ "outputs": [],
479
+ "source": [
480
+ "# TODO(conv-point):\n",
481
+ "# from cortex.routing_policy import DeterministicRouter, TrainableRouter\n",
482
+ "# from cortex.council import Council\n",
483
+ "# from baselines.cortex_fixed_router import B3CortexFixedRouter\n",
484
+ "#\n",
485
+ "# def run_b6_episode(task, seed):\n",
486
+ "# trainable = TrainableRouter(\n",
487
+ "# model=router_model,\n",
488
+ "# tokenizer=router_tokenizer,\n",
489
+ "# system_prompt=ROUTER_SYSTEM_PROMPT,\n",
490
+ "# featurize=metacog_state_to_prompt,\n",
491
+ "# )\n",
492
+ "# council = Council(routing_policy=trainable, env=make_env(), brains=cortex_brains)\n",
493
+ "# return council.run_episode(task=task, seed=seed).total_reward\n",
494
+ "#\n",
495
+ "# def run_b3_episode(task, seed):\n",
496
+ "# b3 = B3CortexFixedRouter(env=make_env())\n",
497
+ "# return b3.run_episode(task=task, seed=seed).total_reward\n",
498
+ "#\n",
499
+ "# b6_results = {t: run_b6_episode(t, 0) for t in TASKS}\n",
500
+ "# b3_results = {t: run_b3_episode(t, 0) for t in TASKS}\n",
501
+ "# print(f\"B6 (trained LLM router): {b6_results}\")\n",
502
+ "# print(f\"B3 (deterministic): {b3_results}\")\n",
503
+ "\n",
504
+ "print(\"Eval cell skeleton — uncomment when Cortex Session 13 ships\")"
505
+ ]
506
+ },
507
+ {
508
+ "cell_type": "markdown",
509
+ "id": "916684f9a58a4a2aa5f864670399430d",
510
+ "metadata": {},
511
+ "source": [
512
+ "## 12. Plot training reward curve\n",
513
+ "\n",
514
+ "Phase 7 hard-exit gate watches this curve. If reward is increasing across training steps → ship B6. Else → ship B3."
515
+ ]
516
+ },
517
+ {
518
+ "cell_type": "code",
519
+ "execution_count": null,
520
+ "id": "1671c31a24314836a5b85d7ef7fbf015",
521
+ "metadata": {},
522
+ "outputs": [],
523
+ "source": [
524
+ "import matplotlib.pyplot as plt\n",
525
+ "\n",
526
+ "log_path = \"/content/cortex_router_grpo_output/trainer_state.json\"\n",
527
+ "if os.path.exists(log_path):\n",
528
+ " import json as _json\n",
529
+ "\n",
530
+ " with open(log_path) as fh:\n",
531
+ " state = _json.load(fh)\n",
532
+ " history = state.get(\"log_history\", [])\n",
533
+ " rewards = [entry[\"reward\"] for entry in history if \"reward\" in entry]\n",
534
+ " if rewards:\n",
535
+ " fig, ax = plt.subplots(figsize=(8, 4))\n",
536
+ " ax.plot(rewards)\n",
537
+ " ax.set_xlabel(\"GRPO step\")\n",
538
+ " ax.set_ylabel(\"Mean reward\")\n",
539
+ " ax.set_title(\"Cortex small-LLM router — training reward\")\n",
540
+ " plt.tight_layout()\n",
541
+ " plt.show()\n",
542
+ " else:\n",
543
+ " print(\"No reward entries yet — uncomment trainer.train() at conv-point\")\n",
544
+ "else:\n",
545
+ " print(\"No trainer state yet — uncomment trainer.train() at conv-point\")"
546
+ ]
547
+ }
548
+ ],
549
+ "metadata": {
550
+ "accelerator": "GPU",
551
+ "colab": {
552
+ "gpuType": "A100",
553
+ "provenance": []
554
+ },
555
+ "kernelspec": {
556
+ "display_name": "Python 3",
557
+ "language": "python",
558
+ "name": "python3"
559
+ },
560
+ "language_info": {
561
+ "name": "python",
562
+ "version": "3.10"
563
+ }
564
+ },
565
+ "nbformat": 4,
566
+ "nbformat_minor": 5
567
+ }
openenv.yaml CHANGED
@@ -1,7 +1,7 @@
1
- spec_version: 1
2
- name: CrisisWorldCortex
3
- type: space
4
- runtime: fastapi
5
- app: server.app:app
6
- port: 8000
7
-
 
1
+ spec_version: 1
2
+ name: CrisisWorldCortex
3
+ type: space
4
+ runtime: fastapi
5
+ app: server.app:app
6
+ port: 8000
7
+
openenv_CrisisWorldCortex.egg-info/PKG-INFO ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Metadata-Version: 2.4
2
+ Name: openenv-CrisisWorldCortex
3
+ Version: 0.1.0
4
+ Summary: Crisisworldcortex environment for OpenEnv
5
+ Requires-Python: >=3.10
6
+ Requires-Dist: openenv-core[core]==0.2.3
7
+ Requires-Dist: openai<3.0,>=2.0
8
+ Requires-Dist: python-dotenv>=1.0.0
9
+ Provides-Extra: dev
10
+ Requires-Dist: pre-commit>=4.0.0; extra == "dev"
11
+ Requires-Dist: pytest>=8.0.0; extra == "dev"
12
+ Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
13
+ Requires-Dist: ruff>=0.8.0; extra == "dev"
openenv_CrisisWorldCortex.egg-info/SOURCES.txt ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ README.md
2
+ __init__.py
3
+ client.py
4
+ inference.py
5
+ models.py
6
+ pyproject.toml
7
+ ./__init__.py
8
+ ./client.py
9
+ ./inference.py
10
+ ./models.py
11
+ openenv_CrisisWorldCortex.egg-info/PKG-INFO
12
+ openenv_CrisisWorldCortex.egg-info/SOURCES.txt
13
+ openenv_CrisisWorldCortex.egg-info/dependency_links.txt
14
+ openenv_CrisisWorldCortex.egg-info/entry_points.txt
15
+ openenv_CrisisWorldCortex.egg-info/requires.txt
16
+ openenv_CrisisWorldCortex.egg-info/top_level.txt
17
+ server/CrisisWorldCortex_environment.py
18
+ server/__init__.py
19
+ server/app.py
20
+ tests/test_actions_round_trip.py
21
+ tests/test_baseline_b1.py
22
+ tests/test_baseline_b2.py
23
+ tests/test_cortex_brain_executive.py
24
+ tests/test_cortex_brain_smoke.py
25
+ tests/test_cortex_lenses.py
26
+ tests/test_cortex_perception.py
27
+ tests/test_cortex_subagents.py
28
+ tests/test_env_reset_kwargs.py
29
+ tests/test_env_step_reward_wiring.py
30
+ tests/test_import_graph.py
31
+ tests/test_legal_constraint_enforcement.py
32
+ tests/test_llm_client.py
33
+ tests/test_observation_no_latent_leak.py
34
+ tests/test_outer_reward_in_range.py
35
+ tests/test_outer_reward_non_constancy.py
36
+ tests/test_outer_reward_terminal_bonus.py
37
+ tests/test_package_exports.py
38
+ tests/test_reward_signal_quality.py
39
+ tests/test_schemas_roundtrip.py
40
+ tests/test_simulator_determinism.py
41
+ tests/test_simulator_random_episode.py
42
+ tests/test_simulator_task_configs.py
43
+ tests/test_smoke_env.py
44
+ tests/test_stdout_format.py
45
+ tests/test_synthetic_rejection_payload.py
46
+ tests/test_training_eval_metrics.py
47
+ tests/test_training_reward_shaping.py
48
+ tests/test_training_rollout_buffer.py
openenv_CrisisWorldCortex.egg-info/dependency_links.txt ADDED
@@ -0,0 +1 @@
 
 
1
+
openenv_CrisisWorldCortex.egg-info/entry_points.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ [console_scripts]
2
+ server = CrisisWorldCortex.server.app:main
openenv_CrisisWorldCortex.egg-info/requires.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ openenv-core[core]==0.2.3
2
+ openai<3.0,>=2.0
3
+ python-dotenv>=1.0.0
4
+
5
+ [dev]
6
+ pre-commit>=4.0.0
7
+ pytest>=8.0.0
8
+ pytest-cov>=4.0.0
9
+ ruff>=0.8.0
openenv_CrisisWorldCortex.egg-info/top_level.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ CrisisWorldCortex
pyproject.toml CHANGED
@@ -25,6 +25,7 @@ dependencies = [
25
  # future uv sync from silently pulling 3.x. Bump explicitly when 3.0
26
  # ships and we've verified compatibility.
27
  "openai>=2.0,<3.0",
 
28
  ]
29
 
30
  [project.optional-dependencies]
 
25
  # future uv sync from silently pulling 3.x. Bump explicitly when 3.0
26
  # ships and we've verified compatibility.
27
  "openai>=2.0,<3.0",
28
+ "python-dotenv>=1.0.0"
29
  ]
30
 
31
  [project.optional-dependencies]
server/CLAUDE.md CHANGED
@@ -15,33 +15,24 @@ Cortex code (brains, subagents, router, metacognition), LLM clients, training lo
15
 
16
  ## Import rule (binding)
17
 
18
- **Inside `server/`, use package-relative imports for server internals**:
19
- `from .simulator import ...` from one-level modules and
20
- `from ..simulator import ...` from nested modules. This works when the
21
- server is loaded as either `CrisisWorldCortex.server.*` (`uv run server`)
22
- or top-level `server.*` (Docker / `uvicorn server.app:app`). Never use
23
- `from CrisisWorldCortex.server...` for server-internal imports.
24
 
25
- For `models`, use the canonical package path in every server module that needs it:
26
 
27
  ```python
28
- from CrisisWorldCortex.models import CrisisworldcortexAction, CrisisworldcortexObservation
 
 
 
29
  ```
30
 
31
- The server runs under multiple import contexts (`uv run server`,
32
- `uvicorn server.app:app`, Docker `cd /app/env && uvicorn server.app:app`).
33
- Canonical wire imports keep Pydantic model identity stable across those modes.
34
 
35
- **Wire-type imports from deep modules (binding)**: two-or-more-levels-deep
36
- files (`server/simulator/*`, `server/graders/*`) must still use
37
- `from CrisisWorldCortex.models import …`; `..models` from those depths
38
- resolves to a non-existent `CrisisWorldCortex.server.models`, and a bare
39
- fallback loads a second `models` module.
40
 
41
  ## Allowed imports
42
 
43
- `CrisisWorldCortex.models`, package-relative `server/simulator/*` and
44
- `server/graders/*`, `openenv.core.*`, stdlib, FastAPI, Pydantic, numpy.
45
 
46
  ## Forbidden imports
47
 
 
15
 
16
  ## Import rule (binding)
17
 
18
+ **Inside `server/`, use `from server.simulator import …`, never `from CrisisWorldCortex.server.simulator import …`.** Same rule for `server.graders`.
 
 
 
 
 
19
 
20
+ For `models`, use the dual-import fallback in every new server module that needs it:
21
 
22
  ```python
23
+ try:
24
+ from ..models import CrisisworldcortexAction, CrisisworldcortexObservation
25
+ except (ImportError, ModuleNotFoundError):
26
+ from models import CrisisworldcortexAction, CrisisworldcortexObservation
27
  ```
28
 
29
+ The server runs under 3 import contexts (`python -m server.app`, `uvicorn server.app:app`, Docker `cd /app/env && uvicorn server.app:app`); the fallback is load-bearing.
 
 
30
 
31
+ **Wire-type imports from deep modules (binding)**: the dual-import fallback above only works for files **one level** inside `server/` (`server/CrisisWorldCortex_environment.py`, `server/app.py`). Two-or-more-levels-deep files (`server/simulator/*`, `server/graders/*`) **must** use `from CrisisWorldCortex.models import …` directly — `..models` from those depths resolves to a non-existent `CrisisWorldCortex.server.models`, the fallback fires, and bare `models` loads as a separate `sys.modules` entry, breaking Pydantic discriminator validation against types imported via the canonical path. Session 5a's `server/simulator/seir_model.py` and `server/simulator/tasks.py` document this with inline comments.
 
 
 
 
32
 
33
  ## Allowed imports
34
 
35
+ `models` (via dual-import fallback), `openenv.core.*`, `server/simulator/*`, `server/graders/*`, stdlib, FastAPI, Pydantic, numpy.
 
36
 
37
  ## Forbidden imports
38
 
server/CrisisWorldCortex_environment.py CHANGED
@@ -121,10 +121,26 @@ class CrisisworldcortexEnvironment(Environment):
121
  )
122
  self._state.step_count += 1
123
  self._world_state = apply_tick(self._world_state, action.action)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
  obs = make_observation(self._world_state)
125
- # Per design §15: r_outer is the only env-side reward signal, in [0,1].
126
- # Terminal bonus (+/-0.20) is composed downstream by the trainer per
127
- # design §14.3 — never bundled into obs.reward.
 
128
  obs.reward = outer_reward(self._world_state, action.action)
129
  return obs
130
 
 
121
  )
122
  self._state.step_count += 1
123
  self._world_state = apply_tick(self._world_state, action.action)
124
+ # Parse-failure terminal contract (design §19, Phase-1 restoration):
125
+ # the synthetic parse_failure_marker (PublicCommunication with
126
+ # honesty=0.0, magic-string discriminator per Phase-A M3-B) ends
127
+ # the episode as state.terminal = "failure". apply_tick may have
128
+ # set terminal to "none"/"success"/"timeout" via the SEIR rules;
129
+ # we override here because parse-failure is a harness-level fault,
130
+ # not a simulator-level event.
131
+ payload = action.action
132
+ if (
133
+ payload.kind == "public_communication"
134
+ and getattr(payload, "honesty", None) == 0.0
135
+ and self._world_state.recent_action_log
136
+ and not self._world_state.recent_action_log[-1].accepted
137
+ ):
138
+ self._world_state.terminal = "failure"
139
  obs = make_observation(self._world_state)
140
+ # Per design §15: r_outer is the only env-side reward signal, in
141
+ # [-1.0, 1.0] post-Phase-1 (was [0, 1]). Terminal bonus (+/-0.20)
142
+ # is composed downstream by the trainer per design §14.3 — never
143
+ # bundled into obs.reward.
144
  obs.reward = outer_reward(self._world_state, action.action)
145
  return obs
146
 
server/__init__.py CHANGED
@@ -1,11 +1,11 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
- #
4
- # This source code is licensed under the BSD-style license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- """Crisisworldcortex environment server components."""
8
-
9
- from .CrisisWorldCortex_environment import CrisisworldcortexEnvironment
10
-
11
- __all__ = ["CrisisworldcortexEnvironment"]
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Crisisworldcortex environment server components."""
8
+
9
+ from .CrisisWorldCortex_environment import CrisisworldcortexEnvironment
10
+
11
+ __all__ = ["CrisisworldcortexEnvironment"]
server/app.py CHANGED
@@ -29,8 +29,8 @@ Usage:
29
  """
30
 
31
  try:
32
- # from openenv.core.env_server.http_server import create_app
33
- from openenv.core.env_server import create_web_interface_app as create_app
34
  except Exception as e: # pragma: no cover
35
  raise ImportError(
36
  "openenv is required for the web interface. Install dependencies with '\n uv sync\n'"
 
29
  """
30
 
31
  try:
32
+ from openenv.core.env_server.http_server import create_app
33
+ # from openenv.core.env_server import create_web_interface_app as create_app
34
  except Exception as e: # pragma: no cover
35
  raise ImportError(
36
  "openenv is required for the web interface. Install dependencies with '\n uv sync\n'"
server/graders/outer_reward.py CHANGED
@@ -5,16 +5,29 @@
5
  # LICENSE file in the root directory of this source tree.
6
 
7
  """
8
- Outer reward grader for CrisisWorld (design §15).
9
 
10
  Public API (re-exported via ``server/graders/__init__.py``):
11
  - ``outer_reward(state, action) -> float`` — 6-component weighted score in
12
- ``[0.0, 1.0]``, computed on post-``apply_tick`` state. The only env-side
13
- reward signal per ``server/CLAUDE.md`` (binding contract).
 
 
 
14
  - ``terminal_bonus(state) -> float`` — episode-end ±0.20 / 0.0 bonus,
15
  composed by trainer in ``training/reward_shaping.py`` per design §14.3
16
- (``episode_return = Σ_t r_outer + terminal_bonus``). Kept separate so the
17
- per-tick ``r_outer`` stays inside ``[0.0, 1.0]``.
 
 
 
 
 
 
 
 
 
 
18
 
19
  Wire-protocol imports use the absolute path ``CrisisWorldCortex.models``
20
  because this file lives two levels deep inside ``server/`` — see
@@ -38,25 +51,42 @@ from ..simulator import (
38
  )
39
 
40
  # ============================================================================
41
- # Component weights (design §15; sum to 1.00)
42
  # ============================================================================
43
 
44
- W_INFECT = 0.35
45
- W_TIME = 0.18
46
- W_HOSP = 0.17
47
- W_CASC = 0.15
48
- W_POLICY = 0.12
49
- W_FAIR = 0.03
 
 
 
 
 
 
 
 
 
 
 
 
50
 
51
  # Terminal-bonus magnitudes (design §14.3 / §15).
52
  TERMINAL_BONUS_SUCCESS = 0.20
53
  TERMINAL_BONUS_FAILURE = -0.20
54
 
55
- # r_casc threshold: any region above this contributes a "hot" count;
56
- # matches the catastrophic-region threshold from design §6.4 so the
57
- # cascade-control signal aligns with the failure terminal.
58
  R_CASC_HOT_THRESHOLD = CATASTROPHIC_INFECTION_THRESHOLD
59
 
 
 
 
 
 
60
 
61
  def _clamp01(x: float) -> float:
62
  if x < 0.0:
@@ -76,34 +106,56 @@ def _hospital_load(region_I: float) -> float:
76
  return _clamp01(region_I * HOSPITALIZATION_FRACTION_OF_I / HOSPITAL_CAPACITY_FRACTION)
77
 
78
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79
  def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
80
- """Compute per-tick outer reward in ``[0.0, 1.0]``.
81
 
82
  Read post-``apply_tick`` state: ``state.regions[*].I`` is the
83
  just-stepped ground-truth infection fraction. ``recent_action_log[-1]``
84
  holds the just-dispatched action's acceptance flag.
85
 
86
- Six components (design §15):
87
- r_infect = 1 - mean(I) # weight 0.35
88
- r_time = 1 - tick / max_ticks # weight 0.18
89
- r_hosp = 1 - mean(hospital_load) # weight 0.17
90
- r_casc = 1 if no region exceeds 0.30 else 0 # weight 0.15
91
- r_policy = 1 if last action accepted else 0 # weight 0.12
92
- r_fair = 1 - pstdev(I) # weight 0.03
93
 
94
  The ``action`` argument is the action just dispatched. We read its
95
  acceptance flag from ``state.recent_action_log[-1]`` rather than
96
  re-dispatching — the simulator already recorded it, and re-dispatch
97
- would mutate state. ``action`` is kept in the signature for design
98
- parity with ``training_reward(trajectory)`` and to support future
99
- action-shape-dependent shaping without a signature break.
100
  """
101
  if not state.regions:
102
  return 0.0
103
 
104
  I_values = [r.I for r in state.regions]
 
105
 
106
- r_infect = _clamp01(1.0 - sum(I_values) / len(I_values))
 
107
 
108
  if state.max_ticks > 0:
109
  r_time = _clamp01(1.0 - state.tick / state.max_ticks)
@@ -111,16 +163,22 @@ def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
111
  r_time = 0.0
112
 
113
  hosp_loads = [_hospital_load(I) for I in I_values]
114
- r_hosp = _clamp01(1.0 - sum(hosp_loads) / len(hosp_loads))
 
115
 
116
- hot_regions = sum(1 for I in I_values if I > R_CASC_HOT_THRESHOLD)
117
- r_casc = 1.0 if hot_regions == 0 else 0.0
 
 
 
 
118
 
 
119
  if state.recent_action_log:
120
  last_entry = state.recent_action_log[-1]
121
- r_policy = 1.0 if last_entry.accepted else 0.0
122
  else:
123
- r_policy = 1.0 # No action dispatched yet → no penalty.
124
 
125
  if len(I_values) >= 2:
126
  r_fair = _clamp01(1.0 - statistics.pstdev(I_values))
@@ -135,8 +193,12 @@ def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
135
  + W_POLICY * r_policy
136
  + W_FAIR * r_fair
137
  )
138
- # Final clamp guards against floating-point drift past 1.0.
139
- return _clamp01(score)
 
 
 
 
140
 
141
 
142
  def terminal_bonus(state: WorldState) -> float:
 
5
  # LICENSE file in the root directory of this source tree.
6
 
7
  """
8
+ Outer reward grader for CrisisWorld (design §15 + §19, post-Phase-1 fix).
9
 
10
  Public API (re-exported via ``server/graders/__init__.py``):
11
  - ``outer_reward(state, action) -> float`` — 6-component weighted score in
12
+ ``[-1.0, 1.0]`` (Workstream-B Phase-1 range relaxation; was
13
+ ``[0.0, 1.0]``). Negative values come from ``r_policy = -0.5`` on rejected
14
+ actions and ``r_policy = -1.0`` on parse-failure markers (per design §19).
15
+ Computed on post-``apply_tick`` state. Still the only env-side reward
16
+ signal per ``server/CLAUDE.md``.
17
  - ``terminal_bonus(state) -> float`` — episode-end ±0.20 / 0.0 bonus,
18
  composed by trainer in ``training/reward_shaping.py`` per design §14.3
19
+ (``episode_return = Σ_t r_outer + terminal_bonus``). Kept separate from
20
+ the per-tick scalar.
21
+
22
+ Phase-1 changes (Workstream B):
23
+ - Steeper sensitivity on ``r_infect`` (``× 20``) and ``r_hosp`` (``× 10``)
24
+ so the gentle outbreak_easy task still produces gradient.
25
+ - Continuous ``r_casc`` (``1 - max(I)/0.30`` clamped) replaces binary.
26
+ - ``r_policy`` ∈ {-1.0 (parse-failure), -0.5 (rejected), 0.0 (accepted
27
+ no_op), +1.0 (accepted real action)} restoring §19 magnitudes.
28
+ - Weight redistribution: W_POLICY 0.12 → 0.35 (signal-driver); W_TIME
29
+ 0.18 → 0.05 (action-independent noise); other components rebalanced.
30
+ - Final ``[0,1]`` clamp dropped (Phase-A M2-A).
31
 
32
  Wire-protocol imports use the absolute path ``CrisisWorldCortex.models``
33
  because this file lives two levels deep inside ``server/`` — see
 
51
  )
52
 
53
  # ============================================================================
54
+ # Component weights (design §15 + Phase-1 redistribution; sum to 1.00)
55
  # ============================================================================
56
 
57
+ W_INFECT = 0.15 # was 0.35; iter-1 reduced because outbreak_easy keeps mean(I) tiny
58
+ W_TIME = 0.05 # was 0.18; action-independent, no signal value
59
+ W_HOSP = 0.10 # was 0.17; iter-1 reduction (gentle env keeps hosp_load low)
60
+ W_CASC = 0.10 # was 0.15; iter-1 reduction
61
+ W_POLICY = 0.55 # was 0.12; iter-1 dominant signal driver — accepted-real vs no_op vs rejected
62
+ W_FAIR = 0.05 # was 0.03; tiny boost
63
+
64
+ # Steepness coefficients (Phase-A M6, ONE-iteration tentative).
65
+ # r_infect ≈ 0 when mean(I) >= 0.05; near 1 when mean(I) <= 0.0 → strong gradient.
66
+ # r_hosp similarly sensitive to mean hospital_load.
67
+ R_INFECT_STEEPNESS = 20.0
68
+ R_HOSP_STEEPNESS = 10.0
69
+
70
+ # r_policy values per design §19 (Phase-1 restoration).
71
+ R_POLICY_PARSE_FAILURE = -1.0 # synthetic parse-failure marker
72
+ R_POLICY_REJECTED = -0.5 # well-formed-illegal (V2 / legal-violation)
73
+ R_POLICY_NOOP_ACCEPTED = 0.0 # accepted no-op (legal but inactive)
74
+ R_POLICY_REAL_ACCEPTED = 1.0 # accepted real intervention
75
 
76
  # Terminal-bonus magnitudes (design §14.3 / §15).
77
  TERMINAL_BONUS_SUCCESS = 0.20
78
  TERMINAL_BONUS_FAILURE = -0.20
79
 
80
+ # r_casc threshold: at max(I) >= this, r_casc = 0 (catastrophe imminent).
81
+ # Matches design §6.4's catastrophic-infection threshold so cascade
82
+ # signal aligns with the failure terminal.
83
  R_CASC_HOT_THRESHOLD = CATASTROPHIC_INFECTION_THRESHOLD
84
 
85
+ # Magic-string discriminator for parse-failure marker (Phase-A M3-B):
86
+ # baselines.flat_agent.parse_failure_marker emits PublicCommunication with
87
+ # honesty=0.0; intentional V2 attempts use honesty > 0.0.
88
+ PARSE_FAILURE_HONESTY_SENTINEL = 0.0
89
+
90
 
91
  def _clamp01(x: float) -> float:
92
  if x < 0.0:
 
106
  return _clamp01(region_I * HOSPITALIZATION_FRACTION_OF_I / HOSPITAL_CAPACITY_FRACTION)
107
 
108
 
109
+ def _r_policy_value(action: OuterActionPayload, accepted: bool) -> float:
110
+ """Compute ``r_policy`` per design §19 four-state contract.
111
+
112
+ Returns one of {-1.0, -0.5, 0.0, +1.0} based on (action.kind, accepted).
113
+ Parse-failure detection uses the ``honesty == 0.0`` sentinel on a
114
+ rejected ``PublicCommunication`` payload (Phase-A M3-B magic string).
115
+ """
116
+ if not accepted:
117
+ # Rejected branch. Distinguish parse-failure marker from intentional
118
+ # V2-PublicCommunication / legal-violation rejection.
119
+ if (
120
+ action.kind == "public_communication"
121
+ and getattr(action, "honesty", None) == PARSE_FAILURE_HONESTY_SENTINEL
122
+ ):
123
+ return R_POLICY_PARSE_FAILURE
124
+ return R_POLICY_REJECTED
125
+ # Accepted branch.
126
+ if action.kind == "no_op":
127
+ return R_POLICY_NOOP_ACCEPTED
128
+ return R_POLICY_REAL_ACCEPTED
129
+
130
+
131
  def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
132
+ """Compute per-tick outer reward in ``[-1.0, 1.0]`` (post-Phase-1 range).
133
 
134
  Read post-``apply_tick`` state: ``state.regions[*].I`` is the
135
  just-stepped ground-truth infection fraction. ``recent_action_log[-1]``
136
  holds the just-dispatched action's acceptance flag.
137
 
138
+ Six components (design §15 + Phase-1 fix):
139
+ r_infect = max(0, 1 - 20 × mean(I)) # weight 0.25
140
+ r_time = 1 - tick / max_ticks # weight 0.05
141
+ r_hosp = max(0, 1 - 10 × mean(hospital_load)) # weight 0.15
142
+ r_casc = max(0, 1 - max(I) / 0.30) # weight 0.15
143
+ r_policy = {-1.0, -0.5, 0.0, +1.0} per §19 # weight 0.35
144
+ r_fair = 1 - pstdev(I) # weight 0.05
145
 
146
  The ``action`` argument is the action just dispatched. We read its
147
  acceptance flag from ``state.recent_action_log[-1]`` rather than
148
  re-dispatching — the simulator already recorded it, and re-dispatch
149
+ would mutate state.
 
 
150
  """
151
  if not state.regions:
152
  return 0.0
153
 
154
  I_values = [r.I for r in state.regions]
155
+ mean_I = sum(I_values) / len(I_values)
156
 
157
+ # r_infect: steepened so gentle outbreak_easy still produces gradient.
158
+ r_infect = _clamp01(1.0 - R_INFECT_STEEPNESS * mean_I)
159
 
160
  if state.max_ticks > 0:
161
  r_time = _clamp01(1.0 - state.tick / state.max_ticks)
 
163
  r_time = 0.0
164
 
165
  hosp_loads = [_hospital_load(I) for I in I_values]
166
+ mean_hosp = sum(hosp_loads) / len(hosp_loads)
167
+ r_hosp = _clamp01(1.0 - R_HOSP_STEEPNESS * mean_hosp)
168
 
169
+ # r_casc: continuous ramp (1.0 at max(I)=0; 0.0 at max(I) >= threshold).
170
+ max_I = max(I_values)
171
+ if R_CASC_HOT_THRESHOLD > 0:
172
+ r_casc = _clamp01(1.0 - max_I / R_CASC_HOT_THRESHOLD)
173
+ else:
174
+ r_casc = 0.0
175
 
176
+ # r_policy: design §19 four-state contract.
177
  if state.recent_action_log:
178
  last_entry = state.recent_action_log[-1]
179
+ r_policy = _r_policy_value(last_entry.action, last_entry.accepted)
180
  else:
181
+ r_policy = R_POLICY_REAL_ACCEPTED # No action dispatched yet → no penalty.
182
 
183
  if len(I_values) >= 2:
184
  r_fair = _clamp01(1.0 - statistics.pstdev(I_values))
 
193
  + W_POLICY * r_policy
194
  + W_FAIR * r_fair
195
  )
196
+ # Final clamp to [-1.0, 1.0] (no longer [0, 1] — M2-A drops the floor).
197
+ if score < -1.0:
198
+ return -1.0
199
+ if score > 1.0:
200
+ return 1.0
201
+ return score
202
 
203
 
204
  def terminal_bonus(state: WorldState) -> float:
server/requirements.txt CHANGED
@@ -1,6 +1,6 @@
1
- openenv[core]>=0.2.0
2
- fastapi>=0.115.0
3
- uvicorn>=0.24.0
4
-
5
-
6
-
 
1
+ openenv[core]>=0.2.0
2
+ fastapi>=0.115.0
3
+ uvicorn>=0.24.0
4
+
5
+
6
+
tests/CLAUDE.md CHANGED
@@ -1,48 +1,48 @@
1
- # tests/CLAUDE.md
2
-
3
- Test surface per subsystem. Smoke bar first, boundary tests next, coverage last.
4
-
5
- ## Belongs here
6
-
7
- - `conftest.py` — repo root on `sys.path` for bare-name imports.
8
- - One test module per subsystem boundary (table below).
9
-
10
- ## Does not belong here
11
-
12
- Helpers that mutate real graders, simulator state, or disk. Fixtures that hit the live HF Space — mock or `pytest.skip`.
13
-
14
- ## Run commands
15
-
16
- ```bash
17
- uv run python -m pytest tests/ -v # all
18
- uv run python -m pytest tests/test_smoke_env.py::test_reset_returns_valid_observation -v # one
19
- uv run python -m pytest --cov tests/ # coverage
20
- ```
21
-
22
- ## Required tests — each maps to exactly one subsystem contract
23
-
24
- | File | Scope | Asserts |
25
- |---|---|---|
26
- | `test_package_exports.py` | wire package | Root `__init__` re-exports `CrisisworldcortexAction/Observation/Env`. |
27
- | `test_smoke_env.py` | `server/` env | `reset()` / `step()` return a valid `CrisisworldcortexObservation`. |
28
- | `test_actions_round_trip.py` | `server/` env | 6 MVP outer actions round-trip; `public_communication` is rejected at runtime. |
29
- | `test_reward_shape.py` | `server/graders/` | Every grader returns values in `[0.0, 1.0]`. |
30
- | `test_reward_non_constancy.py` | `server/graders/` | Grader output varies across ≥ 2 synthetic episodes. |
31
- | `test_anti_hivemind_protocol.py` | `cortex/` | 5 protocol steps fire in order; caps enforced (2 rounds, 1 cross-brain challenge, 1 Critic/brain/tick). |
32
- | `test_collapse_detector.py` | `cortex/` | Metacognition flags when all brains recommend the same action. |
33
- | `test_import_graph.py` | repo-wide | No `import server` under `cortex/**`; no `import cortex` under `server/**`; no `import server.simulator` under `training/**`. |
34
- | `test_baselines_smoke.py` | `baselines/` | B1 / B2 / B3 each run one episode on `outbreak_easy`. |
35
- | `test_training_smoke.py` | `training/` | `train_router.main()` runs one episode against a mocked env under 5 s. |
36
-
37
- ## Binding rules
38
-
39
- - Every public API in a subsystem's CLAUDE.md has ≥ 1 test here.
40
- - Coverage target: 80% per subsystem; 100% for `server/graders/` and `cortex/anti_hivemind.py`.
41
- - No test may take > 10 s unless marked `@pytest.mark.slow` and gated behind `--runslow`.
42
- - `test_import_graph.py` uses a fresh subprocess import, not `sys.modules` monkey-patching — the latter passes under contamination.
43
-
44
- ## Common failure modes
45
-
46
- - Smoke test asserting on current-echo values — breaks when real env logic lands. Assert on shape, not value.
47
- - Module-scope env instantiation in tests — slows collection and hides init errors until runtime.
48
- - Tests that hit the HF Space without a skip guard — CI flakes on rate limits.
 
1
+ # tests/CLAUDE.md
2
+
3
+ Test surface per subsystem. Smoke bar first, boundary tests next, coverage last.
4
+
5
+ ## Belongs here
6
+
7
+ - `conftest.py` — repo root on `sys.path` for bare-name imports.
8
+ - One test module per subsystem boundary (table below).
9
+
10
+ ## Does not belong here
11
+
12
+ Helpers that mutate real graders, simulator state, or disk. Fixtures that hit the live HF Space — mock or `pytest.skip`.
13
+
14
+ ## Run commands
15
+
16
+ ```bash
17
+ uv run python -m pytest tests/ -v # all
18
+ uv run python -m pytest tests/test_smoke_env.py::test_reset_returns_valid_observation -v # one
19
+ uv run python -m pytest --cov tests/ # coverage
20
+ ```
21
+
22
+ ## Required tests — each maps to exactly one subsystem contract
23
+
24
+ | File | Scope | Asserts |
25
+ |---|---|---|
26
+ | `test_package_exports.py` | wire package | Root `__init__` re-exports `CrisisworldcortexAction/Observation/Env`. |
27
+ | `test_smoke_env.py` | `server/` env | `reset()` / `step()` return a valid `CrisisworldcortexObservation`. |
28
+ | `test_actions_round_trip.py` | `server/` env | 6 MVP outer actions round-trip; `public_communication` is rejected at runtime. |
29
+ | `test_reward_shape.py` | `server/graders/` | Every grader returns values in `[0.0, 1.0]`. |
30
+ | `test_reward_non_constancy.py` | `server/graders/` | Grader output varies across ≥ 2 synthetic episodes. |
31
+ | `test_anti_hivemind_protocol.py` | `cortex/` | 5 protocol steps fire in order; caps enforced (2 rounds, 1 cross-brain challenge, 1 Critic/brain/tick). |
32
+ | `test_collapse_detector.py` | `cortex/` | Metacognition flags when all brains recommend the same action. |
33
+ | `test_import_graph.py` | repo-wide | No `import server` under `cortex/**`; no `import cortex` under `server/**`; no `import server.simulator` under `training/**`. |
34
+ | `test_baselines_smoke.py` | `baselines/` | B1 / B2 / B3 each run one episode on `outbreak_easy`. |
35
+ | `test_training_smoke.py` | `training/` | `train_router.main()` runs one episode against a mocked env under 5 s. |
36
+
37
+ ## Binding rules
38
+
39
+ - Every public API in a subsystem's CLAUDE.md has ≥ 1 test here.
40
+ - Coverage target: 80% per subsystem; 100% for `server/graders/` and `cortex/anti_hivemind.py`.
41
+ - No test may take > 10 s unless marked `@pytest.mark.slow` and gated behind `--runslow`.
42
+ - `test_import_graph.py` uses a fresh subprocess import, not `sys.modules` monkey-patching — the latter passes under contamination.
43
+
44
+ ## Common failure modes
45
+
46
+ - Smoke test asserting on current-echo values — breaks when real env logic lands. Assert on shape, not value.
47
+ - Module-scope env instantiation in tests — slows collection and hides init errors until runtime.
48
+ - Tests that hit the HF Space without a skip guard — CI flakes on rate limits.
tests/_helpers/__init__.py ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ """Reusable test helpers for cortex/baselines tests.
2
+
3
+ Lives alongside ``tests/`` (not under ``cortex/``) so production code
4
+ never imports test doubles. Sessions 9-13 reuse the LLMClient stub
5
+ defined here.
6
+ """
tests/_helpers/llm_stub.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """LLMClient-level test double.
2
+
3
+ Drop-in replacement for ``cortex.llm_client.LLMClient`` for tests that
4
+ exercise consumers of the client (subagents, brains, council). Differs
5
+ from the SDK-level stub in ``tests/test_llm_client.py``: that one
6
+ intercepts the OpenAI SDK; this one intercepts the ``LLMClient.chat``
7
+ surface directly, which is what subagents and harnesses see.
8
+
9
+ Reused by sessions 9-13 — do not bury role-specific test logic here.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from dataclasses import dataclass, field
15
+ from typing import Dict, List, Optional
16
+
17
+ from cortex.llm_client import ChatMessage, ChatResponse
18
+
19
+
20
+ @dataclass
21
+ class RecordedCall:
22
+ """One captured ``chat()`` invocation. Tests assert on these."""
23
+
24
+ caller_id: str
25
+ messages: List[ChatMessage]
26
+ max_tokens: Optional[int]
27
+ temperature: Optional[float]
28
+
29
+
30
+ @dataclass
31
+ class StubLLMClient:
32
+ """Quacks like ``LLMClient`` for the ``chat()`` and counter surface.
33
+
34
+ Args:
35
+ scripted_responses: Yielded one per ``chat()`` call, in order.
36
+ Each entry is the response ``content`` string.
37
+ prompt_tokens_per_call: Fake prompt-token billing per call.
38
+ completion_tokens_per_call: Fake completion-token billing per call.
39
+ """
40
+
41
+ scripted_responses: List[str]
42
+ prompt_tokens_per_call: int = 50
43
+ completion_tokens_per_call: int = 30
44
+ calls: List[RecordedCall] = field(default_factory=list)
45
+ _counters: Dict[str, int] = field(default_factory=dict)
46
+
47
+ def chat(
48
+ self,
49
+ caller_id: str,
50
+ messages: List[ChatMessage],
51
+ max_tokens: Optional[int] = None,
52
+ temperature: Optional[float] = None,
53
+ ) -> ChatResponse:
54
+ if not self.scripted_responses:
55
+ raise RuntimeError(
56
+ f"StubLLMClient exhausted: caller_id={caller_id!r}; "
57
+ f"add more scripted_responses or assert call_count earlier."
58
+ )
59
+ content = self.scripted_responses.pop(0)
60
+ self.calls.append(
61
+ RecordedCall(
62
+ caller_id=caller_id,
63
+ messages=list(messages),
64
+ max_tokens=max_tokens,
65
+ temperature=temperature,
66
+ )
67
+ )
68
+ billed = self.prompt_tokens_per_call + self.completion_tokens_per_call
69
+ self._counters[caller_id] = self._counters.get(caller_id, 0) + billed
70
+ return ChatResponse(
71
+ content=content,
72
+ finish_reason="stop",
73
+ prompt_tokens=self.prompt_tokens_per_call,
74
+ completion_tokens=self.completion_tokens_per_call,
75
+ )
76
+
77
+ def tokens_used_for(self, caller_id: str) -> int:
78
+ return self._counters.get(caller_id, 0)
79
+
80
+ @property
81
+ def call_count(self) -> int:
82
+ return len(self.calls)
tests/test_baseline_b1.py CHANGED
@@ -218,7 +218,13 @@ def test_b1_runs_episode_with_valid_json() -> None:
218
  def test_b1_parse_failure_submits_synthetic_rejection() -> None:
219
  """When the LLM emits unparseable text, B1 submits a synthetic
220
  PublicCommunication so the env rejects with accepted=False — landing
221
- r_policy=0 in outer_reward, and the action log shows the rejection.
 
 
 
 
 
 
222
  """
223
  env_inner = CrisisworldcortexEnvironment()
224
  env = _InProcessEnvAdapter(env_inner)
@@ -233,36 +239,26 @@ def test_b1_parse_failure_submits_synthetic_rejection() -> None:
233
 
234
  trajectory = agent.run_episode(task="outbreak_easy", seed=0, max_ticks=5)
235
 
236
- # Both parse failures were detected and counted.
237
- assert trajectory["parse_failure_count"] == 2
238
 
239
- # B1 did NOT crash on parse failure — at least 3 ticks ran, even
240
- # though the env may then have hit a terminal (success-on-3-safe-ticks
241
- # or otherwise). What's binding: parse failure does not raise.
242
- assert trajectory["steps_taken"] >= 3, (
243
  f"steps_taken={trajectory['steps_taken']!r} - parse failure "
244
- f"shouldn't kill the agent before tick 3"
245
  )
246
 
247
- # The first two action-log entries show V2 rejection (synthetic
248
- # public_communication was submitted; env returned accepted=False).
249
  log = env_inner._world_state.recent_action_log
250
- assert len(log) >= 2
251
  assert log[0].action.kind == "public_communication"
252
  assert log[0].accepted is False, "parse-failure synthetic must be rejected by env"
253
- assert log[1].action.kind == "public_communication"
254
- assert log[1].accepted is False
255
-
256
- # The third entry should be the first parsed NoOp.
257
- assert log[2].action.kind == "no_op"
258
- assert log[2].accepted is True
259
 
260
- # B1's local trajectory carries the raw snippets for forensic use.
261
  assert trajectory["action_history"][0]["parse_failure"] is True
262
  assert trajectory["action_history"][0]["raw_llm"] == "I cannot help with that."
263
- assert trajectory["action_history"][1]["parse_failure"] is True
264
- assert trajectory["action_history"][1]["raw_llm"] == "Sorry, no JSON."
265
- assert trajectory["action_history"][2]["parse_failure"] is False
266
 
267
 
268
  def test_b1_caller_id_format_short_colon_separated() -> None:
@@ -335,7 +331,10 @@ def test_b1_step_event_carries_rich_context() -> None:
335
  from baselines.flat_agent import B1StepEvent
336
 
337
  env = _InProcessEnvAdapter(CrisisworldcortexEnvironment())
338
- llm = _StubLLMClient(["I cannot help with that."] + ['{"kind": "no_op"}'] * 5)
 
 
 
339
  agent = B1FlatAgent(env=env, llm=llm)
340
 
341
  events: list[B1StepEvent] = []
@@ -348,18 +347,17 @@ def test_b1_step_event_carries_rich_context() -> None:
348
 
349
  assert len(events) >= 2
350
 
351
- # Tick 1: parse failure. Submitted action is the synthetic V2-rejected
352
- # PublicCommunication marker; env returns accepted=False; reward in [0,1].
353
  e1 = events[0]
354
  assert e1.tick == 1
355
- assert e1.parse_failure is True
356
- assert e1.error == "parse_failure"
357
- assert e1.raw_llm == "I cannot help with that."
358
- assert e1.action.kind == "public_communication"
359
- assert 0.0 <= e1.reward <= 1.0
360
  assert isinstance(e1.done, bool)
361
 
362
- # Tick 2: clean parse. Submitted is NoOp. error must be None.
363
  e2 = events[1]
364
  assert e2.tick == 2
365
  assert e2.parse_failure is False
 
218
  def test_b1_parse_failure_submits_synthetic_rejection() -> None:
219
  """When the LLM emits unparseable text, B1 submits a synthetic
220
  PublicCommunication so the env rejects with accepted=False — landing
221
+ r_policy=-1.0 in outer_reward (Phase-1 fix per design §19) and the
222
+ action log shows the rejection.
223
+
224
+ Phase-1 contract: parse-failure now TERMINATES the episode at the
225
+ rejection tick (state.terminal = "failure" → obs.done = True). So
226
+ only the first parse-failure lands; the second LLM response in the
227
+ stub queue never gets dispatched.
228
  """
229
  env_inner = CrisisworldcortexEnvironment()
230
  env = _InProcessEnvAdapter(env_inner)
 
239
 
240
  trajectory = agent.run_episode(task="outbreak_easy", seed=0, max_ticks=5)
241
 
242
+ # Phase-1: parse-failure terminates only the first marker lands.
243
+ assert trajectory["parse_failure_count"] == 1
244
 
245
+ # B1 did NOT crash on parse failure — exactly 1 tick ran (parse-failure
246
+ # marker submitted, env terminated episode).
247
+ assert trajectory["steps_taken"] == 1, (
 
248
  f"steps_taken={trajectory['steps_taken']!r} - parse failure "
249
+ f"should terminate at tick 1 under the §19 contract"
250
  )
251
 
252
+ # The action-log entry shows synthetic public_communication
253
+ # (parse-failure marker, honesty=0.0) rejected by env.
254
  log = env_inner._world_state.recent_action_log
255
+ assert len(log) == 1
256
  assert log[0].action.kind == "public_communication"
257
  assert log[0].accepted is False, "parse-failure synthetic must be rejected by env"
 
 
 
 
 
 
258
 
259
+ # B1's local trajectory carries the raw snippet for forensic use.
260
  assert trajectory["action_history"][0]["parse_failure"] is True
261
  assert trajectory["action_history"][0]["raw_llm"] == "I cannot help with that."
 
 
 
262
 
263
 
264
  def test_b1_caller_id_format_short_colon_separated() -> None:
 
331
  from baselines.flat_agent import B1StepEvent
332
 
333
  env = _InProcessEnvAdapter(CrisisworldcortexEnvironment())
334
+ # Use a clean-parse first response; parse-failure now terminates the
335
+ # episode at tick 1 under the §19 contract, so we exercise the
336
+ # tick-1 + tick-2 sequence with both responses being valid JSON.
337
+ llm = _StubLLMClient(['{"kind": "no_op"}'] * 5)
338
  agent = B1FlatAgent(env=env, llm=llm)
339
 
340
  events: list[B1StepEvent] = []
 
347
 
348
  assert len(events) >= 2
349
 
350
+ # Tick 1: clean parse. Submitted is NoOp. error must be None.
 
351
  e1 = events[0]
352
  assert e1.tick == 1
353
+ assert e1.parse_failure is False
354
+ assert e1.error is None
355
+ assert e1.action.kind == "no_op"
356
+ assert e1.raw_llm == '{"kind": "no_op"}'
357
+ assert -1.0 <= e1.reward <= 1.0 # Phase-1 range relaxed from [0,1].
358
  assert isinstance(e1.done, bool)
359
 
360
+ # Tick 2: another clean parse. Submitted is NoOp. error must be None.
361
  e2 = events[1]
362
  assert e2.tick == 2
363
  assert e2.parse_failure is False
tests/test_cortex_brain_executive.py ADDED
@@ -0,0 +1,206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Session 11 - Brain Executive aggregation tests.
2
+
3
+ Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 15-21 + M-FR-3
4
+ (partial evidence union; CandidatePlan and CriticReport schemas have
5
+ no evidence field).
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import pytest
11
+
12
+ from cortex.brains import aggregate_brain_outputs
13
+ from cortex.schemas import (
14
+ BeliefState,
15
+ CandidatePlan,
16
+ CriticReport,
17
+ EvidenceCitation,
18
+ Hypothesis,
19
+ PerceptionReport,
20
+ RegionBeliefEstimate,
21
+ )
22
+ from CrisisWorldCortex.models import (
23
+ DeployResource,
24
+ NoOp,
25
+ RestrictMovement,
26
+ )
27
+
28
+
29
+ def _belief(uncertainty: float = 0.4, evidence_count: int = 1) -> BeliefState:
30
+ return BeliefState(
31
+ brain="epidemiology",
32
+ latent_estimates={
33
+ "R1": RegionBeliefEstimate(
34
+ estimated_infection_rate=0.05,
35
+ estimated_r_effective=1.2,
36
+ estimated_compliance=0.85,
37
+ ),
38
+ },
39
+ hypotheses=[Hypothesis(label="h1", weight=0.6, explanation="rising")],
40
+ uncertainty=uncertainty,
41
+ reducible_by_more_thought=0.3,
42
+ evidence=[
43
+ EvidenceCitation(source="telemetry", ref=f"R1.cases@{i}", excerpt=f"e{i}")
44
+ for i in range(evidence_count)
45
+ ],
46
+ )
47
+
48
+
49
+ def _plan(action=None, confidence: float = 0.75, expected_value: float = 0.6) -> CandidatePlan:
50
+ if action is None:
51
+ action = DeployResource(region="R1", resource_type="test_kits", quantity=100)
52
+ return CandidatePlan(
53
+ action_sketch="Deploy 100 test_kits to R1",
54
+ expected_outer_action=action,
55
+ expected_value=expected_value,
56
+ cost=200.0,
57
+ assumptions=["kits available"],
58
+ falsifiers=["R1 cases drop without intervention"],
59
+ confidence=confidence,
60
+ )
61
+
62
+
63
+ def _critic(severity: float = 0.3) -> CriticReport:
64
+ return CriticReport(
65
+ brain="epidemiology",
66
+ target_plan_id="plan-0",
67
+ attacks=["limited reach"],
68
+ missing_considerations=[],
69
+ would_change_mind_if=[],
70
+ severity=severity,
71
+ )
72
+
73
+
74
+ def _perception(evidence_count: int = 1) -> PerceptionReport:
75
+ return PerceptionReport(
76
+ brain="epidemiology",
77
+ salient_signals=["R1 cases rising"],
78
+ anomalies=[],
79
+ confidence=0.7,
80
+ evidence=[
81
+ EvidenceCitation(source="telemetry", ref=f"R1.perception@{i}", excerpt=f"p{i}")
82
+ for i in range(evidence_count)
83
+ ],
84
+ )
85
+
86
+
87
+ # T4
88
+ def test_brain_executive_aggregates_subagent_outputs() -> None:
89
+ rec = aggregate_brain_outputs(
90
+ brain_id="epidemiology",
91
+ perception=_perception(),
92
+ beliefs=[_belief()],
93
+ plans=[_plan()],
94
+ critics=[_critic()],
95
+ )
96
+ assert rec.brain == "epidemiology"
97
+ assert rec.top_action.kind == "deploy_resource"
98
+ assert rec.top_confidence > 0.0
99
+ assert rec.tokens_used == 0
100
+
101
+
102
+ # T5 -- Decision 16
103
+ def test_brain_executive_top_confidence_includes_uncertainty() -> None:
104
+ rec = aggregate_brain_outputs(
105
+ brain_id="epidemiology",
106
+ perception=_perception(),
107
+ beliefs=[_belief(uncertainty=0.4)],
108
+ plans=[_plan(confidence=0.75)],
109
+ critics=[_critic()],
110
+ )
111
+ # D16: top_confidence == confidence x (1 - uncertainty) == 0.75 x 0.6 == 0.45
112
+ assert rec.top_confidence == pytest.approx(0.45)
113
+
114
+
115
+ # T6 -- Decision 17
116
+ def test_brain_executive_minority_actions_excludes_top() -> None:
117
+ plan_a = _plan(
118
+ action=DeployResource(region="R1", resource_type="test_kits", quantity=100),
119
+ confidence=0.8,
120
+ expected_value=0.7,
121
+ )
122
+ plan_b = _plan(
123
+ action=RestrictMovement(region="R1", severity="moderate"),
124
+ confidence=0.5,
125
+ expected_value=0.4,
126
+ )
127
+ # plan_a wins: 0.8 * 0.7 = 0.56 > 0.5 * 0.4 = 0.20
128
+
129
+ rec = aggregate_brain_outputs(
130
+ brain_id="epidemiology",
131
+ perception=_perception(),
132
+ beliefs=[_belief(), _belief(uncertainty=0.5)],
133
+ plans=[plan_a, plan_b],
134
+ critics=[_critic(), _critic()],
135
+ )
136
+
137
+ assert rec.top_action.kind == "deploy_resource"
138
+ assert len(rec.minority_actions) == 1
139
+ assert rec.minority_actions[0].kind == "restrict_movement"
140
+
141
+
142
+ # T7 -- Decision 20 + M-FR-3
143
+ def test_brain_executive_evidence_union() -> None:
144
+ perception = _perception(evidence_count=1)
145
+ belief = _belief(evidence_count=2)
146
+
147
+ rec = aggregate_brain_outputs(
148
+ brain_id="epidemiology",
149
+ perception=perception,
150
+ beliefs=[belief],
151
+ plans=[_plan()],
152
+ critics=[_critic()],
153
+ )
154
+
155
+ # M-FR-3: union of perception.evidence + belief.evidence (3 total)
156
+ assert len(rec.evidence) == 3
157
+ assert rec.evidence[0].ref.startswith("R1.perception")
158
+ assert rec.evidence[1].ref.startswith("R1.cases")
159
+
160
+
161
+ # T8 -- empty fallback
162
+ def test_brain_executive_handles_empty_subagent_outputs() -> None:
163
+ empty_belief = BeliefState(
164
+ brain="epidemiology",
165
+ latent_estimates={},
166
+ hypotheses=[],
167
+ uncertainty=1.0,
168
+ reducible_by_more_thought=0.0,
169
+ evidence=[],
170
+ )
171
+ empty_plan = CandidatePlan(
172
+ action_sketch="(empty)",
173
+ expected_outer_action=NoOp(),
174
+ expected_value=0.0,
175
+ cost=0.0,
176
+ assumptions=[],
177
+ falsifiers=[],
178
+ confidence=0.0,
179
+ )
180
+ empty_critic = CriticReport(
181
+ brain="epidemiology",
182
+ target_plan_id="",
183
+ attacks=[],
184
+ missing_considerations=[],
185
+ would_change_mind_if=[],
186
+ severity=0.0,
187
+ )
188
+ empty_perception = PerceptionReport(
189
+ brain="epidemiology",
190
+ salient_signals=[],
191
+ anomalies=[],
192
+ confidence=0.0,
193
+ evidence=[],
194
+ )
195
+
196
+ rec = aggregate_brain_outputs(
197
+ brain_id="epidemiology",
198
+ perception=empty_perception,
199
+ beliefs=[empty_belief],
200
+ plans=[empty_plan],
201
+ critics=[empty_critic],
202
+ )
203
+
204
+ assert rec.top_action.kind == "no_op"
205
+ assert rec.top_confidence == 0.0
206
+ assert rec.uncertainty == 1.0
tests/test_cortex_brain_smoke.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Session 11 - Single-brain end-to-end smoke + no-LLM-in-Python-layers test.
2
+
3
+ T9 asserts Brain.compute_perception and Brain.compute_lens are pure
4
+ Python (zero LLM calls). T10 is the integration smoke gate per Phase A
5
+ section 10: a single brain runs end-to-end on a real observation,
6
+ returns a BrainRecommendation. Three LLM calls in canonical
7
+ WorldModeler -> Planner -> Critic order with locked caller_id format.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import json
13
+
14
+ from cortex.brains import EpiBrain
15
+ from cortex.schemas import BrainRecommendation
16
+ from CrisisWorldCortex.models import (
17
+ CrisisworldcortexObservation,
18
+ RegionTelemetry,
19
+ ResourceInventory,
20
+ )
21
+ from tests._helpers.llm_stub import StubLLMClient
22
+
23
+
24
+ def _make_obs() -> CrisisworldcortexObservation:
25
+ return CrisisworldcortexObservation(
26
+ regions=[
27
+ RegionTelemetry(
28
+ region=f"R{i + 1}",
29
+ reported_cases_d_ago=5 if i == 0 else 1,
30
+ hospital_load=0.3 if i == 0 else 0.1,
31
+ compliance_proxy=0.85,
32
+ )
33
+ for i in range(4)
34
+ ],
35
+ resources=ResourceInventory(
36
+ test_kits=1000,
37
+ hospital_beds_free=500,
38
+ mobile_units=20,
39
+ vaccine_doses=2000,
40
+ ),
41
+ active_restrictions=[],
42
+ legal_constraints=[],
43
+ tick=3,
44
+ ticks_remaining=9,
45
+ cognition_budget_remaining=5200,
46
+ recent_action_log=[],
47
+ )
48
+
49
+
50
+ _VALID_BELIEF = json.dumps(
51
+ {
52
+ "brain": "epidemiology",
53
+ "latent_estimates": {
54
+ "R1": {
55
+ "estimated_infection_rate": 0.05,
56
+ "estimated_r_effective": 1.2,
57
+ "estimated_compliance": 0.85,
58
+ "confidence_intervals": {},
59
+ }
60
+ },
61
+ "hypotheses": [{"label": "rising", "weight": 0.6, "explanation": "R1 cases up"}],
62
+ "uncertainty": 0.4,
63
+ "reducible_by_more_thought": 0.3,
64
+ "evidence": [
65
+ {"source": "telemetry", "ref": "R1.cases", "excerpt": "5"},
66
+ {"source": "policy", "ref": "R1.restriction", "excerpt": "none"},
67
+ ],
68
+ }
69
+ )
70
+
71
+ _VALID_PLAN = json.dumps(
72
+ {
73
+ "action_sketch": "Deploy 100 test_kits to R1",
74
+ "expected_outer_action": {
75
+ "kind": "deploy_resource",
76
+ "region": "R1",
77
+ "resource_type": "test_kits",
78
+ "quantity": 100,
79
+ },
80
+ "expected_value": 0.6,
81
+ "cost": 200.0,
82
+ "assumptions": ["kits inventory > 100"],
83
+ "falsifiers": ["R1 cases drop without intervention"],
84
+ "confidence": 0.75,
85
+ }
86
+ )
87
+
88
+ _VALID_CRITIC = json.dumps(
89
+ {
90
+ "brain": "epidemiology",
91
+ "target_plan_id": "plan-0",
92
+ "attacks": ["ignores R3 hospital saturation"],
93
+ "missing_considerations": [],
94
+ "would_change_mind_if": [],
95
+ "severity": 0.3,
96
+ }
97
+ )
98
+
99
+
100
+ # T9
101
+ def test_brain_runs_zero_llm_calls_in_python_layers() -> None:
102
+ """Perception and Lens are pure Python; no LLMClient invocation."""
103
+ stub = StubLLMClient(scripted_responses=[]) # any chat() would raise
104
+ brain = EpiBrain(stub)
105
+ obs = _make_obs()
106
+
107
+ perception = brain.compute_perception(obs)
108
+ lensed = brain.compute_lens(obs, last_reward=0.0)
109
+
110
+ assert stub.call_count == 0
111
+ assert perception.brain == "epidemiology"
112
+ assert lensed.brain == "epidemiology"
113
+
114
+
115
+ # T10 -- integration smoke gate (Phase A section 10)
116
+ def test_brain_smoke_one_tick_three_llm_calls_in_order() -> None:
117
+ """Full round-1 tick: WorldModeler -> Planner -> Critic, in that order."""
118
+ stub = StubLLMClient(scripted_responses=[_VALID_BELIEF, _VALID_PLAN, _VALID_CRITIC])
119
+ brain = EpiBrain(stub)
120
+ obs = _make_obs()
121
+
122
+ rec = brain.run_tick(obs, last_reward=0.0, tick=3)
123
+
124
+ assert isinstance(rec, BrainRecommendation)
125
+ assert rec.brain == "epidemiology"
126
+ assert rec.top_action.kind == "deploy_resource"
127
+ assert stub.call_count == 3, "exactly 3 LLM calls per Phase A 'WM + Planner + Critic'"
128
+
129
+ # Order pin per user adjustment: WM (s0) -> Planner (s1) -> Critic (s2)
130
+ assert stub.calls[0].caller_id.endswith(":world_modeler:t3:r1:s0"), (
131
+ f"first call must be WorldModeler, got {stub.calls[0].caller_id!r}"
132
+ )
133
+ assert stub.calls[1].caller_id.endswith(":planner:t3:r1:s1"), (
134
+ f"second call must be Planner, got {stub.calls[1].caller_id!r}"
135
+ )
136
+ assert stub.calls[2].caller_id.endswith(":critic:t3:r1:s2"), (
137
+ f"third call must be Critic, got {stub.calls[2].caller_id!r}"
138
+ )
139
+
140
+ # Brain prefix locks
141
+ for call in stub.calls:
142
+ assert call.caller_id.startswith("cortex:epidemiology:")
tests/test_cortex_lenses.py ADDED
@@ -0,0 +1,238 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Session 10 - Cortex lens tests.
2
+
3
+ Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 9-14 + §2 A1 and the
4
+ user's Session 10 proposal acceptance with 8 tests + the M-FR-4 rename
5
+ (epi_pressure) and T7 tightening (non-bool float).
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import Iterable
11
+
12
+ import pytest
13
+
14
+ from cortex.lenses import lens_for
15
+ from cortex.schemas import BrainLensedObservation
16
+ from CrisisWorldCortex.models import (
17
+ CrisisworldcortexObservation,
18
+ Escalate,
19
+ ExecutedAction,
20
+ LegalConstraint,
21
+ NoOp,
22
+ RegionTelemetry,
23
+ ResourceInventory,
24
+ Restriction,
25
+ )
26
+
27
+ # ============================================================================
28
+ # Test fixtures
29
+ # ============================================================================
30
+
31
+
32
+ def _make_obs(
33
+ cases_per_region: Iterable[int] = (5, 1, 1, 1),
34
+ hospital_loads: Iterable[float] = (0.3, 0.1, 0.1, 0.1),
35
+ compliance_proxies: Iterable[float] = (0.85, 0.95, 0.95, 0.95),
36
+ test_kits: int = 1000,
37
+ hospital_beds_free: int = 500,
38
+ mobile_units: int = 20,
39
+ vaccine_doses: int = 2000,
40
+ restrictions: Iterable[Restriction] = (),
41
+ legal_constraints: Iterable[LegalConstraint] = (),
42
+ recent_action_log: Iterable[ExecutedAction] = (),
43
+ ) -> CrisisworldcortexObservation:
44
+ cases = list(cases_per_region)
45
+ loads = list(hospital_loads)
46
+ comps = list(compliance_proxies)
47
+ return CrisisworldcortexObservation(
48
+ regions=[
49
+ RegionTelemetry(
50
+ region=f"R{i + 1}",
51
+ reported_cases_d_ago=cases[i],
52
+ hospital_load=loads[i],
53
+ compliance_proxy=comps[i],
54
+ )
55
+ for i in range(4)
56
+ ],
57
+ resources=ResourceInventory(
58
+ test_kits=test_kits,
59
+ hospital_beds_free=hospital_beds_free,
60
+ mobile_units=mobile_units,
61
+ vaccine_doses=vaccine_doses,
62
+ ),
63
+ active_restrictions=list(restrictions),
64
+ legal_constraints=list(legal_constraints),
65
+ tick=3,
66
+ ticks_remaining=9,
67
+ cognition_budget_remaining=5200,
68
+ recent_action_log=list(recent_action_log),
69
+ )
70
+
71
+
72
+ # ============================================================================
73
+ # T1 - Epi lens emphasizes telemetry; uses epi_pressure (M-FR-4 rename)
74
+ # ============================================================================
75
+
76
+
77
+ def test_epi_lens_emphasizes_telemetry() -> None:
78
+ obs = _make_obs()
79
+ lensed = lens_for("epidemiology", obs, last_reward=0.5)
80
+
81
+ assert isinstance(lensed, BrainLensedObservation)
82
+ assert lensed.brain == "epidemiology"
83
+ assert lensed.last_reward == 0.5
84
+
85
+ keys = set(lensed.derived_features.keys())
86
+ assert {"epi_pressure", "worst_region_infection", "transmission_rate_trend"} <= keys
87
+
88
+ # M-FR-2: trend is 0.0 in MVP (no history available in single-obs lens)
89
+ assert lensed.derived_features["transmission_rate_trend"] == 0.0
90
+
91
+ assert "regions[*].reported_cases_d_ago" in lensed.salient_field_ids
92
+ assert "regions[*].hospital_load" in lensed.salient_field_ids
93
+
94
+
95
+ # ============================================================================
96
+ # T2 - Logistics lens emphasizes resources
97
+ # ============================================================================
98
+
99
+
100
+ def test_logistics_lens_emphasizes_resources() -> None:
101
+ obs = _make_obs(test_kits=100, hospital_beds_free=50, mobile_units=10, vaccine_doses=200)
102
+ lensed = lens_for("logistics", obs, last_reward=0.5)
103
+
104
+ assert lensed.brain == "logistics"
105
+ keys = set(lensed.derived_features.keys())
106
+ expected_keys = {
107
+ "total_inventory",
108
+ "hospital_load_max",
109
+ "deployment_feasibility_R1",
110
+ "deployment_feasibility_R2",
111
+ "deployment_feasibility_R3",
112
+ "deployment_feasibility_R4",
113
+ }
114
+ assert expected_keys <= keys
115
+
116
+ # 100 + 50 + 10 + 200 = 360
117
+ assert lensed.derived_features["total_inventory"] == 360.0
118
+
119
+ assert "resources.test_kits" in lensed.salient_field_ids
120
+
121
+
122
+ # ============================================================================
123
+ # T3 - Governance lens emphasizes legal
124
+ # ============================================================================
125
+
126
+
127
+ def test_governance_lens_emphasizes_legal() -> None:
128
+ obs = _make_obs(
129
+ restrictions=[Restriction(region="R1", severity="moderate", ticks_remaining=3)],
130
+ legal_constraints=[
131
+ LegalConstraint(rule_id="L1", blocked_action="restrict_movement.strict")
132
+ ],
133
+ )
134
+ lensed = lens_for("governance", obs, last_reward=0.5)
135
+
136
+ assert lensed.brain == "governance"
137
+ keys = set(lensed.derived_features.keys())
138
+ assert {
139
+ "escalation_unlocked_strict",
140
+ "legal_constraints_count",
141
+ "restrictions_active_count",
142
+ } <= keys
143
+
144
+ assert lensed.derived_features["legal_constraints_count"] == 1.0
145
+ assert lensed.derived_features["restrictions_active_count"] == 1.0
146
+ assert "active_restrictions[*]" in lensed.salient_field_ids
147
+
148
+
149
+ # ============================================================================
150
+ # T4 - Lens does NOT strip raw_obs (D13)
151
+ # ============================================================================
152
+
153
+
154
+ def test_lens_does_not_strip_raw_obs() -> None:
155
+ obs = _make_obs(restrictions=[Restriction(region="R1", severity="moderate", ticks_remaining=3)])
156
+
157
+ for brain in ("epidemiology", "logistics", "governance"):
158
+ lensed = lens_for(brain, obs, last_reward=0.0)
159
+ # Pydantic deep-equality on the full observation
160
+ assert lensed.raw_obs == obs
161
+
162
+
163
+ # ============================================================================
164
+ # T5 - V2 brain ids raise KeyError (Decision 9, post-review)
165
+ # ============================================================================
166
+
167
+
168
+ def test_lens_for_v2_brain_raises_key_error() -> None:
169
+ obs = _make_obs()
170
+
171
+ for v2_brain in ("communications", "equity"):
172
+ with pytest.raises(KeyError):
173
+ lens_for(v2_brain, obs, last_reward=0.0)
174
+
175
+ with pytest.raises(KeyError):
176
+ lens_for("not_a_brain", obs, last_reward=0.0)
177
+
178
+
179
+ # ============================================================================
180
+ # T7 - All derived_features values are non-bool floats (D14 + tightened)
181
+ # ============================================================================
182
+
183
+
184
+ def test_lens_derived_features_all_floats() -> None:
185
+ obs = _make_obs(
186
+ restrictions=[Restriction(region="R1", severity="moderate", ticks_remaining=3)],
187
+ legal_constraints=[
188
+ LegalConstraint(rule_id="L1", blocked_action="restrict_movement.strict")
189
+ ],
190
+ )
191
+
192
+ for brain in ("epidemiology", "logistics", "governance"):
193
+ lensed = lens_for(brain, obs, last_reward=0.0)
194
+ for key, value in lensed.derived_features.items():
195
+ assert isinstance(value, float) and not isinstance(value, bool), (
196
+ f"derived_features[{key!r}] = {value!r} ({type(value).__name__}) "
197
+ f"is not a non-bool float"
198
+ )
199
+
200
+
201
+ # ============================================================================
202
+ # T8 - Governance lens detects accepted escalate(national)
203
+ # ============================================================================
204
+
205
+
206
+ def test_governance_lens_detects_escalation_unlocked() -> None:
207
+ obs_with_accepted = _make_obs(
208
+ recent_action_log=[
209
+ ExecutedAction(tick=1, action=NoOp(), accepted=True),
210
+ ExecutedAction(tick=2, action=Escalate(to_authority="national"), accepted=True),
211
+ ]
212
+ )
213
+ lensed = lens_for("governance", obs_with_accepted, last_reward=0.0)
214
+ assert lensed.derived_features["escalation_unlocked_strict"] == 1.0
215
+
216
+ obs_without = _make_obs(
217
+ recent_action_log=[ExecutedAction(tick=1, action=NoOp(), accepted=True)]
218
+ )
219
+ lensed_w = lens_for("governance", obs_without, last_reward=0.0)
220
+ assert lensed_w.derived_features["escalation_unlocked_strict"] == 0.0
221
+
222
+ # Rejected escalate must NOT count
223
+ obs_rejected = _make_obs(
224
+ recent_action_log=[
225
+ ExecutedAction(tick=1, action=Escalate(to_authority="national"), accepted=False),
226
+ ]
227
+ )
228
+ lensed_r = lens_for("governance", obs_rejected, last_reward=0.0)
229
+ assert lensed_r.derived_features["escalation_unlocked_strict"] == 0.0
230
+
231
+ # Accepted escalate(regional) must NOT count -- only "national" unlocks strict
232
+ obs_regional = _make_obs(
233
+ recent_action_log=[
234
+ ExecutedAction(tick=1, action=Escalate(to_authority="regional"), accepted=True),
235
+ ]
236
+ )
237
+ lensed_re = lens_for("governance", obs_regional, last_reward=0.0)
238
+ assert lensed_re.derived_features["escalation_unlocked_strict"] == 0.0
tests/test_cortex_perception.py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Session 11 - Perception subagent tests (deterministic Python; no LLM).
2
+
3
+ Per cortex/CLAUDE.md binding (Perception is pure Python) and Phase A
4
+ Decisions 9 (V2 KeyError) + 63 (salient_signals cap at 5).
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from typing import Iterable
10
+
11
+ import pytest
12
+
13
+ from cortex.schemas import PerceptionReport
14
+ from cortex.subagents import perception_for
15
+ from CrisisWorldCortex.models import (
16
+ CrisisworldcortexObservation,
17
+ LegalConstraint,
18
+ RegionTelemetry,
19
+ ResourceInventory,
20
+ Restriction,
21
+ )
22
+
23
+
24
+ def _make_obs(
25
+ cases_per_region: Iterable[int] = (5, 1, 0, 0),
26
+ hospital_loads: Iterable[float] = (0.7, 0.1, 0.1, 0.1),
27
+ test_kits: int = 200,
28
+ hospital_beds_free: int = 50,
29
+ mobile_units: int = 2,
30
+ vaccine_doses: int = 300,
31
+ restrictions: Iterable[Restriction] = (),
32
+ legal_constraints: Iterable[LegalConstraint] = (),
33
+ ) -> CrisisworldcortexObservation:
34
+ cases = list(cases_per_region)
35
+ loads = list(hospital_loads)
36
+ return CrisisworldcortexObservation(
37
+ regions=[
38
+ RegionTelemetry(
39
+ region=f"R{i + 1}",
40
+ reported_cases_d_ago=cases[i],
41
+ hospital_load=loads[i],
42
+ compliance_proxy=0.85,
43
+ )
44
+ for i in range(4)
45
+ ],
46
+ resources=ResourceInventory(
47
+ test_kits=test_kits,
48
+ hospital_beds_free=hospital_beds_free,
49
+ mobile_units=mobile_units,
50
+ vaccine_doses=vaccine_doses,
51
+ ),
52
+ active_restrictions=list(restrictions),
53
+ legal_constraints=list(legal_constraints),
54
+ tick=3,
55
+ ticks_remaining=9,
56
+ cognition_budget_remaining=5200,
57
+ recent_action_log=[],
58
+ )
59
+
60
+
61
+ # T1
62
+ def test_perception_runs_without_llm_call() -> None:
63
+ """Perception is pure Python; the function does not take an LLMClient."""
64
+ obs = _make_obs()
65
+ for brain in ("epidemiology", "logistics", "governance"):
66
+ report = perception_for(brain, obs)
67
+ assert isinstance(report, PerceptionReport)
68
+ assert report.brain == brain
69
+ assert isinstance(report.confidence, float)
70
+ assert 0.0 <= report.confidence <= 1.0
71
+
72
+
73
+ # T2
74
+ def test_perception_for_v2_brain_raises_key_error() -> None:
75
+ obs = _make_obs()
76
+ for v2_brain in ("communications", "equity"):
77
+ with pytest.raises(KeyError):
78
+ perception_for(v2_brain, obs)
79
+ with pytest.raises(KeyError):
80
+ perception_for("not_a_brain", obs)
81
+
82
+
83
+ # T3
84
+ @pytest.mark.parametrize("brain", ["epidemiology", "logistics", "governance"])
85
+ def test_perception_brain_specific_signals(brain: str) -> None:
86
+ obs = _make_obs(
87
+ cases_per_region=(20, 1, 0, 0),
88
+ hospital_loads=(0.7, 0.1, 0.1, 0.1),
89
+ test_kits=100, # below threshold (300)
90
+ hospital_beds_free=50, # below threshold (100)
91
+ mobile_units=2, # below threshold (5)
92
+ vaccine_doses=200, # below threshold (500)
93
+ restrictions=[
94
+ Restriction(region="R1", severity="moderate", ticks_remaining=3),
95
+ ],
96
+ legal_constraints=[
97
+ LegalConstraint(rule_id="L1", blocked_action="restrict_movement.strict"),
98
+ ],
99
+ )
100
+ report = perception_for(brain, obs)
101
+
102
+ assert report.brain == brain
103
+ # Decision 63 / OQ-2 cap: at most 5 entries
104
+ assert len(report.salient_signals) <= 5
105
+
106
+ if brain == "epidemiology":
107
+ assert any("R1" in s for s in report.salient_signals), (
108
+ f"epi salient_signals should reference R1, got {report.salient_signals}"
109
+ )
110
+ elif brain == "logistics":
111
+ joined = " ".join(report.salient_signals).lower()
112
+ assert "kits" in joined or "mobile" in joined or "vaccine" in joined or "beds" in joined
113
+ elif brain == "governance":
114
+ assert any("R1" in s and "moderate" in s.lower() for s in report.salient_signals), (
115
+ f"governance salient_signals should mention R1 moderate, got {report.salient_signals}"
116
+ )
tests/test_cortex_subagents.py ADDED
@@ -0,0 +1,387 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Session 9 - Cortex subagent tests (WorldModeler, Planner, Critic).
2
+
3
+ Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 1-8 + 62 and the user's
4
+ proposal acceptance with 11 tests total. RED-tests-first.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import json
10
+ import re
11
+ from typing import Any, Dict, Optional
12
+
13
+ import pytest
14
+
15
+ from cortex.schemas import (
16
+ BeliefState,
17
+ CandidatePlan,
18
+ CriticReport,
19
+ EvidenceCitation,
20
+ PerceptionReport,
21
+ SubagentInput,
22
+ )
23
+ from cortex.subagents import (
24
+ PROMPTS_DIR,
25
+ CriticSubagent,
26
+ PlannerSubagent,
27
+ WorldModelerSubagent,
28
+ )
29
+ from tests._helpers.llm_stub import StubLLMClient
30
+
31
+ # ============================================================================
32
+ # Test fixtures
33
+ # ============================================================================
34
+
35
+
36
+ _VALID_BELIEF_PAYLOAD: Dict[str, Any] = {
37
+ "brain": "epidemiology",
38
+ "latent_estimates": {
39
+ "R1": {
40
+ "estimated_infection_rate": 0.05,
41
+ "estimated_r_effective": 1.2,
42
+ "estimated_compliance": 0.85,
43
+ "confidence_intervals": {},
44
+ },
45
+ },
46
+ "hypotheses": [{"label": "rising-r1", "weight": 0.6, "explanation": "telemetry trending up"}],
47
+ "uncertainty": 0.4,
48
+ "reducible_by_more_thought": 0.3,
49
+ "evidence": [
50
+ {"source": "telemetry", "ref": "R1.reported_cases@t3", "excerpt": "rising"},
51
+ {"source": "policy", "ref": "R1.restriction", "excerpt": "moderate"},
52
+ ],
53
+ }
54
+
55
+
56
+ _VALID_PLAN_PAYLOAD: Dict[str, Any] = {
57
+ "action_sketch": "Deploy 100 test_kits to R1",
58
+ "expected_outer_action": {
59
+ "kind": "deploy_resource",
60
+ "region": "R1",
61
+ "resource_type": "test_kits",
62
+ "quantity": 100,
63
+ },
64
+ "expected_value": 0.6,
65
+ "cost": 200.0,
66
+ "assumptions": ["kits inventory > 100"],
67
+ "falsifiers": ["R1 cases drop without intervention"],
68
+ "confidence": 0.75,
69
+ }
70
+
71
+
72
+ _VALID_CRITIC_PAYLOAD: Dict[str, Any] = {
73
+ "brain": "logistics",
74
+ "target_plan_id": "plan-1",
75
+ "attacks": ["ignores R3 hospital saturation"],
76
+ "missing_considerations": ["compliance decay over 4 ticks"],
77
+ "would_change_mind_if": ["new R3 telemetry shows under-utilisation"],
78
+ "severity": 0.6,
79
+ }
80
+
81
+
82
+ def _valid_json_for(role: str, brain: str = "epidemiology") -> str:
83
+ """Return a valid JSON-schema response for ``role``, brain-substituted."""
84
+ if role == "world_modeler":
85
+ payload = dict(_VALID_BELIEF_PAYLOAD)
86
+ payload["brain"] = brain
87
+ return json.dumps(payload)
88
+ if role == "planner":
89
+ return json.dumps(_VALID_PLAN_PAYLOAD)
90
+ if role == "critic":
91
+ payload = dict(_VALID_CRITIC_PAYLOAD)
92
+ payload["brain"] = brain
93
+ return json.dumps(payload)
94
+ raise ValueError(f"unknown role: {role}")
95
+
96
+
97
+ def _make_subagent_input(
98
+ brain: str = "epidemiology",
99
+ role: str = "world_modeler",
100
+ tick: int = 3,
101
+ round_: int = 1,
102
+ target_plan_id: Optional[str] = None,
103
+ ) -> SubagentInput:
104
+ """Minimal valid SubagentInput. ``round_`` arg avoids shadowing builtin."""
105
+ return SubagentInput(
106
+ brain=brain,
107
+ role=role,
108
+ tick=tick,
109
+ round=round_,
110
+ perception=PerceptionReport(
111
+ brain=brain,
112
+ salient_signals=["R1 cases rising"],
113
+ anomalies=[],
114
+ confidence=0.7,
115
+ evidence=[EvidenceCitation(source="telemetry", ref="R1.cases", excerpt="rising")],
116
+ ),
117
+ prior_belief=None,
118
+ prior_plans=[],
119
+ target_plan_id=target_plan_id,
120
+ last_reward=0.5,
121
+ recent_action_log_excerpt=[],
122
+ )
123
+
124
+
125
+ # ============================================================================
126
+ # T1 - WorldModeler emits BeliefState
127
+ # ============================================================================
128
+
129
+
130
+ def test_world_modeler_emits_belief_state() -> None:
131
+ stub = StubLLMClient(scripted_responses=[_valid_json_for("world_modeler", "epidemiology")])
132
+ agent = WorldModelerSubagent(llm_client=stub)
133
+ input_data = _make_subagent_input(brain="epidemiology", role="world_modeler")
134
+
135
+ result = agent.run(input_data, step_idx=0)
136
+
137
+ assert isinstance(result, BeliefState)
138
+ assert result.brain == "epidemiology"
139
+ assert "R1" in result.latent_estimates
140
+ assert len(result.evidence) >= 1
141
+ assert stub.call_count == 1
142
+
143
+
144
+ # ============================================================================
145
+ # T2 - Planner emits CandidatePlan
146
+ # ============================================================================
147
+
148
+
149
+ def test_planner_emits_candidate_plan() -> None:
150
+ stub = StubLLMClient(scripted_responses=[_valid_json_for("planner", "epidemiology")])
151
+ agent = PlannerSubagent(llm_client=stub)
152
+ input_data = _make_subagent_input(brain="epidemiology", role="planner")
153
+
154
+ result = agent.run(input_data, step_idx=1)
155
+
156
+ assert isinstance(result, CandidatePlan)
157
+ assert result.expected_outer_action.kind == "deploy_resource"
158
+ assert result.confidence == 0.75
159
+
160
+
161
+ # ============================================================================
162
+ # T3 - Critic emits CriticReport
163
+ # ============================================================================
164
+
165
+
166
+ def test_critic_emits_critic_report() -> None:
167
+ stub = StubLLMClient(scripted_responses=[_valid_json_for("critic", "logistics")])
168
+ agent = CriticSubagent(llm_client=stub)
169
+ input_data = _make_subagent_input(brain="logistics", role="critic", target_plan_id="plan-1")
170
+
171
+ result = agent.run(input_data, step_idx=2)
172
+
173
+ assert isinstance(result, CriticReport)
174
+ assert result.brain == "logistics"
175
+ assert result.target_plan_id == "plan-1"
176
+ assert result.severity == 0.6
177
+
178
+
179
+ # ============================================================================
180
+ # T4 - Parse failure then retry succeeds (2 LLM calls)
181
+ # ============================================================================
182
+
183
+
184
+ def test_subagent_parse_failure_then_retry_succeeds() -> None:
185
+ stub = StubLLMClient(
186
+ scripted_responses=["not-json-garbage", _valid_json_for("world_modeler", "epidemiology")]
187
+ )
188
+ agent = WorldModelerSubagent(llm_client=stub)
189
+
190
+ result = agent.run(_make_subagent_input(), step_idx=0)
191
+
192
+ assert isinstance(result, BeliefState)
193
+ assert result.brain == "epidemiology"
194
+ assert stub.call_count == 2, "expected 1 initial call + 1 retry"
195
+
196
+
197
+ # ============================================================================
198
+ # T5 - Parse failure twice -> empty fallback (no third LLM call)
199
+ # ============================================================================
200
+
201
+
202
+ def test_subagent_parse_failure_then_retry_fails_returns_empty() -> None:
203
+ stub = StubLLMClient(scripted_responses=["garbage-1", "garbage-2"])
204
+ agent = WorldModelerSubagent(llm_client=stub)
205
+
206
+ result = agent.run(_make_subagent_input(), step_idx=0)
207
+
208
+ assert isinstance(result, BeliefState)
209
+ assert result.brain == "epidemiology"
210
+ assert result.latent_estimates == {}
211
+ assert result.hypotheses == []
212
+ assert result.evidence == []
213
+ assert result.uncertainty == 1.0
214
+ assert result.reducible_by_more_thought == 0.0
215
+ assert stub.call_count == 2, "must NOT make a 3rd call after retry failure"
216
+
217
+
218
+ # ============================================================================
219
+ # T6 - caller_id format matches Phase A Decision 7
220
+ # ============================================================================
221
+
222
+
223
+ _CALLER_ID_RE = re.compile(
224
+ r"^cortex:(epidemiology|logistics|governance):"
225
+ r"(world_modeler|planner|critic):"
226
+ r"t\d+:r[12]:s\d+$"
227
+ )
228
+
229
+
230
+ @pytest.mark.parametrize(
231
+ "role_cls,brain,role_name",
232
+ [
233
+ (WorldModelerSubagent, "epidemiology", "world_modeler"),
234
+ (PlannerSubagent, "logistics", "planner"),
235
+ (CriticSubagent, "governance", "critic"),
236
+ ],
237
+ )
238
+ def test_subagent_caller_id_format(
239
+ role_cls: type,
240
+ brain: str,
241
+ role_name: str,
242
+ ) -> None:
243
+ stub = StubLLMClient(scripted_responses=[_valid_json_for(role_name, brain)])
244
+ agent = role_cls(llm_client=stub)
245
+ input_data = _make_subagent_input(
246
+ brain=brain,
247
+ role=role_name,
248
+ tick=7,
249
+ round_=2,
250
+ target_plan_id="plan-X" if role_name == "critic" else None,
251
+ )
252
+
253
+ agent.run(input_data, step_idx=4)
254
+
255
+ caller_id = stub.calls[0].caller_id
256
+ assert _CALLER_ID_RE.match(caller_id), (
257
+ f"caller_id={caller_id!r} does not match the locked format"
258
+ )
259
+ assert caller_id == f"cortex:{brain}:{role_name}:t7:r2:s4"
260
+
261
+
262
+ # ============================================================================
263
+ # T8 - SYS prompt loaded from prompts/<role>.txt and brain-formatted
264
+ # (folds in the prompt-formatting refinement: format() must not raise)
265
+ # ============================================================================
266
+
267
+
268
+ @pytest.mark.parametrize(
269
+ "role_cls,role_name",
270
+ [
271
+ (WorldModelerSubagent, "world_modeler"),
272
+ (PlannerSubagent, "planner"),
273
+ (CriticSubagent, "critic"),
274
+ ],
275
+ )
276
+ def test_subagent_uses_loaded_prompt_from_file(role_cls: type, role_name: str) -> None:
277
+ raw = (PROMPTS_DIR / f"{role_name}.txt").read_text(encoding="utf-8")
278
+
279
+ # Refinement: format() must not raise even with extra kwargs ignored.
280
+ # Catches {{/}}-escape regressions in JSON-schema sections of the prompts.
281
+ formatted = raw.format(brain="epidemiology", target_plan_id="plan-X")
282
+ assert isinstance(formatted, str)
283
+ assert "epidemiology" in formatted
284
+
285
+ stub = StubLLMClient(scripted_responses=[_valid_json_for(role_name, "epidemiology")])
286
+ agent = role_cls(llm_client=stub)
287
+ input_data = _make_subagent_input(
288
+ brain="epidemiology",
289
+ role=role_name,
290
+ target_plan_id="plan-X" if role_name == "critic" else None,
291
+ )
292
+
293
+ agent.run(input_data, step_idx=0)
294
+
295
+ sys_msg = stub.calls[0].messages[0]
296
+ assert sys_msg.role == "system"
297
+ assert sys_msg.content == formatted
298
+
299
+
300
+ # ============================================================================
301
+ # T9 - Token counter is billed to the expected caller_id
302
+ # ============================================================================
303
+
304
+
305
+ def test_subagent_token_counter_billed_correctly() -> None:
306
+ stub = StubLLMClient(scripted_responses=[_valid_json_for("world_modeler", "epidemiology")])
307
+ agent = WorldModelerSubagent(llm_client=stub)
308
+ input_data = _make_subagent_input(brain="epidemiology", role="world_modeler", tick=3, round_=1)
309
+
310
+ agent.run(input_data, step_idx=0)
311
+
312
+ expected_caller_id = "cortex:epidemiology:world_modeler:t3:r1:s0"
313
+ assert stub.tokens_used_for(expected_caller_id) > 0, (
314
+ "tokens must be billed to the per-role caller_id, not silently lost"
315
+ )
316
+ assert stub.tokens_used_for("never:called") == 0
317
+
318
+
319
+ # ============================================================================
320
+ # T10 - empty_fallback shape locked per Phase A Decisions 6 + 62
321
+ # ============================================================================
322
+
323
+
324
+ def test_subagent_empty_fallback_shape_locked() -> None:
325
+ # WorldModeler: empty BeliefState
326
+ bs = WorldModelerSubagent.empty_fallback("epidemiology")
327
+ assert isinstance(bs, BeliefState)
328
+ assert bs.brain == "epidemiology"
329
+ assert bs.latent_estimates == {}
330
+ assert bs.hypotheses == []
331
+ assert bs.uncertainty == 1.0
332
+ assert bs.reducible_by_more_thought == 0.0
333
+ assert bs.evidence == []
334
+
335
+ # Planner: empty CandidatePlan with NoOp + confidence=0
336
+ cp = PlannerSubagent.empty_fallback("epidemiology")
337
+ assert isinstance(cp, CandidatePlan)
338
+ assert cp.expected_outer_action.kind == "no_op"
339
+ assert cp.expected_value == 0.0
340
+ assert cp.cost == 0.0
341
+ assert cp.assumptions == []
342
+ assert cp.falsifiers == []
343
+ assert cp.confidence == 0.0
344
+
345
+ # Critic: empty CriticReport with severity=0
346
+ cr = CriticSubagent.empty_fallback("epidemiology", target_plan_id="plan-X")
347
+ assert isinstance(cr, CriticReport)
348
+ assert cr.brain == "epidemiology"
349
+ assert cr.target_plan_id == "plan-X"
350
+ assert cr.attacks == []
351
+ assert cr.missing_considerations == []
352
+ assert cr.would_change_mind_if == []
353
+ assert cr.severity == 0.0
354
+
355
+
356
+ # ============================================================================
357
+ # T11 - retry call uses chat-history continuation (sys + usr + bad + retry)
358
+ # ============================================================================
359
+
360
+
361
+ def test_subagent_run_uses_chat_history_on_retry() -> None:
362
+ stub = StubLLMClient(
363
+ scripted_responses=["bad-json", _valid_json_for("world_modeler", "epidemiology")]
364
+ )
365
+ agent = WorldModelerSubagent(llm_client=stub)
366
+
367
+ agent.run(_make_subagent_input(brain="epidemiology", role="world_modeler"), step_idx=0)
368
+
369
+ assert stub.call_count == 2, "expected 2 LLM calls (initial + retry)"
370
+ call1, call2 = stub.calls
371
+
372
+ # call 1 has the original sys + user (2 messages).
373
+ assert len(call1.messages) == 2
374
+ assert call1.messages[0].role == "system"
375
+ assert call1.messages[1].role == "user"
376
+
377
+ # call 2 must contain: original sys + original user + assistant(bad-json) + retry-user.
378
+ # Without chat-history continuation the LLM loses schema context on retry.
379
+ assert len(call2.messages) == 4, "retry must append to the chat history, not start fresh"
380
+ assert call2.messages[0].role == "system"
381
+ assert call2.messages[0].content == call1.messages[0].content
382
+ assert call2.messages[1].role == "user"
383
+ assert call2.messages[1].content == call1.messages[1].content
384
+ assert call2.messages[2].role == "assistant"
385
+ assert call2.messages[2].content == "bad-json"
386
+ assert call2.messages[3].role == "user"
387
+ assert "failed to parse" in call2.messages[3].content.lower()
tests/test_outer_reward_in_range.py CHANGED
@@ -1,8 +1,10 @@
1
- """Outer reward stays in ``[0.0, 1.0]`` across diverse states.
2
 
3
- Required by ``server/CLAUDE.md`` ("Every grader scalar reward component
4
- lives in ``[0.0, 1.0]``") and the ``tests/CLAUDE.md`` row
5
- ``test_reward_shape.py``. Covers:
 
 
6
 
7
  - Initial state (right after ``load_task``, no ticks applied).
8
  - Mid-episode rollouts on all 3 tasks with varied actions.
@@ -29,7 +31,7 @@ def test_outer_reward_in_range_at_episode_start() -> None:
29
  for name in TASKS:
30
  state = load_task(name, episode_seed=0)
31
  r = outer_reward(state, NoOp())
32
- assert 0.0 <= r <= 1.0, f"{name}: r={r!r} out of [0,1] at tick 0"
33
 
34
 
35
  def test_outer_reward_in_range_during_rollout() -> None:
@@ -51,8 +53,8 @@ def test_outer_reward_in_range_during_rollout() -> None:
51
  for action in actions:
52
  state = apply_tick(state, action)
53
  r = outer_reward(state, action)
54
- assert 0.0 <= r <= 1.0, (
55
- f"{name} tick={state.tick}: r={r!r} out of [0,1] after action kind={action.kind!r}"
56
  )
57
 
58
 
@@ -67,7 +69,7 @@ def test_outer_reward_in_range_with_high_infection() -> None:
67
  region.S, region.E, region.I, region.R = 0.0, 0.0, 0.95, 0.05
68
  state.tick = state.max_ticks # r_time → 0
69
  r = outer_reward(state, NoOp())
70
- assert 0.0 <= r <= 1.0, f"high-I worst case: r={r!r}"
71
 
72
 
73
  def test_outer_reward_in_range_with_rejected_actions() -> None:
@@ -81,7 +83,7 @@ def test_outer_reward_in_range_with_rejected_actions() -> None:
81
  )
82
  state = apply_tick(state, a_v2)
83
  r = outer_reward(state, a_v2)
84
- assert 0.0 <= r <= 1.0, f"V2-rejected: r={r!r}"
85
 
86
  # Legal-violation: strict severity before escalate-national on hard.
87
  state2 = load_task("outbreak_hard", episode_seed=0)
@@ -89,4 +91,4 @@ def test_outer_reward_in_range_with_rejected_actions() -> None:
89
  state2 = apply_tick(state2, a_legal)
90
  assert state2.recent_action_log[-1].accepted is False
91
  r2 = outer_reward(state2, a_legal)
92
- assert 0.0 <= r2 <= 1.0, f"legal-violation: r={r2!r}"
 
1
+ """Outer reward stays in ``[-1.0, 1.0]`` across diverse states.
2
 
3
+ Range relaxed by Workstream-B Phase-1 fix (M2-A): rejected actions land
4
+ ``r_policy = -0.5`` and parse-failure markers land ``r_policy = -1.0``,
5
+ so the per-tick total can go negative. Upper bound stays at 1.0.
6
+
7
+ Covers:
8
 
9
  - Initial state (right after ``load_task``, no ticks applied).
10
  - Mid-episode rollouts on all 3 tasks with varied actions.
 
31
  for name in TASKS:
32
  state = load_task(name, episode_seed=0)
33
  r = outer_reward(state, NoOp())
34
+ assert -1.0 <= r <= 1.0, f"{name}: r={r!r} out of [-1,1] at tick 0"
35
 
36
 
37
  def test_outer_reward_in_range_during_rollout() -> None:
 
53
  for action in actions:
54
  state = apply_tick(state, action)
55
  r = outer_reward(state, action)
56
+ assert -1.0 <= r <= 1.0, (
57
+ f"{name} tick={state.tick}: r={r!r} out of [-1,1] after action kind={action.kind!r}"
58
  )
59
 
60
 
 
69
  region.S, region.E, region.I, region.R = 0.0, 0.0, 0.95, 0.05
70
  state.tick = state.max_ticks # r_time → 0
71
  r = outer_reward(state, NoOp())
72
+ assert -1.0 <= r <= 1.0, f"high-I worst case: r={r!r}"
73
 
74
 
75
  def test_outer_reward_in_range_with_rejected_actions() -> None:
 
83
  )
84
  state = apply_tick(state, a_v2)
85
  r = outer_reward(state, a_v2)
86
+ assert -1.0 <= r <= 1.0, f"V2-rejected: r={r!r}"
87
 
88
  # Legal-violation: strict severity before escalate-national on hard.
89
  state2 = load_task("outbreak_hard", episode_seed=0)
 
91
  state2 = apply_tick(state2, a_legal)
92
  assert state2.recent_action_log[-1].accepted is False
93
  r2 = outer_reward(state2, a_legal)
94
+ assert -1.0 <= r2 <= 1.0, f"legal-violation: r={r2!r}"
tests/test_outer_reward_non_constancy.py CHANGED
@@ -15,6 +15,11 @@ from CrisisWorldCortex.models import (
15
  RestrictMovement,
16
  )
17
  from CrisisWorldCortex.server.graders import outer_reward
 
 
 
 
 
18
  from CrisisWorldCortex.server.simulator import apply_tick, load_task
19
 
20
 
@@ -87,8 +92,12 @@ def test_reward_differs_for_accepted_vs_rejected_action() -> None:
87
 
88
  Compare two episodes from the same starting state: one issues NoOp
89
  (accepted), the other issues PublicCommunication (V2-rejected). The
90
- SEIR step runs identically; only ``r_policy`` differs. With weight
91
- 0.12, the gap should be exactly 0.12 (modulo float rounding).
 
 
 
 
92
  """
93
  s_a = load_task("outbreak_easy", episode_seed=42)
94
  s_b = load_task("outbreak_easy", episode_seed=42)
@@ -109,7 +118,7 @@ def test_reward_differs_for_accepted_vs_rejected_action() -> None:
109
  assert r_ok > r_bad, (
110
  f"accepted action should score higher than rejected: ok={r_ok!r} bad={r_bad!r}"
111
  )
112
- # Exact gap = 0.12 because the only differing component is r_policy.
113
- assert abs((r_ok - r_bad) - 0.12) < 1e-9, (
114
- f"r_policy gap mismatch: ok-bad={r_ok - r_bad!r}, expected 0.12"
115
  )
 
15
  RestrictMovement,
16
  )
17
  from CrisisWorldCortex.server.graders import outer_reward
18
+ from CrisisWorldCortex.server.graders.outer_reward import (
19
+ R_POLICY_NOOP_ACCEPTED,
20
+ R_POLICY_REJECTED,
21
+ W_POLICY,
22
+ )
23
  from CrisisWorldCortex.server.simulator import apply_tick, load_task
24
 
25
 
 
92
 
93
  Compare two episodes from the same starting state: one issues NoOp
94
  (accepted), the other issues PublicCommunication (V2-rejected). The
95
+ SEIR step runs identically; only ``r_policy`` differs. After the
96
+ Workstream-B Phase-1 four-state contract, NoOp(accepted)
97
+ R_POLICY_NOOP_ACCEPTED (0.0) and PublicCommunication(honesty=0.9,
98
+ rejected as legal-violation) → R_POLICY_REJECTED (-0.5). The exact
99
+ gap is therefore ``(R_POLICY_NOOP_ACCEPTED - R_POLICY_REJECTED) *
100
+ W_POLICY`` (modulo float rounding).
101
  """
102
  s_a = load_task("outbreak_easy", episode_seed=42)
103
  s_b = load_task("outbreak_easy", episode_seed=42)
 
118
  assert r_ok > r_bad, (
119
  f"accepted action should score higher than rejected: ok={r_ok!r} bad={r_bad!r}"
120
  )
121
+ expected_gap = (R_POLICY_NOOP_ACCEPTED - R_POLICY_REJECTED) * W_POLICY
122
+ assert abs((r_ok - r_bad) - expected_gap) < 1e-9, (
123
+ f"r_policy gap mismatch: ok-bad={r_ok - r_bad!r}, expected {expected_gap!r}"
124
  )