Fix issues from repo review

AGENTS.md

@@ -19,7 +19,13 @@ messages += [{"role": "assistant", "content": result.raw_output}]
 messages += result.tool_messages()        # tool results back to the model
 ```
 
-The agent harness calls `agent.step(messages)` each turn instead of `model.generate()`, then appends the assistant output and `result.tool_messages()` (role `tool`) so the model sees whether its calls worked. If there are tool messages, give the model the next turn. Steering is transparent — the harness doesn't need to know about hooks or vectors.
+Install it into the harness's environment with `pip install -e /path/to/modulo`. Keep the install editable so the default `vectors/` directory still resolves to the repo (or set `MODULO_VECTORS_DIR`).
+
+The agent harness calls `agent.step(messages)` each turn instead of `model.generate()`, then appends the assistant output and `result.tool_messages()` (role `tool`) so the model sees whether its calls worked. If there are tool messages, give the model the next turn. If `messages` starts with a system message, it is merged into the steering system prompt. Steering is transparent — the harness doesn't need to know about hooks or vectors.
+
+`agent.apply()` / `agent.drop()` called between steps behave like the model's own tool calls: an applied vector gets one full-strength steered turn, and a drop blocks re-applying for `COOLDOWN_TURNS` turns.
+
+Chat templates must accept `system` and `tool` roles (Qwen, Llama-3.x, recent Mistral do). Gemma-2 rejects both, so it isn't supported yet.
 
 ### 2. As a standalone demo
 

agent/loop.py

@@ -71,7 +71,8 @@ def run_turn(agent: SteeringAgent, messages: list, logger: SessionLogger, turn:
         else:
             logger.log({"type": "tool_error", "turn": turn, "name": r.name, "message": r.message})
     if result.steering_before != result.steering_after:
-        logger.log({"type": "decay", "turn": turn, "active_vectors": result.steering_after})
+        # Steering changed this turn (steer/drop calls, decay, or expiry)
+        logger.log({"type": "steering_state", "turn": turn, "active_vectors": result.steering_after})
     return result
 
 
@@ -90,8 +91,10 @@ def main():
     messages = []
     if args.seed:
         messages.append({"role": "user", "content": args.seed})
+        logger.log_turn(1, "user", args.seed)
         print(f"\n[seed] {args.seed}\n")
 
+    # Inputs (user, tool results, nudges) are logged under the turn they feed into
     turn = 0
     while turn < args.max_turns and not kill.killed:
         if not messages or messages[-1]["role"] not in ("user", "tool"):
@@ -105,7 +108,7 @@ def main():
             if not user_input or user_input.lower() in ("quit", "exit"):
                 break
             messages.append({"role": "user", "content": user_input})
-            logger.log_turn(turn, "user", user_input)
+            logger.log_turn(turn + 1, "user", user_input)
 
         print(f"\n--- Turn {turn+1} ---")
         print(f"Steering status:\n{agent.status()}")
@@ -119,12 +122,12 @@ def main():
             # Tool results go straight back to the model, which gets the next turn
             for m in tool_messages:
                 print(f"[tool] {m['content']}")
-                logger.log_turn(turn, "tool", m["content"])
+                logger.log_turn(turn + 1, "tool", m["content"])
             messages.extend(tool_messages)
-        elif args.auto:
+        elif args.auto and turn < args.max_turns and not kill.killed:  # no nudge after the last turn
             nudge = "(You can continue experimenting with steer if you like, or just talk.)"
             messages.append({"role": "user", "content": nudge})
-            logger.log_turn(turn, "user", nudge)
+            logger.log_turn(turn + 1, "user", nudge)
 
     print(f"\n[session ended after {turn} turns]")
     print(f"[log saved to {logger.log_path}]")

config.py

@@ -9,22 +9,15 @@
 
 # Qwen2.5-7B has 28 decoder layers (indices 0..27). Mid layers are good for steering.
 # CAA on Llama-2-7B (32 layers) used layer 13; for Qwen2.5-7B we start at layer 13.
-NUM_LAYERS = 28
 DEFAULT_STEER_LAYER = 13
 LAYER_SWEEP = [10, 13, 16, 20]  # tried if default doesn't shift behavior
 
-# Steering
-DEFAULT_MULTIPLIER = 2.0
-MULTIPLIER_SWEEP = [-2.0, -1.0, 0.0, 1.0, 2.0]
-
 # Paths (resolved relative to project root)
 PROJECT_ROOT = Path(__file__).resolve().parent
-DATA_DIR = PROJECT_ROOT / "data"
 VECTORS_DIR = PROJECT_ROOT / "vectors"
 LOGS_DIR = PROJECT_ROOT / "logs"
-AGENT_DIR = PROJECT_ROOT / "agent"
 
-for d in (DATA_DIR, VECTORS_DIR, LOGS_DIR, AGENT_DIR):
+for d in (VECTORS_DIR, LOGS_DIR):
     d.mkdir(parents=True, exist_ok=True)
 
 # HF cache on the pod's large /workspace disk — override via env if needed.

modulo/agent.py

@@ -12,6 +12,7 @@
     # agent.status() -> active vector status string
 """
 import json
+import math
 import re
 from collections import Counter
 from contextlib import nullcontext
@@ -293,6 +294,9 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None,
         self.active: dict[str, ActiveVector] = {}
         self.turn = 0
         self._dropped_recently: dict[str, int] = {}
+        # True between this turn's generation and decay(). apply()/drop() calls
+        # before generation (manual control) count as the end of the previous turn.
+        self._turn_generated = False
         self._unit_vectors: dict[str, torch.Tensor] = {}
         self._residual_norms: dict[int, float] = {}
 
@@ -347,6 +351,7 @@ def apply(self, name: str, strength: float) -> tuple[bool, str]:
             return False, f"Too many active vectors (max {MAX_ACTIVE_VECTORS}). Drop one first."
         self.active[name] = ActiveVector(
             entry=entry, base_strength=strength, current_strength=strength,
+            turns_active=0 if self._turn_generated else 1,
             last_applied_turn=self.turn,
         )
         sign = "+" if strength >= 0 else ""
@@ -357,11 +362,12 @@ def drop(self, name: str) -> tuple[bool, str]:
         if name not in self.active:
             return False, f"'{name}' is not active."
         del self.active[name]
-        self._dropped_recently[name] = self.turn
+        self._dropped_recently[name] = self.turn if self._turn_generated else self.turn - 1
         return True, f"Dropped '{name}'. On cooldown for the next {COOLDOWN_TURNS} turns."
 
     def decay(self):
         self.turn += 1
+        self._turn_generated = False
         expired = []
         for name, av in self.active.items():
             av.turns_active += 1
@@ -377,7 +383,10 @@ def decay(self):
     def _unit_vector(self, entry: VectorEntry) -> torch.Tensor:
         if entry.name not in self._unit_vectors:
             v = entry.load_vector(self.registry.vectors_dir).float()
-            self._unit_vectors[entry.name] = v / v.norm()
+            norm = v.norm()
+            if norm == 0 or not torch.isfinite(norm):
+                raise ValueError(f"Vector '{entry.name}' has zero or non-finite norm; re-extract it.")
+            self._unit_vectors[entry.name] = v / norm
         return self._unit_vectors[entry.name]
 
     def residual_norm(self, layer: int) -> float:
@@ -426,7 +435,7 @@ def _execute(self, tc: ToolCall) -> ToolResult:
                 strength = float(args.get("strength", 0))
             except (TypeError, ValueError):
                 return ToolResult(tc.name, args, False, f"'strength' must be a number, got {args.get('strength')!r}.")
-            if strength != strength:  # NaN
+            if not math.isfinite(strength):
                 return ToolResult(tc.name, args, False, "'strength' must be a finite number.")
             ok, msg = self.apply(concept, strength)
         elif tc.name == "drop":
@@ -450,6 +459,20 @@ def _generate(self, inputs, max_new_tokens, temperature, top_p):
             )
         return self.tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
 
+    def _full_messages(self, messages: list[dict]) -> list[dict]:
+        sys_prompt = self.build_system_prompt()
+        if messages and messages[0].get("role") == "system":
+            # Merge the harness's own system prompt; many chat templates reject
+            # or mishandle a second system message.
+            sys_prompt = messages[0]["content"] + "\n\n" + sys_prompt
+            messages = messages[1:]
+        return [{"role": "system", "content": sys_prompt}] + messages
+
+    def _encode(self, messages: list[dict]):
+        text = self.tokenizer.apply_chat_template(
+            self._full_messages(messages), tokenize=False, add_generation_prompt=True)
+        return self.tokenizer([text], return_tensors="pt").to(self.model.device)
+
     def step(self, messages: list[dict], max_new_tokens: int = 300,
              temperature: float = 0.7, top_p: float = 0.9) -> StepResult:
         """Run one turn: build prompt, generate (steered), parse tool calls, apply them, decay.
@@ -460,24 +483,23 @@ def step(self, messages: list[dict], max_new_tokens: int = 300,
 
         Returns: StepResult with prose, raw output, tool calls/results, and steering state
         """
-        sys_prompt = self.build_system_prompt()
-        full_messages = [{"role": "system", "content": sys_prompt}] + messages
-        text = self.tokenizer.apply_chat_template(full_messages, tokenize=False, add_generation_prompt=True)
-        inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
+        inputs = self._encode(messages)
         prompt_len = inputs.input_ids.shape[1]
 
         steering_before = {n: av.current_strength for n, av in self.active.items()}
 
         with self._build_steering_context(prompt_len):
             raw_output = self._generate(inputs, max_new_tokens, temperature, top_p)
+        self._turn_generated = True
 
         # Degeneration check (only meaningful if steering was applied)
         degenerate_output = None
         if steering_before and is_degenerate(strip_tool_calls(raw_output)):
             degenerate_output = raw_output
             for name in list(self.active.keys()):
                 self.drop(name)
-            raw_output = self._generate(inputs, max_new_tokens, temperature, top_p)
+            # Re-encode so the regenerated turn's system prompt shows the drops
+            raw_output = self._generate(self._encode(messages), max_new_tokens, temperature, top_p)
 
         tool_calls = parse_tool_calls(raw_output)
         prose = strip_tool_calls(raw_output)

modulo/extraction.py

@@ -85,9 +85,11 @@ def hook(module, input, output):
         activation = act[0, -1].detach().clone()
 
     handle = model.model.layers[layer].register_forward_hook(hook)
-    with torch.inference_mode():
-        model(**full_inputs)
-    handle.remove()
+    try:
+        with torch.inference_mode():
+            model(**full_inputs)
+    finally:
+        handle.remove()
     return activation
 
 

modulo/registry.py

@@ -128,14 +128,14 @@ def load_vector(self, vectors_dir: Path) -> torch.Tensor:
         """Load the actual vector tensor for this entry's layer."""
         path = vectors_dir / self.file
         if self.format == "per_layer":
-            data = torch.load(path, weights_only=True)
+            data = torch.load(path, map_location="cpu", weights_only=True)
             return data[self.layer].to(torch.bfloat16)
         elif self.file.endswith(".npz"):
             import numpy as np
             data = np.load(path)
             return torch.tensor(data["vector"], dtype=torch.bfloat16)
         elif self.file.endswith(".pt"):
-            data = torch.load(path, weights_only=True)
+            data = torch.load(path, map_location="cpu", weights_only=True)
             return data.to(torch.bfloat16)
         else:
             raise ValueError(f"Unknown vector format: {self.file}")

pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "modulo"
+version = "0.1.0"
+description = "Self-steering for local language models: let a model steer its own activations."
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.4",
+    "transformers>=4.46,<5",
+    "accelerate>=1.0",
+    "huggingface_hub",
+    "numpy",
+    "pyyaml",
+]
+
+[project.optional-dependencies]
+# Extra packages used by the repo's scripts (CAA extraction, quantization) and tests
+scripts = ["steering-vectors>=0.12", "datasets>=3.0", "scipy", "bitsandbytes"]
+test = ["pytest"]
+
+[tool.setuptools]
+# Only the library is installable; the demo, scripts and config.py stay repo-local.
+# Install editable (pip install -e .) so the default vectors/ dir resolves to this repo.
+packages = ["modulo"]

replay.py

@@ -44,8 +44,8 @@ def replay(log_path: str):
             print(f"[T{turn}] REPORT: {e.get('concept')} worked={e.get('worked')} — {e.get('notes')}")
         elif t == "tool_error":
             print(f"[T{turn}] TOOL ERROR: {e.get('name')} — {e.get('message', '')}")
-        elif t == "decay":
-            print(f"[T{turn}] DECAY: {e.get('active_vectors')}")
+        elif t in ("steering_state", "decay"):  # "decay" in older logs
+            print(f"[T{turn}] STEERING: {e.get('active_vectors')}")
         elif t == "degeneration":
             print(f"[T{turn}] DEGENERATION DETECTED (dropped {e.get('dropped', '?')}): {e.get('snippet', '')[:100]}...")
         print()

setup.sh

@@ -30,7 +30,7 @@ if ! pip install --break-system-packages --quiet -r requirements.txt; then
 fi
 
 # 2. Create directories
-mkdir -p "$VEC_DIR" data logs agent
+mkdir -p "$VEC_DIR" logs
 
 # 3. Download the assistant-axis steering vector (the one confirmed-working vector)
 echo ""

tests/test_agent_logic.py

@@ -2,7 +2,10 @@
 
 No model is loaded; SteeringAgent runs against a stub registry.
 """
+from contextlib import nullcontext
+
 import pytest
+import torch
 
 import modulo.agent as A
 from modulo.agent import (COOLDOWN_TURNS, DECAY_FACTOR, SteeringAgent, ToolCall,
@@ -85,7 +88,13 @@ def test_strength_clamped_to_entry_and_global_cap(agent):
     assert agent.active["a"].current_strength == cap
 
 
+def in_step(agent):
+    """Put the agent where step() runs tool calls: after generation, before decay."""
+    agent._turn_generated = True
+
+
 def test_full_strength_on_first_turn_then_decays(agent):
+    in_step(agent)
     agent.apply("a", 5.0)
     agent.decay()  # end of the turn steer was called; next turn runs at full strength
     assert agent.active["a"].current_strength == 5.0
@@ -94,6 +103,7 @@ def test_full_strength_on_first_turn_then_decays(agent):
 
 
 def test_cooldown_blocks_exactly_cooldown_turns(agent):
+    in_step(agent)
     agent.apply("a", 1.0)
     assert agent.drop("a")[0]
     agent.decay()
@@ -106,6 +116,31 @@ def test_cooldown_blocks_exactly_cooldown_turns(agent):
     assert blocked == COOLDOWN_TURNS
 
 
+def test_manual_apply_between_steps_gets_one_full_strength_turn(agent):
+    agent.apply("a", 5.0)  # harness call before step()
+    assert "turn 1" in agent.status()
+    assert agent.active["a"].current_strength == 5.0  # strength used for this turn's generation
+    in_step(agent)
+    agent.decay()
+    assert agent.active["a"].current_strength == pytest.approx(5.0 * DECAY_FACTOR)
+
+
+def test_manual_drop_between_steps_blocks_exactly_cooldown_turns(agent):
+    in_step(agent)
+    agent.apply("a", 1.0)
+    agent.decay()
+    assert agent.drop("a")[0]  # harness call before step()
+    assert f"{COOLDOWN_TURNS} more turns" in agent.status()
+    blocked = 0
+    for _ in range(COOLDOWN_TURNS + 3):
+        in_step(agent)
+        if agent.apply("a", 1.0)[0]:
+            break
+        blocked += 1
+        agent.decay()
+    assert blocked == COOLDOWN_TURNS
+
+
 def test_natural_expiry_has_no_cooldown(agent):
     agent.apply("b", 0.2)
     for _ in range(10):
@@ -127,6 +162,8 @@ def test_max_active_vectors(agent):
 @pytest.mark.parametrize("args", [
     {"concept": "a", "strength": "high"},
     {"concept": "a", "strength": "nan"},
+    {"concept": "a", "strength": "inf"},
+    {"concept": "a", "strength": "-inf"},
     {"concept": ["a"], "strength": 1},
     {"strength": 1},
     {"concept": "nope", "strength": 1},
@@ -157,3 +194,37 @@ def test_tool_messages_include_failures_and_degeneration():
     assert [m["role"] for m in msgs] == ["tool", "tool"]
     assert "degeneration_detected" in msgs[0]["content"]
     assert '"ok": false' in msgs[1]["content"] and "on cooldown" in msgs[1]["content"]
+
+
+# --- prompt building ---
+
+def test_harness_system_prompt_is_merged(agent):
+    full = agent._full_messages([{"role": "system", "content": "HARNESS RULES"},
+                                 {"role": "user", "content": "hi"}])
+    assert [m["role"] for m in full] == ["system", "user"]
+    assert full[0]["content"].startswith("HARNESS RULES\n\n")
+    assert "self-steering experiment" in full[0]["content"]
+
+
+def test_degenerate_turn_regenerates_with_updated_prompt(agent, monkeypatch):
+    """The unsteered retry must be built after the drops, so its prompt doesn't list them as active."""
+    class FakeInputs:
+        def __init__(self, prompt):
+            self.prompt = prompt
+            self.input_ids = torch.zeros(1, 1)
+
+    prompts, outputs = [], ["loop " * 50, "Fine now."]
+
+    def fake_generate(inputs, *args):
+        prompts.append(inputs.prompt)
+        return outputs.pop(0)
+
+    monkeypatch.setattr(agent, "_encode", lambda msgs: FakeInputs(agent._full_messages(msgs)[0]["content"]))
+    monkeypatch.setattr(agent, "_build_steering_context", lambda prompt_len: nullcontext())
+    monkeypatch.setattr(agent, "_generate", fake_generate)
+
+    agent.apply("a", 5.0)
+    result = agent.step([{"role": "user", "content": "hi"}])
+    assert result.degenerate_output is not None and result.raw_output == "Fine now."
+    assert "a: +5.00" in prompts[0]
+    assert "a: +5.00" not in prompts[1] and "a: on cooldown" in prompts[1]