personal memory agent
at scratch/segment-sense-rd 493 lines 16 kB view raw
1# SPDX-License-Identifier: AGPL-3.0-only 2# Copyright (c) 2026 sol pbc 3 4"""Tests for the generator output hooks system. 5 6Tests cover: 7- Hook loading and validation via load_post_hook / load_pre_hook 8- Hook invocation via NDJSON protocol 9- Hook error handling 10""" 11 12import importlib 13import io 14import json 15import os 16import shutil 17from pathlib import Path 18 19from think.muse import load_post_hook, load_pre_hook 20from think.utils import day_path 21 22FIXTURES = Path("tests/fixtures") 23 24 25def copy_day(tmp_path: Path) -> Path: 26 os.environ["_SOLSTONE_JOURNAL_OVERRIDE"] = str(tmp_path) 27 dest = day_path("20240101") 28 src = FIXTURES / "journal" / "20240101" 29 for item in src.iterdir(): 30 if item.is_dir(): 31 shutil.copytree(item, dest / item.name, dirs_exist_ok=True) 32 else: 33 shutil.copy2(item, dest / item.name) 34 return dest 35 36 37MOCK_RESULT = { 38 "text": "## Original Result\n\nThis is the original output content.", 39 "usage": {"input_tokens": 100, "output_tokens": 50}, 40} 41 42 43def run_generator_with_config(mod, config: dict, monkeypatch) -> list[dict]: 44 """Run generator with NDJSON config and capture output events.""" 45 # Mock argv to prevent argparse from seeing pytest args 46 monkeypatch.setattr("sys.argv", ["sol"]) 47 48 stdin_data = json.dumps(config) + "\n" 49 monkeypatch.setattr("sys.stdin", io.StringIO(stdin_data)) 50 51 captured_output = io.StringIO() 52 monkeypatch.setattr("sys.stdout", captured_output) 53 54 mod.main() 55 56 events = [] 57 captured_output.seek(0) 58 for line in captured_output: 59 line = line.strip() 60 if line: 61 events.append(json.loads(line)) 62 63 return events 64 65 66def test_load_post_hook_success(tmp_path): 67 """Test loading a valid hook with post_process function.""" 68 hook_file = tmp_path / "test_hook.py" 69 hook_file.write_text(""" 70def post_process(result, context): 71 return result + "\\n\\n## Added by hook" 72""") 73 74 # Config with explicit path 75 config = {"hook": {"post": str(hook_file)}} 76 hook_fn = load_post_hook(config) 77 assert callable(hook_fn) 78 79 # Test the hook transforms content 80 output = hook_fn("Original", {"day": "20240101"}) 81 assert output == "Original\n\n## Added by hook" 82 83 84def test_load_post_hook_missing_post_process(tmp_path): 85 """Test that hook without post_process function raises ValueError.""" 86 hook_file = tmp_path / "bad_hook.py" 87 hook_file.write_text(""" 88def other_function(): 89 pass 90""") 91 92 config = {"hook": {"post": str(hook_file)}} 93 try: 94 load_post_hook(config) 95 assert False, "Should have raised ValueError" 96 except ValueError as e: 97 assert "must define a 'post_process' function" in str(e) 98 99 100def test_load_post_hook_not_callable(tmp_path): 101 """Test that hook with non-callable post_process raises ValueError.""" 102 hook_file = tmp_path / "bad_hook.py" 103 hook_file.write_text(""" 104post_process = "not a function" 105""") 106 107 config = {"hook": {"post": str(hook_file)}} 108 try: 109 load_post_hook(config) 110 assert False, "Should have raised ValueError" 111 except ValueError as e: 112 assert "'post_process' must be callable" in str(e) 113 114 115def test_load_post_hook_no_hook_config(): 116 """Test that missing hook config returns None.""" 117 assert load_post_hook({}) is None 118 assert load_post_hook({"hook": {}}) is None 119 assert load_post_hook({"hook": {"pre": "something"}}) is None 120 121 122def test_load_post_hook_named_resolution(): 123 """Test that named hooks resolve to muse/{name}.py.""" 124 # occurrence.py exists in muse/ 125 config = {"hook": {"post": "occurrence"}} 126 hook_fn = load_post_hook(config) 127 assert callable(hook_fn) 128 129 130def test_load_post_hook_file_not_found(tmp_path): 131 """Test that nonexistent hook file raises ImportError.""" 132 config = {"hook": {"post": str(tmp_path / "nonexistent.py")}} 133 try: 134 load_post_hook(config) 135 assert False, "Should have raised ImportError" 136 except ImportError as e: 137 assert "not found" in str(e) 138 139 140def test_prompt_metadata_no_hook_path(tmp_path): 141 """Test that _load_prompt_metadata no longer sets hook_path.""" 142 muse = importlib.import_module("think.muse") 143 144 md_file = tmp_path / "test_generator.md" 145 md_file.write_text( 146 '{\n "title": "Test",\n "hook": {"post": "entities"}\n}\n\nTest prompt' 147 ) 148 149 # Create a co-located .py file 150 hook_file = tmp_path / "test_generator.py" 151 hook_file.write_text("def post_process(r, c): return r") 152 153 meta = muse._load_prompt_metadata(md_file) 154 155 # hook_path should no longer be set (hooks are loaded via load_post_hook) 156 assert "hook_path" not in meta 157 assert meta["path"] == str(md_file) 158 assert meta["title"] == "Test" 159 160 161def test_output_hook_invocation(tmp_path, monkeypatch): 162 """Test that agents.py invokes hook and uses transformed result.""" 163 mod = importlib.import_module("think.agents") 164 copy_day(tmp_path) 165 166 # Use tmp_path as muse directory to avoid polluting real muse/ 167 import think.muse 168 169 monkeypatch.setattr(think.muse, "MUSE_DIR", tmp_path) 170 171 prompt_file = tmp_path / "hooked_test.md" 172 prompt_file.write_text( 173 '{\n "type": "generate",\n "title": "Hooked",\n "schedule": "daily",\n "priority": 10,\n "output": "md",\n "hook": {"post": "hooked_test"},\n "instructions": {"system": "journal", "sources": {"transcripts": true, "percepts": true}}\n}\n\nTest prompt' 174 ) 175 176 hook_file = tmp_path / "hooked_test.py" 177 hook_file.write_text(""" 178def post_process(result, context): 179 # Verify context has expected fields 180 assert "day" in context 181 assert "transcript" in context 182 assert "name" in context 183 return result + "\\n\\n## Hook was here" 184""") 185 186 # Mock the underlying generation function in think.models 187 import think.models 188 189 monkeypatch.setattr( 190 think.models, 191 "generate_with_result", 192 lambda *a, **k: MOCK_RESULT, 193 ) 194 monkeypatch.setenv("GOOGLE_API_KEY", "x") 195 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(tmp_path)) 196 197 config = { 198 "name": "hooked_test", 199 "day": "20240101", 200 "output": "md", 201 "provider": "google", 202 "model": "gemini-2.0-flash", 203 } 204 205 events = run_generator_with_config(mod, config, monkeypatch) 206 207 # Find finish event 208 finish_events = [e for e in events if e["event"] == "finish"] 209 assert len(finish_events) == 1 210 211 content = finish_events[0]["result"] 212 assert "## Original Result" in content 213 assert "## Hook was here" in content 214 215 216def test_output_hook_returns_none(tmp_path, monkeypatch): 217 """Test that hook returning None uses original result.""" 218 mod = importlib.import_module("think.agents") 219 copy_day(tmp_path) 220 221 import think.muse 222 223 monkeypatch.setattr(think.muse, "MUSE_DIR", tmp_path) 224 225 prompt_file = tmp_path / "noop_test.md" 226 prompt_file.write_text( 227 '{\n "type": "generate",\n "title": "Noop",\n "schedule": "daily",\n "priority": 10,\n "output": "md",\n "hook": {"post": "noop_test"},\n "instructions": {"system": "journal", "sources": {"transcripts": true, "percepts": true}}\n}\n\nTest prompt' 228 ) 229 230 hook_file = tmp_path / "noop_test.py" 231 hook_file.write_text(""" 232def post_process(result, context): 233 return None # Signal to use original 234""") 235 236 # Mock the underlying generation function in think.models 237 import think.models 238 239 monkeypatch.setattr( 240 think.models, 241 "generate_with_result", 242 lambda *a, **k: MOCK_RESULT, 243 ) 244 monkeypatch.setenv("GOOGLE_API_KEY", "x") 245 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(tmp_path)) 246 247 config = { 248 "name": "noop_test", 249 "day": "20240101", 250 "output": "md", 251 "provider": "google", 252 "model": "gemini-2.0-flash", 253 } 254 255 events = run_generator_with_config(mod, config, monkeypatch) 256 257 finish_events = [e for e in events if e["event"] == "finish"] 258 assert len(finish_events) == 1 259 assert finish_events[0]["result"] == MOCK_RESULT["text"] 260 261 262def test_output_hook_error_fallback(tmp_path, monkeypatch): 263 """Test that hook errors fall back to original result.""" 264 mod = importlib.import_module("think.agents") 265 copy_day(tmp_path) 266 267 import think.muse 268 269 monkeypatch.setattr(think.muse, "MUSE_DIR", tmp_path) 270 271 prompt_file = tmp_path / "broken_test.md" 272 prompt_file.write_text( 273 '{\n "type": "generate",\n "title": "Broken",\n "schedule": "daily",\n "priority": 10,\n "output": "md",\n "hook": {"post": "broken_test"},\n "instructions": {"system": "journal", "sources": {"transcripts": true, "percepts": true}}\n}\n\nTest prompt' 274 ) 275 276 hook_file = tmp_path / "broken_test.py" 277 hook_file.write_text(""" 278def post_process(result, context): 279 raise RuntimeError("Hook exploded!") 280""") 281 282 # Mock the underlying generation function in think.models 283 import think.models 284 285 monkeypatch.setattr( 286 think.models, 287 "generate_with_result", 288 lambda *a, **k: MOCK_RESULT, 289 ) 290 monkeypatch.setenv("GOOGLE_API_KEY", "x") 291 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(tmp_path)) 292 293 config = { 294 "name": "broken_test", 295 "day": "20240101", 296 "output": "md", 297 "provider": "google", 298 "model": "gemini-2.0-flash", 299 } 300 301 # Should not raise, should fall back gracefully 302 events = run_generator_with_config(mod, config, monkeypatch) 303 304 finish_events = [e for e in events if e["event"] == "finish"] 305 assert len(finish_events) == 1 306 assert finish_events[0]["result"] == MOCK_RESULT["text"] 307 308 309# ============================================================================= 310# Pre-hook Tests 311# ============================================================================= 312 313 314def test_load_pre_hook_success(tmp_path): 315 """Test loading a valid hook with pre_process function.""" 316 hook_file = tmp_path / "test_pre_hook.py" 317 hook_file.write_text(""" 318def pre_process(context): 319 return {"prompt": context["prompt"] + " [modified]"} 320""") 321 322 config = {"hook": {"pre": str(hook_file)}} 323 hook_fn = load_pre_hook(config) 324 assert callable(hook_fn) 325 326 # Test the hook returns modifications 327 result = hook_fn({"prompt": "original"}) 328 assert result == {"prompt": "original [modified]"} 329 330 331def test_load_pre_hook_missing_pre_process(tmp_path): 332 """Test that hook without pre_process function raises ValueError.""" 333 hook_file = tmp_path / "bad_hook.py" 334 hook_file.write_text(""" 335def other_function(): 336 pass 337""") 338 339 config = {"hook": {"pre": str(hook_file)}} 340 try: 341 load_pre_hook(config) 342 assert False, "Should have raised ValueError" 343 except ValueError as e: 344 assert "must define a 'pre_process' function" in str(e) 345 346 347def test_load_pre_hook_not_callable(tmp_path): 348 """Test that hook with non-callable pre_process raises ValueError.""" 349 hook_file = tmp_path / "bad_hook.py" 350 hook_file.write_text(""" 351pre_process = "not a function" 352""") 353 354 config = {"hook": {"pre": str(hook_file)}} 355 try: 356 load_pre_hook(config) 357 assert False, "Should have raised ValueError" 358 except ValueError as e: 359 assert "'pre_process' must be callable" in str(e) 360 361 362def test_load_pre_hook_no_hook_config(): 363 """Test that missing hook config returns None.""" 364 assert load_pre_hook({}) is None 365 assert load_pre_hook({"hook": {}}) is None 366 assert load_pre_hook({"hook": {"post": "something"}}) is None 367 368 369def test_load_pre_hook_file_not_found(tmp_path): 370 """Test that nonexistent hook file raises ImportError.""" 371 config = {"hook": {"pre": str(tmp_path / "nonexistent.py")}} 372 try: 373 load_pre_hook(config) 374 assert False, "Should have raised ImportError" 375 except ImportError as e: 376 assert "not found" in str(e) 377 378 379def test_pre_hook_invocation(tmp_path, monkeypatch): 380 """Test that agents.py invokes pre-hook and uses modified inputs.""" 381 mod = importlib.import_module("think.agents") 382 copy_day(tmp_path) 383 384 import think.muse 385 386 monkeypatch.setattr(think.muse, "MUSE_DIR", tmp_path) 387 388 prompt_file = tmp_path / "prehooked_test.md" 389 prompt_file.write_text( 390 '{\n "type": "generate",\n "title": "Prehooked",\n "schedule": "daily",\n "priority": 10,\n "output": "md",\n "hook": {"pre": "prehooked_test"},\n "instructions": {"system": "journal", "sources": {"transcripts": true, "percepts": true}}\n}\n\nOriginal prompt' 391 ) 392 393 hook_file = tmp_path / "prehooked_test.py" 394 hook_file.write_text(""" 395def pre_process(context): 396 # Verify context has expected fields 397 assert "transcript" in context 398 assert "prompt" in context 399 assert "system_instruction" in context 400 # Modify the prompt 401 return {"prompt": context["prompt"] + " [pre-processed]"} 402""") 403 404 # Track what generate_with_result receives 405 received_kwargs = {} 406 407 def mock_generate(*args, **kwargs): 408 received_kwargs.update(kwargs) 409 received_kwargs["contents"] = args[0] if args else kwargs.get("contents") 410 return MOCK_RESULT 411 412 import think.models 413 414 monkeypatch.setattr(think.models, "generate_with_result", mock_generate) 415 monkeypatch.setenv("GOOGLE_API_KEY", "x") 416 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(tmp_path)) 417 418 config = { 419 "name": "prehooked_test", 420 "day": "20240101", 421 "output": "md", 422 "provider": "google", 423 "model": "gemini-2.0-flash", 424 } 425 426 events = run_generator_with_config(mod, config, monkeypatch) 427 428 # Verify pre-hook modified the prompt - check in contents 429 contents = received_kwargs.get("contents", []) 430 # The prompt should contain [pre-processed] 431 prompt_found = any("[pre-processed]" in str(c) for c in contents) 432 assert prompt_found, f"Expected [pre-processed] in contents: {contents}" 433 434 # Verify generator still completed successfully 435 finish_events = [e for e in events if e["event"] == "finish"] 436 assert len(finish_events) == 1 437 438 439def test_both_pre_and_post_hooks(tmp_path, monkeypatch): 440 """Test that both pre and post hooks can be configured together.""" 441 mod = importlib.import_module("think.agents") 442 copy_day(tmp_path) 443 444 import think.muse 445 446 monkeypatch.setattr(think.muse, "MUSE_DIR", tmp_path) 447 448 prompt_file = tmp_path / "both_hooks_test.md" 449 prompt_file.write_text( 450 '{\n "type": "generate",\n "title": "Both Hooks",\n "schedule": "daily",\n "priority": 10,\n "output": "md",\n "hook": {"pre": "both_hooks_test", "post": "both_hooks_test"},\n "instructions": {"system": "journal", "sources": {"transcripts": true, "percepts": true}}\n}\n\nOriginal prompt' 451 ) 452 453 hook_file = tmp_path / "both_hooks_test.py" 454 hook_file.write_text(""" 455def pre_process(context): 456 return {"prompt": context["prompt"] + " [pre]"} 457 458def post_process(result, context): 459 return result + "\\n\\n[post]" 460""") 461 462 received_kwargs = {} 463 464 def mock_generate(*args, **kwargs): 465 received_kwargs.update(kwargs) 466 received_kwargs["contents"] = args[0] if args else kwargs.get("contents") 467 return MOCK_RESULT 468 469 import think.models 470 471 monkeypatch.setattr(think.models, "generate_with_result", mock_generate) 472 monkeypatch.setenv("GOOGLE_API_KEY", "x") 473 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(tmp_path)) 474 475 config = { 476 "name": "both_hooks_test", 477 "day": "20240101", 478 "output": "md", 479 "provider": "google", 480 "model": "gemini-2.0-flash", 481 } 482 483 events = run_generator_with_config(mod, config, monkeypatch) 484 485 # Verify pre-hook modified the prompt - check in contents 486 contents = received_kwargs.get("contents", []) 487 prompt_found = any("[pre]" in str(c) for c in contents) 488 assert prompt_found, f"Expected [pre] in contents: {contents}" 489 490 # Verify post-hook modified the result 491 finish_events = [e for e in events if e["event"] == "finish"] 492 assert len(finish_events) == 1 493 assert "[post]" in finish_events[0]["result"]