personal memory agent
Add universal max_output_tokens and thinking_budget support for agents
Unify model configuration across insights and agents by adding support
for max_output_tokens and thinking_budget in agent frontmatter and
runtime config. These parameters now work consistently with the existing
generate() functions.
Changes:
- Extract max_output_tokens (replaces max_tokens) in all provider run_agent()
- Extract thinking_budget from config in Google and Anthropic providers
- Google: uses explicit budget or -1 (dynamic) if not specified
- Anthropic: uses explicit budget or computes default via existing logic
- OpenAI: uses max_output_tokens only (no thinking_budget support)
- Add helper functions to eliminate duplicated token/budget calculation
- Update all tests from max_tokens to max_output_tokens
- Document new fields in CORTEX.md, APPS.md, and PROMPT_TEMPLATES.md
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
+80
-23
+2
-2
docs/APPS.md
+2
-2
docs/APPS.md
···
273
273
- Keys are namespaced as `{app}:{topic}` (e.g., `my_app:weekly_summary`)
274
274
- Outputs go to `JOURNAL/YYYYMMDD/insights/_<app>_<topic>.md` (or `.json` if `output: "json"`)
275
275
276
-
**Metadata format:** Same schema as system insights in `muse/*.md` - JSON frontmatter includes `title`, `description`, `color`, `schedule` (required), `hook`, and `output` fields. The `schedule` field must be `"segment"` or `"daily"` - insights with missing or invalid schedule are skipped with a warning. Set `output: "json"` for structured JSON output instead of markdown.
276
+
**Metadata format:** Same schema as system insights in `muse/*.md` - JSON frontmatter includes `title`, `description`, `color`, `schedule` (required), `hook`, `output`, `max_output_tokens`, and `thinking_budget` fields. The `schedule` field must be `"segment"` or `"daily"` - insights with missing or invalid schedule are skipped with a warning. Set `output: "json"` for structured JSON output instead of markdown. Optional `max_output_tokens` sets the maximum response length; `thinking_budget` sets the model's thinking token budget (provider-specific defaults apply if omitted).
277
277
278
278
**Event extraction via hooks:** To extract structured events from insight output, use the `hook` field:
279
279
···
332
332
- Keys are namespaced as `{app}:{name}` (e.g., `my_app:helper`)
333
333
- Agents inherit all system agent capabilities (tools, scheduling, handoffs, multi-facet)
334
334
335
-
**Metadata format:** Same schema as system agents in `muse/*.md` - JSON frontmatter includes `title`, `provider`, `model`, `tools`, `schedule`, `priority`, and `multi_facet` fields. See [CORTEX.md](CORTEX.md) for agent configuration details.
335
+
**Metadata format:** Same schema as system agents in `muse/*.md` - JSON frontmatter includes `title`, `provider`, `model`, `tools`, `schedule`, `priority`, `multi_facet`, `max_output_tokens`, and `thinking_budget` fields. Optional `max_output_tokens` sets the maximum response length; `thinking_budget` sets the model's thinking token budget (provider-specific defaults apply if omitted; OpenAI uses fixed reasoning and ignores this field). See [CORTEX.md](CORTEX.md) for agent configuration details.
336
336
337
337
**Template variables:** Agent prompts can use template variables like `$name`, `$preferred`, and pronoun variables. See [PROMPT_TEMPLATES.md](PROMPT_TEMPLATES.md) for the complete template system documentation.
338
338
+2
-1
docs/CORTEX.md
+2
-1
docs/CORTEX.md
···
39
39
"prompt": "Analyze this code for security issues", // Required: the task or question
40
40
"persona": "default", // Optional: agent persona from muse/*.md
41
41
"provider": "openai", // Optional: override provider (openai, google, anthropic)
42
-
"max_tokens": 8192, // Optional: token limit (if supported)
42
+
"max_output_tokens": 8192, // Optional: maximum response tokens
43
+
"thinking_budget": 10000, // Optional: thinking token budget (ignored by OpenAI)
43
44
"disable_mcp": false, // Optional: disable MCP tools for this request
44
45
"continue_from": "1234567890122", // Optional: continue from previous agent
45
46
"facet": "my-project", // Optional: project context
+4
docs/PROMPT_TEMPLATES.md
+4
docs/PROMPT_TEMPLATES.md
···
117
117
118
118
The `$segment_insight` or `$daily_insight` template provides standardized context about what's being analyzed, while the rest of the prompt defines the specific analysis task.
119
119
120
+
**Optional model configuration:** Add `max_output_tokens` (response length limit) and `thinking_budget` (model thinking token budget) to override provider defaults.
121
+
120
122
**Reference:** `muse/*.md` for examples (files with `schedule` field but no `tools` field)
121
123
122
124
### For Agents
···
139
141
2. **User instruction** - Agent-specific `.md` file (e.g., `muse/default.md`)
140
142
141
143
The system instruction establishes the journal partnership context. The user instruction defines the agent's specific role and capabilities.
144
+
145
+
**Optional model configuration:** Add `max_output_tokens` (response length limit) and `thinking_budget` (model thinking token budget) to override provider defaults. Note: OpenAI uses fixed reasoning and ignores `thinking_budget`.
142
146
143
147
**Reference:** `think/utils.py` → `get_agent()` for agent configuration loading
144
148
+2
-2
tests/integration/test_anthropic_provider.py
+2
-2
tests/integration/test_anthropic_provider.py
···
56
56
"provider": "anthropic",
57
57
"model": CLAUDE_SONNET_4,
58
58
"persona": "default",
59
-
"max_tokens": 100,
59
+
"max_output_tokens": 100,
60
60
"disable_mcp": True,
61
61
}
62
62
)
···
160
160
"provider": "anthropic",
161
161
"persona": "default",
162
162
"model": CLAUDE_SONNET_4,
163
-
"max_tokens": 2048,
163
+
"max_output_tokens": 2048,
164
164
"disable_mcp": True,
165
165
}
166
166
)
+2
-2
tests/integration/test_google_provider.py
+2
-2
tests/integration/test_google_provider.py
···
56
56
"provider": "google",
57
57
"persona": "default",
58
58
"model": GEMINI_FLASH,
59
-
"max_tokens": 100,
59
+
"max_output_tokens": 100,
60
60
"disable_mcp": True,
61
61
}
62
62
)
···
148
148
"provider": "google",
149
149
"persona": "default",
150
150
"model": GEMINI_PRO, # Pro model for thinking
151
-
"max_tokens": 2000,
151
+
"max_output_tokens": 2000,
152
152
"disable_mcp": True,
153
153
}
154
154
)
+3
-3
tests/integration/test_openai_provider.py
+3
-3
tests/integration/test_openai_provider.py
···
56
56
"provider": "openai",
57
57
"persona": "default",
58
58
"model": GPT_5_MINI, # Use cheap model for testing
59
-
"max_tokens": 100,
59
+
"max_output_tokens": 100,
60
60
"disable_mcp": True,
61
61
}
62
62
)
···
153
153
"provider": "openai",
154
154
"persona": "default",
155
155
"model": GPT_5_MINI,
156
-
"max_tokens": 500,
156
+
"max_output_tokens": 500,
157
157
"disable_mcp": True,
158
158
}
159
159
)
···
239
239
"provider": "openai",
240
240
"persona": "default",
241
241
"model": GPT_5_MINI,
242
-
"max_tokens": 50,
242
+
"max_output_tokens": 50,
243
243
"disable_mcp": True,
244
244
"extra_context": "## Project Context\nYou are working on Project Moonshot.",
245
245
}
+1
-1
tests/test_agents_ndjson.py
+1
-1
tests/test_agents_ndjson.py
+1
-1
think/cortex_client.py
+1
-1
think/cortex_client.py
···
33
33
persona: Agent persona - system (e.g., "default") or app-qualified (e.g., "entities:entity_assist")
34
34
provider: AI provider - openai, google, or anthropic
35
35
handoff_from: Previous agent ID if this is a handoff request
36
-
config: Provider-specific configuration (model, max_tokens, etc.)
36
+
config: Provider-specific configuration (model, max_output_tokens, thinking_budget, etc.)
37
37
38
38
Returns:
39
39
Agent ID (timestamp-based string)
+28
-4
think/providers/anthropic.py
+28
-4
think/providers/anthropic.py
···
81
81
return thinking_budget, max_tokens
82
82
83
83
84
+
def _resolve_agent_thinking_params(
85
+
max_output_tokens: int, thinking_budget_config: int | None
86
+
) -> tuple[int, int]:
87
+
"""Resolve thinking budget and max tokens for agent run.
88
+
89
+
Args:
90
+
max_output_tokens: Maximum output tokens from config.
91
+
thinking_budget_config: Explicit thinking budget from config, or None.
92
+
93
+
Returns:
94
+
Tuple of (thinking_budget, effective_max_tokens).
95
+
If thinking_budget_config is provided and > 0, uses it directly.
96
+
Otherwise computes defaults via _compute_thinking_params.
97
+
"""
98
+
if thinking_budget_config is not None and thinking_budget_config > 0:
99
+
return thinking_budget_config, max_output_tokens
100
+
return _compute_thinking_params(max_output_tokens)
101
+
102
+
84
103
def _emit_thinking_event(
85
104
block: ThinkingBlock | RedactedThinkingBlock,
86
105
model: str,
···
267
286
if not model:
268
287
raise ValueError("Missing 'model' in config - should be set by Cortex")
269
288
270
-
max_tokens = config.get("max_tokens", _DEFAULT_MAX_TOKENS)
289
+
max_output_tokens = config.get("max_output_tokens", _DEFAULT_MAX_TOKENS)
290
+
thinking_budget_config = config.get(
291
+
"thinking_budget"
292
+
) # None = use computed default
271
293
disable_mcp = config.get("disable_mcp", False)
272
294
persona = config.get("persona", "default")
273
295
···
334
356
mcp, callback, agent_id=agent_id, persona=persona
335
357
)
336
358
337
-
thinking_budget, effective_max_tokens = _compute_thinking_params(
338
-
max_tokens
359
+
thinking_budget, effective_max_tokens = _resolve_agent_thinking_params(
360
+
max_output_tokens, thinking_budget_config
339
361
)
340
362
341
363
for _ in range(_MAX_TOOL_ITERATIONS):
···
409
431
return "Done."
410
432
else:
411
433
# No MCP tools - single response only
412
-
thinking_budget, effective_max_tokens = _compute_thinking_params(max_tokens)
434
+
thinking_budget, effective_max_tokens = _resolve_agent_thinking_params(
435
+
max_output_tokens, thinking_budget_config
436
+
)
413
437
create_params = {
414
438
"model": model,
415
439
"max_tokens": effective_max_tokens,
+33
-5
think/providers/google.py
+33
-5
think/providers/google.py
···
81
81
return client
82
82
83
83
84
+
def _compute_agent_thinking_params(
85
+
max_output_tokens: int, thinking_budget: int | None
86
+
) -> tuple[int, int]:
87
+
"""Compute total tokens and effective thinking budget for agent run.
88
+
89
+
Args:
90
+
max_output_tokens: Maximum output tokens from config.
91
+
thinking_budget: Thinking budget from config, or None for dynamic.
92
+
93
+
Returns:
94
+
Tuple of (total_tokens, effective_thinking_budget).
95
+
total_tokens = max_output_tokens + (thinking_budget or 0)
96
+
effective_thinking_budget = thinking_budget if provided, else -1 (dynamic)
97
+
"""
98
+
total_tokens = max_output_tokens + (thinking_budget or 0)
99
+
effective_thinking_budget = thinking_budget if thinking_budget is not None else -1
100
+
return total_tokens, effective_thinking_budget
101
+
102
+
84
103
def _build_generate_config(
85
104
temperature: float,
86
105
max_output_tokens: int,
···
504
523
if not model:
505
524
raise ValueError("Missing 'model' in config - should be set by Cortex")
506
525
507
-
max_tokens = config.get("max_tokens", _DEFAULT_MAX_TOKENS)
526
+
max_output_tokens = config.get("max_output_tokens", _DEFAULT_MAX_TOKENS)
527
+
thinking_budget = config.get("thinking_budget") # None = dynamic (-1)
508
528
disable_mcp = config.get("disable_mcp", False)
509
529
persona = config.get("persona", "default")
510
530
···
601
621
else:
602
622
function_calling_config = types.FunctionCallingConfig(mode="AUTO")
603
623
624
+
total_tokens, effective_thinking_budget = (
625
+
_compute_agent_thinking_params(max_output_tokens, thinking_budget)
626
+
)
627
+
604
628
cfg = types.GenerateContentConfig(
605
-
max_output_tokens=max_tokens,
629
+
max_output_tokens=total_tokens,
606
630
tools=[mcp.session],
607
631
tool_config=types.ToolConfig(
608
632
function_calling_config=function_calling_config
609
633
),
610
634
thinking_config=types.ThinkingConfig(
611
635
include_thoughts=True,
612
-
thinking_budget=-1, # Enable dynamic thinking
636
+
thinking_budget=effective_thinking_budget,
613
637
),
614
638
)
615
639
···
621
645
tool_call_count = tool_hooks._counter
622
646
else:
623
647
# No MCP tools - just basic config
648
+
total_tokens, effective_thinking_budget = _compute_agent_thinking_params(
649
+
max_output_tokens, thinking_budget
650
+
)
651
+
624
652
cfg = types.GenerateContentConfig(
625
-
max_output_tokens=max_tokens,
653
+
max_output_tokens=total_tokens,
626
654
thinking_config=types.ThinkingConfig(
627
655
include_thoughts=True,
628
-
thinking_budget=-1, # Enable dynamic thinking
656
+
thinking_budget=effective_thinking_budget,
629
657
),
630
658
)
631
659
+2
-2
think/providers/openai.py
+2
-2
think/providers/openai.py
···
225
225
if not model:
226
226
raise ValueError("Missing 'model' in config - should be set by Cortex")
227
227
228
-
max_tokens = config.get("max_tokens", _DEFAULT_MAX_TOKENS)
228
+
max_output_tokens = config.get("max_output_tokens", _DEFAULT_MAX_TOKENS)
229
229
max_turns = config.get("max_turns", _DEFAULT_MAX_TURNS)
230
230
disable_mcp = config.get("disable_mcp", False)
231
231
continue_from = config.get("continue_from")
···
249
249
250
250
# Model settings: always enable reasoning with detailed summaries
251
251
model_settings = ModelSettings(
252
-
max_tokens=max_tokens,
252
+
max_tokens=max_output_tokens,
253
253
reasoning=_DEFAULT_REASONING,
254
254
)
255
255