feat: enhance inference and logging capabilities with SWD tracing

Updated .env.example to include new environment variables for master and worker agents, and added support for SWD tracing in inference.py. The SwdTraceWriter class was introduced to log SWD snapshots to a specified file, improving the logging mechanism. Adjusted README.md to reflect changes in API key requirements and SWD tracing options.

.env.example

@@ -1,17 +1,59 @@
-# Master / judge LLM (OpenAI-compatible)
+# ---------------------------------------------------------------------------
+# Global fallbacks (OpenAI-compatible). Used when role-specific vars are unset.
+# ---------------------------------------------------------------------------
 HF_TOKEN=
 OPENAI_API_KEY=
 API_BASE_URL=https://router.huggingface.co/v1
 MODEL_NAME=Qwen/Qwen2.5-72B-Instruct
 
-# Episode selection (e1_launch_readiness | m1_budget_reallocation | h1_acquisition_defence)
+# ---------------------------------------------------------------------------
+# Master agent (inference.py) — planner that reads/writes the SWD
+# ---------------------------------------------------------------------------
+CORP_MASTER_API_KEY=
+CORP_MASTER_BASE_URL=
+CORP_MASTER_MODEL=
+
+# ---------------------------------------------------------------------------
+# Frozen workers (delegate) — optional per-agent key/router/model
+# Naming: CORP_WORKER_<AGENT_ID_UPPER>_API_KEY / _BASE_URL / _MODEL
+# Example for dev_agent -> CORP_WORKER_DEV_AGENT_API_KEY
+# ---------------------------------------------------------------------------
+CORP_WORKER_DEFAULT_API_KEY=
+CORP_WORKER_DEFAULT_BASE_URL=
+CORP_WORKER_DEFAULT_MODEL=
+
+CORP_WORKER_DEV_AGENT_API_KEY=
+CORP_WORKER_DEV_AGENT_BASE_URL=
+CORP_WORKER_DEV_AGENT_MODEL=
+
+CORP_WORKER_HR_AGENT_API_KEY=
+CORP_WORKER_HR_AGENT_BASE_URL=
+CORP_WORKER_HR_AGENT_MODEL=
+
+CORP_WORKER_FINANCE_AGENT_API_KEY=
+CORP_WORKER_FINANCE_AGENT_BASE_URL=
+CORP_WORKER_FINANCE_AGENT_MODEL=
+
+CORP_WORKER_MODEL=
+
+# ---------------------------------------------------------------------------
+# LLM judge (server/reward.py) — separate endpoint from master/workers
+# ---------------------------------------------------------------------------
+CORP_JUDGE_API_KEY=
+CORP_JUDGE_BASE_URL=
+CORP_JUDGE_MODEL=Qwen/Qwen2.5-7B-Instruct
+CORP_DISABLE_LLM_JUDGE=1
+
+# ---------------------------------------------------------------------------
+# Episode + inference logging
+# ---------------------------------------------------------------------------
 CORP_TASK_ID=e1_launch_readiness
 
-# Worker + judge helpers
+# Append-only SWD evolution log (separate from console). Use .jsonl for one JSON
+# object per line, or .txt for human-readable JSON blocks.
+CORP_SWD_TRACE_FILE=logs/swd_trace.jsonl
+
 CORP_STUB_WORKERS=1
-CORP_DISABLE_LLM_JUDGE=1
-CORP_WORKER_MODEL=Qwen/Qwen2.5-7B-Instruct
-CORP_JUDGE_MODEL=Qwen/Qwen2.5-7B-Instruct
 
 # Hugging Face Space
 PORT=7860

.gitignore

@@ -51,3 +51,4 @@ task.md
 walkthrough.md
 72b_eval.txt
 .cursor/plans/corp-env_rewrite_plan_952c3fcd.plan.md
+_test_swd.jsonl

README.md

@@ -53,11 +53,11 @@ uv run server
 
 ## Baseline inference (master agent)
 
-Requires `HF_TOKEN` or `OPENAI_API_KEY` for the **master** model. Without a key, `inference.py` runs a short **deterministic E1** smoke test using stub workers.
+Requires a **master** API key (`CORP_MASTER_API_KEY`, or `HF_TOKEN` / `OPENAI_API_KEY` as fallback). Without it, `inference.py` runs a short **deterministic E1** smoke test using stub workers. Optional **per-worker** and **judge** keys/URLs are in [`.env.example`](.env.example). Set `CORP_SWD_TRACE_FILE` or pass `--swd-trace path.jsonl` to append SWD snapshots to a file separate from console logs.
 
 ```powershell
 uv run python inference.py
-uv run python inference.py --tasks e1_launch_readiness --max-steps 25
+uv run python inference.py --tasks e1_launch_readiness --max-steps 25 --swd-trace logs/run.jsonl
 ```
 
 ## OpenEnv validation
@@ -79,7 +79,7 @@ docker run -p 7860:7860 --env-file .env.example corp-env
 
 ## Configuration
 
-See [`.env.example`](.env.example) for `CORP_TASK_ID`, `CORP_STUB_WORKERS`, judge toggles, and API routing.
+See [`.env.example`](.env.example) for master/worker/judge API routing, `CORP_TASK_ID`, `CORP_STUB_WORKERS`, and `CORP_SWD_TRACE_FILE`.
 
 ## License
 

inference.py

@@ -10,19 +10,22 @@ import os
 import re
 import textwrap
 import time
-from typing import List, Optional
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 from dotenv import load_dotenv
 from openai import OpenAI
 
 from corp_env.models import CorpAction, CorpObservation
 from server.environment import CorpEnvironment
+from server.llm_env import openai_client_kwargs_master
 
 load_dotenv()
 
-API_BASE_URL = os.getenv("API_BASE_URL")
-MODEL_NAME = os.getenv("MODEL_NAME")
-HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY")
+MASTER_KWARGS = openai_client_kwargs_master()
+MASTER_API_KEY = MASTER_KWARGS.get("api_key")
+MODEL_NAME = os.getenv("CORP_MASTER_MODEL") or os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
 
 BENCHMARK = "corp-env"
 MAX_HISTORY_MESSAGES = 40
@@ -78,6 +81,66 @@ def log_end(task: str, steps: int, score: float, rewards: List[float]) -> None:
     print(f"[END] task={task} steps={steps} score={score:.3f} rewards={rs}", flush=True)
 
 
+class SwdTraceWriter:
+    """Append SWD snapshots to a dedicated file (not mixed with console logs)."""
+
+    def __init__(self, path: Optional[str], task_id: str) -> None:
+        self.path = path.strip() if path else None
+        self.task_id = task_id
+        self._jsonl = bool(self.path and self.path.lower().endswith(".jsonl"))
+        if not self.path:
+            return
+        p = Path(self.path)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H-%M-%SZ")
+        with p.open("a", encoding="utf-8") as f:
+            f.write(
+                f"\n{'=' * 72}\n"
+                f"# CORP-ENV SWD trace | task={task_id} | started_utc={ts}\n"
+                f"{'=' * 72}\n"
+            )
+
+    def write(
+        self,
+        *,
+        phase: str,
+        step_index: int,
+        action: Optional[CorpAction],
+        obs: CorpObservation,
+    ) -> None:
+        if not self.path:
+            return
+        action_blob: Dict[str, Any]
+        if action is None:
+            action_blob = {"note": "initial observation after reset"}
+        else:
+            action_blob = action.model_dump(mode="json", exclude_none=True)
+
+        if self._jsonl:
+            record = {
+                "phase": phase,
+                "step_index": step_index,
+                "env_turn": obs.turn,
+                "reward": obs.reward,
+                "done": obs.done,
+                "error": obs.error,
+                "action": action_blob,
+                "swd": obs.swd,
+            }
+            line = json.dumps(record, ensure_ascii=False)
+            with Path(self.path).open("a", encoding="utf-8") as f:
+                f.write(line + "\n")
+            return
+
+        with Path(self.path).open("a", encoding="utf-8") as f:
+            f.write(
+                f"\n--- {phase} step_index={step_index} env_turn={obs.turn} "
+                f"reward={obs.reward} done={obs.done} ---\n"
+            )
+            f.write(f"action: {json.dumps(action_blob, indent=2, ensure_ascii=False)}\n")
+            f.write(f"swd:\n{json.dumps(obs.swd, indent=2, ensure_ascii=False)}\n")
+
+
 def extract_json(raw_text: str) -> dict:
     cleaned = raw_text.strip()
     cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned)
@@ -146,7 +209,12 @@ def trim_history(messages: list, max_messages: int = MAX_HISTORY_MESSAGES) -> No
         messages.pop(1)
 
 
-def run_episode(client: OpenAI, task_id: str, max_steps: int) -> tuple[float, int, List[float]]:
+def run_episode(
+    client: OpenAI,
+    task_id: str,
+    max_steps: int,
+    swd_trace: Optional[SwdTraceWriter],
+) -> tuple[float, int, List[float]]:
     os.environ["CORP_TASK_ID"] = task_id
     os.environ.setdefault("CORP_STUB_WORKERS", "1")
 
@@ -157,6 +225,8 @@ def run_episode(client: OpenAI, task_id: str, max_steps: int) -> tuple[float, in
 
     log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
     obs = env.reset(task_id=task_id)
+    if swd_trace:
+        swd_trace.write(phase="after_reset", step_index=0, action=None, obs=obs)
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": build_observation_message(0, obs)},
@@ -208,6 +278,8 @@ def run_episode(client: OpenAI, task_id: str, max_steps: int) -> tuple[float, in
         total += float(obs.reward or 0.0)
         steps = step
         log_step(step, alog[:200], float(obs.reward or 0.0), obs.done, obs.error)
+        if swd_trace:
+            swd_trace.write(phase="after_step", step_index=step, action=action, obs=obs)
         messages.append({"role": "user", "content": build_observation_message(step, obs)})
         if obs.done:
             break
@@ -216,12 +288,14 @@ def run_episode(client: OpenAI, task_id: str, max_steps: int) -> tuple[float, in
     return total, steps, rewards
 
 
-def deterministic_e1_smoke() -> None:
+def deterministic_e1_smoke(swd_trace: Optional[SwdTraceWriter] = None) -> None:
     """Offline smoke: E1 solved with stub workers (no master LLM)."""
     os.environ["CORP_TASK_ID"] = "e1_launch_readiness"
     os.environ["CORP_STUB_WORKERS"] = "1"
     env = CorpEnvironment()
     obs = env.reset(task_id="e1_launch_readiness")
+    if swd_trace:
+        swd_trace.write(phase="after_reset", step_index=0, action=None, obs=obs)
     seq = [
         CorpAction(action_type="delegate", agent_id="dev_agent", payload="Assess launch readiness"),
         CorpAction(action_type="delegate", agent_id="hr_agent", payload="Staffing sign-off"),
@@ -241,6 +315,8 @@ def deterministic_e1_smoke() -> None:
         total += r
         rlist.append(r)
         log_step(i, act.action_type, r, obs.done, obs.error)
+        if swd_trace:
+            swd_trace.write(phase="after_step", step_index=i, action=act, obs=obs)
     log_end("e1_launch_readiness", len(seq), total, rlist)
 
 
@@ -253,21 +329,31 @@ def main() -> None:
         help="Comma-separated task ids",
     )
     parser.add_argument("--max-steps", type=int, default=30, help="Max steps per episode")
+    parser.add_argument(
+        "--swd-trace",
+        type=str,
+        default=os.getenv("CORP_SWD_TRACE_FILE", ""),
+        help="Append SWD evolution to this file (.jsonl recommended). Overrides CORP_SWD_TRACE_FILE.",
+    )
     args = parser.parse_args()
 
-    if not HF_TOKEN:
+    trace_path = (args.swd_trace or "").strip() or None
+
+    if not MASTER_API_KEY:
         print(
-            "No HF_TOKEN / OPENAI_API_KEY - running deterministic E1 smoke only. "
-            "Set keys to run the LLM master on --tasks.",
+            "No master API key (set CORP_MASTER_API_KEY or HF_TOKEN / OPENAI_API_KEY) - "
+            "running deterministic E1 smoke only. Set keys to run the LLM master on --tasks.",
             flush=True,
         )
-        deterministic_e1_smoke()
+        tw = SwdTraceWriter(trace_path, "e1_launch_readiness") if trace_path else None
+        deterministic_e1_smoke(swd_trace=tw)
         return
 
-    client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL or None)
+    client = OpenAI(**MASTER_KWARGS)
     for tid in [t.strip() for t in args.tasks.split(",") if t.strip()]:
         ms = args.max_steps * 2 if tid == "h1_acquisition_defence" else args.max_steps
-        run_episode(client, tid, max_steps=ms)
+        tw = SwdTraceWriter(trace_path, tid) if trace_path else None
+        run_episode(client, tid, max_steps=ms, swd_trace=tw)
 
 
 if __name__ == "__main__":

server/llm_env.py

@@ -0,0 +1,88 @@
+"""Resolve OpenAI-compatible API key and base URL per role (master, worker, judge)."""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Optional
+
+
+def _first(*values: Optional[str]) -> Optional[str]:
+    for v in values:
+        if v is not None and str(v).strip() != "":
+            return str(v).strip()
+    return None
+
+
+def openai_client_kwargs_master() -> Dict[str, Any]:
+    """Credentials for the master model (inference loop)."""
+    api_key = _first(
+        os.getenv("CORP_MASTER_API_KEY"),
+        os.getenv("HF_TOKEN"),
+        os.getenv("OPENAI_API_KEY"),
+        os.getenv("API_KEY"),
+    )
+    base_url = _first(
+        os.getenv("CORP_MASTER_BASE_URL"),
+        os.getenv("API_BASE_URL"),
+        os.getenv("OPENAI_BASE_URL"),
+    )
+    return _kwargs(api_key, base_url)
+
+
+def openai_client_kwargs_worker(canonical_agent_id: str) -> Dict[str, Any]:
+    """
+    Credentials for a frozen worker (dev_agent, hr_agent, finance_agent).
+
+    Per-agent overrides use uppercase id with hyphens as underscores, e.g.:
+    CORP_WORKER_DEV_AGENT_API_KEY, CORP_WORKER_DEV_AGENT_BASE_URL
+    """
+    suffix = canonical_agent_id.upper().replace("-", "_")
+    api_key = _first(
+        os.getenv(f"CORP_WORKER_{suffix}_API_KEY"),
+        os.getenv("CORP_WORKER_DEFAULT_API_KEY"),
+        os.getenv("OPENAI_API_KEY"),
+        os.getenv("HF_TOKEN"),
+        os.getenv("API_KEY"),
+    )
+    base_url = _first(
+        os.getenv(f"CORP_WORKER_{suffix}_BASE_URL"),
+        os.getenv("CORP_WORKER_DEFAULT_BASE_URL"),
+        os.getenv("API_BASE_URL"),
+        os.getenv("OPENAI_BASE_URL"),
+    )
+    return _kwargs(api_key, base_url)
+
+
+def openai_client_kwargs_judge() -> Dict[str, Any]:
+    """Credentials for the optional LLM judge (reward)."""
+    api_key = _first(
+        os.getenv("CORP_JUDGE_API_KEY"),
+        os.getenv("OPENAI_API_KEY"),
+        os.getenv("HF_TOKEN"),
+        os.getenv("API_KEY"),
+    )
+    base_url = _first(
+        os.getenv("CORP_JUDGE_BASE_URL"),
+        os.getenv("API_BASE_URL"),
+        os.getenv("OPENAI_BASE_URL"),
+    )
+    return _kwargs(api_key, base_url)
+
+
+def _kwargs(api_key: Optional[str], base_url: Optional[str]) -> Dict[str, Any]:
+    out: Dict[str, Any] = {}
+    if api_key:
+        out["api_key"] = api_key
+    if base_url:
+        out["base_url"] = base_url
+    return out
+
+
+def worker_model_for(canonical_agent_id: str) -> str:
+    suffix = canonical_agent_id.upper().replace("-", "_")
+    return _first(
+        os.getenv(f"CORP_WORKER_{suffix}_MODEL"),
+        os.getenv("CORP_WORKER_DEFAULT_MODEL"),
+        os.getenv("CORP_WORKER_MODEL"),
+        os.getenv("MODEL_NAME"),
+    ) or "Qwen/Qwen2.5-7B-Instruct"

server/reward.py

@@ -9,6 +9,7 @@ from typing import Any, Callable, Dict, List, Optional
 
 from openai import OpenAI
 
+from server.llm_env import openai_client_kwargs_judge
 from server.swd import (
     REQUIRED_TOP_LEVEL,
     VALID_PHASES,
@@ -51,16 +52,15 @@ def compute_swd_coherence(swd: Dict[str, Any]) -> float:
 def call_llm_judge(swd: Dict[str, Any], task_goal: str) -> float:
     """
     Fast LLM judge (optional). Returns score in [0, 1] from YES count / 3.
-    Disabled when OPENAI_API_KEY / HF_TOKEN missing unless CORP_FORCE_JUDGE=1.
+    Uses CORP_JUDGE_* then global API keys (see server/llm_env.py). No call without a key.
     """
     if os.getenv("CORP_DISABLE_LLM_JUDGE", "").lower() in ("1", "true", "yes"):
         return 0.0
 
-    api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") or os.getenv("API_KEY")
-    if not api_key and os.getenv("CORP_FORCE_JUDGE", "").lower() not in ("1", "true"):
+    kwargs = openai_client_kwargs_judge()
+    if not kwargs.get("api_key"):
         return 0.0
 
-    base_url = os.getenv("API_BASE_URL") or os.getenv("OPENAI_BASE_URL")
     model = os.getenv("CORP_JUDGE_MODEL", "Qwen/Qwen2.5-7B-Instruct")
 
     prompt = f"""
@@ -82,9 +82,6 @@ Q1: YES/NO
 Q2: YES/NO
 Q3: YES/NO
 """
-    kwargs: Dict[str, Any] = {"api_key": api_key}
-    if base_url:
-        kwargs["base_url"] = base_url
     client = OpenAI(**kwargs)
     resp = client.chat.completions.create(
         model=model,

server/worker_client.py

@@ -8,6 +8,7 @@ from typing import Optional
 from openai import OpenAI
 
 from server.agents.prompts import WORKER_PROMPTS
+from server.llm_env import openai_client_kwargs_worker, worker_model_for
 
 STUB_OUTPUTS = {
     "dev_agent": (
@@ -44,19 +45,15 @@ def call_worker_model(
     if os.getenv("CORP_STUB_WORKERS", "").lower() in ("1", "true", "yes"):
         return call_model_stub(canonical_agent_id, task_description)
 
-    api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") or os.getenv("API_KEY")
-    if not api_key:
+    kwargs = openai_client_kwargs_worker(canonical_agent_id)
+    if not kwargs.get("api_key"):
         return call_model_stub(canonical_agent_id, task_description)
 
-    base_url = os.getenv("API_BASE_URL") or os.getenv("OPENAI_BASE_URL")
-    model = os.getenv("CORP_WORKER_MODEL") or os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-7B-Instruct"
+    model = worker_model_for(canonical_agent_id)
     system = WORKER_PROMPTS.get(
         canonical_agent_id,
         "You are a concise corporate advisor. Plain prose only.",
     )
-    kwargs = {"api_key": api_key}
-    if base_url:
-        kwargs["base_url"] = base_url
     client = OpenAI(**kwargs)
     resp = client.chat.completions.create(
         model=model,