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 # gemini/gemini-2.0-flash
AI_MODEL=anthropic/claude-opus-4-6 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 # 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 --- # --- Build system prompt ---
system = build_system_prompt(_kb, state) system = build_system_prompt(_kb, state)
# --- Call LLM natively async to keep the event loop free --- # --- Call LLM natively async (supports extended thinking; no event-loop blocking) ---
# 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).
try: 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: except Exception:
logger.exception("LLM call failed for session %s", state.conversation_id) logger.exception("LLM call failed for session %s", state.conversation_id)
error_text = fallback_message(state.language) error_text = fallback_message(state.language)
await cl.Message(content=error_text).send() await cl.Message(content=error_text).send()
# Don't update state — let parent retry
return return
# --- Parse and apply LLM response --- # --- Parse and apply LLM response ---
@@ -151,12 +149,12 @@ async def on_message(message: cl.Message) -> None:
state.language = language state.language = language
state.updated_at = now 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 # Append assistant reply to history
state.messages.append(ChatMessage(role="assistant", content=reply_text)) 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 --- # --- Handle registration completion ---
if is_complete and not state.completed: if is_complete and not state.completed:
state.completed = True 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()], 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( channel = EmailChannel(
imap_host=config.imap_host, imap_host=config.imap_host,
+4 -2
View File
@@ -31,8 +31,10 @@ class EmailAgent:
kb: KnowledgeBase, kb: KnowledgeBase,
store: ConversationStore, store: ConversationStore,
notifier: AdminNotifier, notifier: AdminNotifier,
thinking_budget: int | None = None,
) -> None: ) -> None:
self._model = model self._model = model
self._thinking_budget = thinking_budget
self._kb = kb self._kb = kb
self._store = store self._store = store
self._notifier = notifier self._notifier = notifier
@@ -98,7 +100,7 @@ class EmailAgent:
system = build_system_prompt(self._kb, state) system = build_system_prompt(self._kb, state)
try: 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) parsed = self._parse_llm_response(content)
except Exception: except Exception:
logger.exception("LLM call failed for %s", state.conversation_id) logger.exception("LLM call failed for %s", state.conversation_id)
@@ -140,7 +142,7 @@ class EmailAgent:
system = build_system_prompt(self._kb, state) system = build_system_prompt(self._kb, state)
try: 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) parsed = self._parse_llm_response(content)
except Exception: except Exception:
logger.exception("LLM call failed (post-completion) for %s", state.conversation_id) 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.""" """Build the system prompt sent to the LLM on every turn."""
import json import json
from datetime import date
from ..knowledge_base.loader import KnowledgeBase from ..knowledge_base.loader import KnowledgeBase
from ..models.conversation import ConversationState 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) reg_json = json.dumps(state.registration.to_dict(), ensure_ascii=False, indent=2)
step_hint = STEP_DESCRIPTIONS.get(state.flow_step, "Continue the conversation.") 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. 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} {_PERSONALITY}
## Registration Flow (8 steps) ## 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) reg_json = json.dumps(state.registration.to_dict(), ensure_ascii=False, indent=2)
child_name = state.registration.child.full_name or "their child" 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. 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} {_PERSONALITY}
## Context: Registration Already Complete ## Context: Registration Already Complete
+10
View File
@@ -50,6 +50,11 @@ class Config:
# Polling interval in seconds # Polling interval in seconds
poll_interval: int = 60 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 @classmethod
def from_env(cls) -> "Config": def from_env(cls) -> "Config":
return cls( return cls(
@@ -74,4 +79,9 @@ class Config:
) )
), ),
poll_interval=int(os.getenv("POLL_INTERVAL", "60")), 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 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. """Call any LLM asynchronously and return the response text.
This is the async equivalent of ``complete()`` — use this from async 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". model: litellm model string, e.g. "anthropic/claude-opus-4-6".
system: System prompt text. system: System prompt text.
messages: List of objects with .role and .content attributes. messages: List of objects with .role and .content attributes.
thinking_budget: When set, enables extended thinking (Anthropic models
only). See ``complete()`` for details.
Returns: Returns:
The model's reply as a plain string. The model's reply as a plain string.
""" """
api_messages = [{"role": "system", "content": system}] api_messages = [{"role": "system", "content": system}]
api_messages += [{"role": m.role, "content": m.content} for m in messages] 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 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. """Call any LLM and return the response text.
Args: Args:
@@ -35,13 +53,24 @@ def complete(model: str, system: str, messages: list) -> str:
environment variable (ANTHROPIC_API_KEY, OPENAI_API_KEY, …). environment variable (ANTHROPIC_API_KEY, OPENAI_API_KEY, …).
system: System prompt text. system: System prompt text.
messages: List of objects with .role and .content attributes. 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: Returns:
The model's reply as a plain string. The model's reply as a plain string.
""" """
api_messages = [{"role": "system", "content": system}] api_messages = [{"role": "system", "content": system}]
api_messages += [{"role": m.role, "content": m.content} for m in messages] 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 return response.choices[0].message.content