diff --git a/.env.example b/.env.example index e6909f0..7562246 100644 --- a/.env.example +++ b/.env.example @@ -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 # --------------------------------------------------------------- diff --git a/chat_app.py b/chat_app.py index 1174ce3..18c67ab 100644 --- a/chat_app.py +++ b/chat_app.py @@ -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 diff --git a/main.py b/main.py index f91f579..b434cbe 100644 --- a/main.py +++ b/main.py @@ -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, diff --git a/src/agent/core.py b/src/agent/core.py index 2167cd5..403faa1 100644 --- a/src/agent/core.py +++ b/src/agent/core.py @@ -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) diff --git a/src/agent/prompts.py b/src/agent/prompts.py index 4330818..2d15e80 100644 --- a/src/agent/prompts.py +++ b/src/agent/prompts.py @@ -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 diff --git a/src/config.py b/src/config.py index a0371ea..fbba30e 100644 --- a/src/config.py +++ b/src/config.py @@ -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 + ), ) diff --git a/src/llm.py b/src/llm.py index d24b122..271722d 100644 --- a/src/llm.py +++ b/src/llm.py @@ -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