personal memory agent
at main 130 lines 3.9 kB view raw
1# SPDX-License-Identifier: AGPL-3.0-only 2# Copyright (c) 2026 sol pbc 3 4import asyncio 5import importlib 6import json 7import sys 8from unittest.mock import AsyncMock 9 10from tests.conftest import setup_google_genai_stub 11from tests.test_google import make_mock_process 12from think.models import GEMINI_FLASH 13 14 15async def run_main(mod, argv, stdin_data=None): 16 sys.argv = argv 17 if stdin_data: 18 import io 19 20 sys.stdin = io.StringIO(stdin_data) 21 await mod.main_async() 22 23 24def test_google_thinking_events(monkeypatch, tmp_path, capsys): 25 setup_google_genai_stub(monkeypatch, with_thinking=True) 26 27 sys.modules.pop("think.providers.google", None) 28 importlib.reload(importlib.import_module("think.providers.google")) 29 mod = importlib.reload(importlib.import_module("think.agents")) 30 31 journal = tmp_path / "journal" 32 journal.mkdir() 33 34 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(journal)) 35 monkeypatch.setenv("GOOGLE_API_KEY", "x") 36 monkeypatch.setattr( 37 "think.providers.cli.shutil.which", 38 lambda name: "/usr/bin/gemini" if name == "gemini" else None, 39 ) 40 41 # Simulate CLI output: thinking text before a tool call triggers thinking event 42 stdout_lines = [ 43 json.dumps( 44 { 45 "type": "init", 46 "timestamp": 100, 47 "session_id": "sess-think", 48 "model": "gemini-2.5-flash", 49 } 50 ), 51 json.dumps( 52 { 53 "type": "message", 54 "role": "assistant", 55 "delta": True, 56 "content": "I need to analyze this step by step.", 57 } 58 ), 59 json.dumps( 60 { 61 "type": "tool_use", 62 "timestamp": 150, 63 "tool_name": "search_insights", 64 "tool_id": "t1", 65 "parameters": {"query": "hello"}, 66 } 67 ), 68 json.dumps( 69 { 70 "type": "tool_result", 71 "timestamp": 200, 72 "tool_id": "t1", 73 "status": "success", 74 "output": "result data", 75 } 76 ), 77 json.dumps( 78 { 79 "type": "message", 80 "role": "assistant", 81 "delta": True, 82 "content": "ok", 83 } 84 ), 85 json.dumps( 86 { 87 "type": "result", 88 "timestamp": 300, 89 "status": "success", 90 "stats": { 91 "total_tokens": 20, 92 "input_tokens": 10, 93 "output_tokens": 10, 94 }, 95 } 96 ), 97 ] 98 process = make_mock_process(stdout_lines) 99 monkeypatch.setattr( 100 "think.providers.cli.asyncio.create_subprocess_exec", 101 AsyncMock(return_value=process), 102 ) 103 104 ndjson_input = json.dumps( 105 { 106 "prompt": "hello", 107 "provider": "google", 108 "model": GEMINI_FLASH, 109 "tools": ["search_insights"], 110 } 111 ) 112 asyncio.run(run_main(mod, ["sol agents"], stdin_data=ndjson_input)) 113 114 out_lines = capsys.readouterr().out.strip().splitlines() 115 events = [json.loads(line) for line in out_lines] 116 117 # Check that we have start, thinking, and finish events 118 assert events[0]["event"] == "start" 119 assert isinstance(events[0]["ts"], int) 120 assert "hello" in events[0]["prompt"] 121 122 # Look for thinking event (flushed when tool_use arrives) 123 thinking_events = [e for e in events if e["event"] == "thinking"] 124 assert len(thinking_events) == 1 125 assert thinking_events[0]["summary"] == "I need to analyze this step by step." 126 assert thinking_events[0]["model"] == GEMINI_FLASH 127 assert isinstance(thinking_events[0]["ts"], int) 128 129 assert events[-1]["event"] == "finish" 130 assert events[-1]["result"] == "ok"