feat: inject today's date into system prompt and add extended thinking support

Fixes age validation errors caused by the LLM not knowing the current date.

Changes:
- prompts.py: inject date.today() at the top of both system prompts so the
  LLM can accurately calculate a child's age from their date of birth
- llm.py: add optional thinking_budget parameter to complete(); when set,
  passes thinking={"type": "enabled", "budget_tokens": N} to litellm and
  raises max_tokens to thinking_budget + 4096 (Anthropic models only)
- config.py: add thinking_budget field, read from THINKING_BUDGET env var
- .env.example: document the THINKING_BUDGET option
- core.py: pass thinking_budget through to llm.complete()
- main.py: pass thinking_budget when constructing EmailAgent
- chat_app.py: switch from stream_complete to asyncio.to_thread(complete)
  so extended thinking works and so only the reply field is shown to
  the parent (not the raw JSON wrapper)

To enable extended thinking set THINKING_BUDGET=8000 in .env.

https://claude.ai/code/session_01SUWzMzFvSfWiHXA2p6rPg9
This commit is contained in:
Claude
2026-02-22 19:57:32 +00:00
parent 13cb35111d
commit 1bd11e6acd
7 changed files with 76 additions and 16 deletions
+6
View File
@@ -15,6 +15,12 @@
# gemini/gemini-2.0-flash
AI_MODEL=anthropic/claude-opus-4-6
# Extended thinking — Anthropic models only (leave unset to disable).
# Enables a reasoning phase before the model's reply, which improves
# accuracy on age calculations, logic-heavy questions, and edge cases.
# Recommended value: 8000 (tokens). Must be less than max_tokens.
# THINKING_BUDGET=8000
# ---------------------------------------------------------------
# API Keys — set the one matching your chosen model's provider
# ---------------------------------------------------------------
+7 -9
View File
@@ -123,17 +123,15 @@ async def on_message(message: cl.Message) -> None:
# --- Build system prompt ---
system = build_system_prompt(_kb, state)
# --- Call LLM natively async to keep the event loop free ---
# litellm.acompletion() is a true coroutine — it does not block the event
# loop and does not require asyncio.to_thread (which can drop Chainlit's
# contextvars, causing the session to reset mid-conversation).
# --- Call LLM natively async (supports extended thinking; no event-loop blocking) ---
try:
full_content = await llm.acomplete(_config.ai_model, system, state.messages)
full_content = await llm.acomplete(
_config.ai_model, system, state.messages, _config.thinking_budget
)
except Exception:
logger.exception("LLM call failed for session %s", state.conversation_id)
error_text = fallback_message(state.language)
await cl.Message(content=error_text).send()
# Don't update state — let parent retry
return
# --- Parse and apply LLM response ---
@@ -151,12 +149,12 @@ async def on_message(message: cl.Message) -> None:
state.language = language
state.updated_at = now
# Send the reply text to the parent (only the human-readable reply, not the JSON wrapper)
await cl.Message(content=reply_text).send()
# Append assistant reply to history
state.messages.append(ChatMessage(role="assistant", content=reply_text))
# Send the reply text to the user (parsed from the LLM's JSON response)
await cl.Message(content=reply_text).send()
# --- Handle registration completion ---
if is_complete and not state.completed:
state.completed = True
+7 -1
View File
@@ -63,7 +63,13 @@ def build_components(config: Config):
cc_emails=[e.strip() for e in config.admin_email_cc.split(",") if e.strip()],
)
agent = EmailAgent(model=config.ai_model, kb=kb, store=store, notifier=notifier)
agent = EmailAgent(
model=config.ai_model,
kb=kb,
store=store,
notifier=notifier,
thinking_budget=config.thinking_budget,
)
channel = EmailChannel(
imap_host=config.imap_host,
+4 -2
View File
@@ -31,8 +31,10 @@ class EmailAgent:
kb: KnowledgeBase,
store: ConversationStore,
notifier: AdminNotifier,
thinking_budget: int | None = None,
) -> None:
self._model = model
self._thinking_budget = thinking_budget
self._kb = kb
self._store = store
self._notifier = notifier
@@ -98,7 +100,7 @@ class EmailAgent:
system = build_system_prompt(self._kb, state)
try:
content = llm.complete(self._model, system, state.messages)
content = llm.complete(self._model, system, state.messages, self._thinking_budget)
parsed = self._parse_llm_response(content)
except Exception:
logger.exception("LLM call failed for %s", state.conversation_id)
@@ -140,7 +142,7 @@ class EmailAgent:
system = build_system_prompt(self._kb, state)
try:
content = llm.complete(self._model, system, state.messages)
content = llm.complete(self._model, system, state.messages, self._thinking_budget)
parsed = self._parse_llm_response(content)
except Exception:
logger.exception("LLM call failed (post-completion) for %s", state.conversation_id)
+9
View File
@@ -1,6 +1,7 @@
"""Build the system prompt sent to the LLM on every turn."""
import json
from datetime import date
from ..knowledge_base.loader import KnowledgeBase
from ..models.conversation import ConversationState
@@ -154,8 +155,12 @@ def _build_registration_prompt(kb: KnowledgeBase, state: ConversationState) -> s
reg_json = json.dumps(state.registration.to_dict(), ensure_ascii=False, indent=2)
step_hint = STEP_DESCRIPTIONS.get(state.flow_step, "Continue the conversation.")
today = date.today().isoformat()
return f"""You are the registration assistant for Spielgruppe Pumuckl, run by Familienverein Fällanden in Fällanden, Switzerland. You help parents register their children for the playgroup and answer questions about it.
**Today's date is {today}.** Use this exact date when calculating a child's age from their date of birth.
{_PERSONALITY}
## Registration Flow (8 steps)
@@ -200,8 +205,12 @@ def _build_post_completion_prompt(kb: KnowledgeBase, state: ConversationState) -
reg_json = json.dumps(state.registration.to_dict(), ensure_ascii=False, indent=2)
child_name = state.registration.child.full_name or "their child"
today = date.today().isoformat()
return f"""You are the registration assistant for Spielgruppe Pumuckl, run by Familienverein Fällanden in Fällanden, Switzerland.
**Today's date is {today}.**
{_PERSONALITY}
## Context: Registration Already Complete
+10
View File
@@ -50,6 +50,11 @@ class Config:
# Polling interval in seconds
poll_interval: int = 60
# Extended thinking — Anthropic models only.
# When set, enables the thinking phase before the LLM replies.
# Recommended value: 8000 (tokens). Set to None/unset to disable.
thinking_budget: int | None = None
@classmethod
def from_env(cls) -> "Config":
return cls(
@@ -74,4 +79,9 @@ class Config:
)
),
poll_interval=int(os.getenv("POLL_INTERVAL", "60")),
thinking_budget=(
int(os.getenv("THINKING_BUDGET"))
if os.getenv("THINKING_BUDGET")
else None
),
)
+33 -4
View File
@@ -5,7 +5,12 @@ from collections.abc import Generator
import litellm
async def acomplete(model: str, system: str, messages: list) -> str:
async def acomplete(
model: str,
system: str,
messages: list,
thinking_budget: int | None = None,
) -> str:
"""Call any LLM asynchronously and return the response text.
This is the async equivalent of ``complete()`` — use this from async
@@ -16,17 +21,30 @@ async def acomplete(model: str, system: str, messages: list) -> str:
model: litellm model string, e.g. "anthropic/claude-opus-4-6".
system: System prompt text.
messages: List of objects with .role and .content attributes.
thinking_budget: When set, enables extended thinking (Anthropic models
only). See ``complete()`` for details.
Returns:
The model's reply as a plain string.
"""
api_messages = [{"role": "system", "content": system}]
api_messages += [{"role": m.role, "content": m.content} for m in messages]
response = await litellm.acompletion(model=model, messages=api_messages, max_tokens=2048)
kwargs: dict = {"model": model, "messages": api_messages, "max_tokens": 2048}
if thinking_budget is not None:
kwargs["thinking"] = {"type": "enabled", "budget_tokens": thinking_budget}
kwargs["max_tokens"] = thinking_budget + 4096
response = await litellm.acompletion(**kwargs)
return response.choices[0].message.content
def complete(model: str, system: str, messages: list) -> str:
def complete(
model: str,
system: str,
messages: list,
thinking_budget: int | None = None,
) -> str:
"""Call any LLM and return the response text.
Args:
@@ -35,13 +53,24 @@ def complete(model: str, system: str, messages: list) -> str:
environment variable (ANTHROPIC_API_KEY, OPENAI_API_KEY, …).
system: System prompt text.
messages: List of objects with .role and .content attributes.
thinking_budget: When set, enables extended thinking (Anthropic models
only). The value is the token budget for the thinking phase; the
final ``max_tokens`` is set to ``thinking_budget + 4096`` so the
model has enough room to both think and reply.
Returns:
The model's reply as a plain string.
"""
api_messages = [{"role": "system", "content": system}]
api_messages += [{"role": m.role, "content": m.content} for m in messages]
response = litellm.completion(model=model, messages=api_messages, max_tokens=2048)
kwargs: dict = {"model": model, "messages": api_messages, "max_tokens": 2048}
if thinking_budget is not None:
kwargs["thinking"] = {"type": "enabled", "budget_tokens": thinking_budget}
# max_tokens must exceed budget_tokens or the API returns an error.
kwargs["max_tokens"] = thinking_budget + 4096
response = litellm.completion(**kwargs)
return response.choices[0].message.content