nation_optimizer / llm_integration /context_builder.py
Kr0issant
(train) pre-test train
f72fde9
Raw
History Blame Contribute Delete
3.91 kB
"""
Logic to convert the raw NationEnvironment observation into a clean LLM context dictionary.
"""
from dataclasses import asdict, fields, is_dataclass
from typing import Any, Mapping
from schemas.observations import Observation
def _to_dict(obj: Any) -> Any:
"""Serialize a Pydantic model or a frozen dataclass to a plain dict."""
if hasattr(obj, "model_dump"):
return obj.model_dump()
if is_dataclass(obj) and not isinstance(obj, type):
return asdict(obj)
return obj
HIDDEN_EVENT_FIELDS = frozenset(
{
"cost",
"exact_cost",
"base_cost",
"base_cost_impact",
"severity_multiplier",
"random_variance",
}
)
PUBLIC_SECTOR_FIELDS = ("critical", "demand", "surplus", "wastage")
def build_public_observation(observation: Any) -> dict[str, Any]:
"""Builds a public-facing observation for ministers."""
phase_str = getattr(
observation,
"phase_name",
getattr(getattr(observation, "phase", None), "name", "UNKNOWN"),
)
target_id = getattr(observation, "target_proposal_id", None)
all_proposals = [_to_dict(p) for p in observation.proposals]
if phase_str == "VOTING" and target_id:
visible_proposals = [
p for p in all_proposals if p.get("proposal_id") == target_id
]
else:
visible_proposals = all_proposals
return {
"round": observation.round,
"phase": phase_str,
"treasury": observation.treasury,
"total_critical": getattr(observation, "total_critical", 0.0),
"max_rounds": getattr(observation, "max_rounds", 0),
"event_ledger": [_sanitize_event(event) for event in observation.event_ledger],
"proposals": visible_proposals,
"votes": [_to_dict(v) for v in observation.votes],
"debate_messages": list(observation.debate_messages),
"own_department": _to_dict(observation.own_department) if observation.own_department else None,
"target_proposal_id": target_id,
"termination": dict(observation.termination) if observation.termination else {},
}
def build_oracle_observation(observation: Observation) -> dict[str, Any]:
"""Builds an oracle observation for the dictator (sees private metrics and event costs)."""
return {
**build_public_observation(observation),
"oracle_own_department": _to_dict(observation.own_department) if observation.own_department else None,
"event_ledger": [dict(event) for event in observation.event_ledger],
}
def build_sector_thresholds(state: Mapping[str, Any]) -> dict[str, dict[str, float]]:
"""Return public per-sector ``(critical, demand, surplus)`` thresholds.
The reward function uses these to score proposed allocations against the
same piecewise revenue curve the engine evaluates during the budget
execution phase. Only public fields are exposed; hidden event costs and
private metrics are deliberately omitted.
"""
sectors = state.get("sectors") or {}
if not isinstance(sectors, Mapping):
raise TypeError("state['sectors'] must be a mapping of sector name to sector dict.")
thresholds: dict[str, dict[str, float]] = {}
for name, sector in sectors.items():
if not isinstance(sector, Mapping):
raise TypeError(f"sector entry for {name!r} must be a mapping.")
thresholds[str(name)] = {
field: float(sector[field])
for field in PUBLIC_SECTOR_FIELDS
}
return thresholds
def _sanitize_event(event: dict[str, Any]) -> dict[str, Any]:
"""Removes hidden fields from events unless cost is already public."""
if event.get("cost") is not None:
return {str(key): value for key, value in event.items()}
return {
str(key): value
for key, value in event.items()
if str(key).lower() not in HIDDEN_EVENT_FIELDS
}