""" core/think.py Aiko's cognitive loop. - Retrieves relevant memories before each turn (scoped by user_id) - Tool routing: LLM-driven tool calling (preferred), with regex-based intent detection as fallback for weather, timezone, currency, joke, anime, and web search - Streams LLM response via token_callback - Stores the turn into long-term memory after each response (background thread) - Supports single-shot reasoning mode via set_reasoning(True) / /think command """ import os import json from datetime import datetime import httpx from pathlib import Path import queue import re import threading from core.memorize import AikoMemorize from core.log import get_logger log = get_logger(__name__) # ── boot labels ─────────────────────────────────────────────────────────────── BOOT_LABELS = { 'think_start': 'Loading LLM client + persona...', 'think_warmup': 'Warming up language model...', } # ── config ──────────────────────────────────────────────────────────────────── LLAMA_BASE_URL = os.getenv("LLAMA_BASE_URL") if not LLAMA_BASE_URL: raise RuntimeError("LLAMA_BASE_URL is not set") LLAMA_API_KEY = os.getenv("LLAMA_API_KEY", "") CONTEXT_WINDOW_TURNS = int(os.getenv("CONTEXT_WINDOW_TURNS", 20)) _BASE_PREDICT = 400 _REASONING_SCALE = 3 _PERSONA_PATH = Path(__file__).resolve().parent.parent / "persona" / "soul.md" _DEFAULT_USER_ID = os.getenv("USER_ID", "Guest") def _render_persona(template: str, user_id: str) -> str: now = datetime.now().astimezone() today = now.strftime("%B %d, %Y") current_time = now.strftime("%Y-%m-%d %H:%M:%S %Z") tz_name = now.astimezone().tzinfo.tzname(now) or "UTC" time_line = f"\nCurrent server time: {current_time} ({tz_name})" return ( template .replace("USER_ID_HERE", user_id) .replace("TODAY_HERE", today + time_line) ) # ── think ───────────────────────────────────────────────────────────────────── class AikoThink: """ Aiko's conversational core. LLM warmup starts immediately on init in a background thread. wakeup.py calls join_warmup() to block until the model is hot. speak is accepted but ignored (kept for BootResult API compatibility). user_id is tracked per-instance and updated via set_system_prompt() when the HF OAuth login resolves the real username. All memory operations (search + store) are scoped to the current user_id so memories never bleed between users. """ def __init__(self, memorize: AikoMemorize, speak=None) -> None: headers = {"Content-Type": "application/json"} if LLAMA_API_KEY: headers["Authorization"] = f"Bearer {LLAMA_API_KEY}" self._client = httpx.Client( base_url=LLAMA_BASE_URL, headers=headers, timeout=120.0, ) self._memorize = memorize self._user_id = _DEFAULT_USER_ID # updated on HF login via set_system_prompt() if not _PERSONA_PATH.exists(): raise FileNotFoundError(f"soul.md not found at {_PERSONA_PATH}") self._persona_raw = _PERSONA_PATH.read_text(encoding="utf-8").strip() # Live rendered system prompt — set_system_prompt() replaces this self._system_prompts: dict[str, str] = {} self._histories: dict[str, list[dict]] = {} self._reasoning = False self._token_callback = None self._mem_queue = queue.Queue() self._mem_worker = threading.Thread(target=self._mem_write_loop, daemon=True) self._mem_worker.start() # Internal warmup — just checks the LLM endpoint is reachable. # The real KV-cache warmup with soul.md is handled by wakeup._warmup_llm() # which runs concurrently. This one only fires a tiny probe so # join_warmup() can confirm the network path is alive. self._warmup_thread = threading.Thread(target=self._probe_llm, daemon=True) self._warmup_thread.start() def _probe_llm(self) -> None: """Lightweight connectivity probe — does NOT duplicate the full warmup.""" try: self._client.post( "/", json={ "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}], "temperature": 0.1, }, timeout=60, ) log.info("LLM probe complete") except Exception as e: log.warning("LLM probe failed (non-fatal): %s", e) # ── public api ──────────────────────────────────────────────────────────── def join_warmup(self) -> None: if self._warmup_thread.is_alive(): self._warmup_thread.join() def set_system_prompt(self, rendered_soul: str, user_id: str | None = None) -> None: """ Replace the active system prompt with a fully rendered soul.md string. Called by app.py _check_login() after HF OAuth resolves the username. Also updates user_id so memory ops are scoped to the logged-in user. Clears conversation history so the new persona starts fresh. """ if user_id: self._user_id = user_id log.info("System prompt updated for user: %s", user_id) effective_user_id = user_id or self._user_id or _DEFAULT_USER_ID self._system_prompts[effective_user_id] = rendered_soul self._histories.setdefault(effective_user_id, []).clear() def chat(self, user_input: str, user_id: str | None = None, token_callback=None) -> str: self._token_callback = token_callback # Resolve effective user_id: explicit arg > instance state > env default effective_user_id = user_id or self._user_id or _DEFAULT_USER_ID # 1. retrieve relevant long-term memories (scoped to this user) if self._memorize: memories = self._memorize.search(user_input, user_id=effective_user_id, limit=int(os.getenv("MEMORY_RECALL_LIMIT", 5))) memory_block = self._memorize.format_for_context(memories) else: memories = [] memory_block = None # 2. build system prompt — rendered persona + injected memories # Use the live _system_prompts if set (post-login), otherwise render fresh if effective_user_id in self._system_prompts: system = self._system_prompts[effective_user_id] else: system = _render_persona(self._persona_raw, effective_user_id) if memory_block: system = f"{system}\n\n{memory_block}" # 3. tool routing — try LLM-driven tool calling first, fall back to regex tool_result = None tool_tag = None user_history = self._histories.setdefault(effective_user_id, []) history_for_check = self._sanitize_history( user_history[-(CONTEXT_WINDOW_TURNS * 2):] + [{"role": "user", "content": user_input}] ) tool_tag, tool_result = self._try_tool_call(history_for_check, system) if tool_result is None: # fallback: regex-based intent detection (unchanged behavior) from core.tools import ( is_search_intent, is_weather_intent, is_timezone_intent, is_currency_intent, is_joke_intent, is_anime_intent, is_crypto_intent, is_camera_see_intent, extract_search_query, extract_location, extract_currency_parts, extract_anime_query, extract_crypto_parts, web_search_and_fetch, get_weather, get_timezone, get_currency, get_joke, get_anime, get_crypto_price, capture_camera_image, ) if is_joke_intent(user_input): tool_result = get_joke() tool_tag = "joke" elif is_weather_intent(user_input): location = extract_location(user_input) if location: if token_callback: token_callback(f"__TOOL__:Checking weather for {location}...") tool_result = get_weather(location) tool_tag = "weather_data" elif is_timezone_intent(user_input): location = extract_location(user_input) if location: if token_callback: token_callback(f"__TOOL__:Looking up time in {location}...") tool_result = get_timezone(location) tool_tag = "time_data" elif is_currency_intent(user_input): amount, from_cur, to_cur = extract_currency_parts(user_input) if token_callback: token_callback(f"__TOOL__:Converting {from_cur} to {to_cur}...") tool_result = get_currency(amount, from_cur, to_cur) tool_tag = "currency_data" elif is_crypto_intent(user_input): coin, currency = extract_crypto_parts(user_input) if token_callback: token_callback(f"__TOOL__:Checking {coin} price...") tool_result = get_crypto_price(coin, currency) tool_tag = "crypto_data" elif is_anime_intent(user_input): query = extract_anime_query(user_input) if token_callback: token_callback(f"__TOOL__:Searching anime for {query}...") tool_result = get_anime(query) tool_tag = "anime_data" elif is_camera_see_intent(user_input): if token_callback: token_callback("__TOOL__:Opening camera...") tool_result = capture_camera_image() tool_tag = "camera_view" elif is_search_intent(user_input): query = extract_search_query(user_input) if token_callback: token_callback(f"__SEARCHING__:{query}") try: tool_result = web_search_and_fetch(query) tool_tag = "search_results" except Exception as exc: log.warning("Web search failed: %s", exc) if tool_result and tool_tag: system = ( f"{system}\n\n" f"<{tool_tag}>\n{tool_result}\n\n\n" f"Use the above {tool_tag.replace('_', ' ')} to inform your response naturally. " f"Don't recite raw data — weave it into your answer as Aiko would." ) # 4. wrap user turn with reasoning instruction if active if self._reasoning: prompt = ( f"{user_input}\n\n" "Think through this carefully before answering. " "Show your reasoning inside tags, then give your final answer." ) else: prompt = user_input # If the camera tool was triggered, signal to open it immediately and bypass LLM generation if tool_result == "__OPEN_CAMERA__": if token_callback: token_callback("__OPEN_CAMERA__") response_text = "Sure! Let me open the camera so I can take a look~ 📷" user_history.append({"role": "user", "content": prompt}) user_history.append({"role": "assistant", "content": response_text}) self._store_async(user_input, response_text, effective_user_id) self._reasoning = False return response_text # 5. append user turn user_history.append({"role": "user", "content": prompt}) # 6. trim history to context window trimmed = self._sanitize_history(user_history[-(CONTEXT_WINDOW_TURNS * 2):]) # 7. LLM call response_text, _ = self._stream_response(trimmed, system=system) # 8. remove orphaned user turn on empty response if not response_text: if user_history and user_history[-1]["role"] == "user": user_history.pop() # 9. append assistant turn to history user_history.append({"role": "assistant", "content": response_text}) # 10. persist to memory (background), scoped to effective user self._store_async(user_input, response_text, effective_user_id) # 11. auto-reset reasoning mode self._reasoning = False return response_text def reset_context(self, user_id: str | None = None) -> None: """Clear the in-memory conversation history for a fresh session.""" effective_user_id = user_id or self._user_id or _DEFAULT_USER_ID if effective_user_id in self._histories: self._histories[effective_user_id].clear() def last_turn(self, user_id: str | None = None) -> tuple[str, str] | None: """Return the latest complete user/assistant exchange, or None.""" effective_user_id = user_id or self._user_id or _DEFAULT_USER_ID user_history = self._histories.get(effective_user_id, []) assistant_text: str | None = None for message in reversed(user_history): role = message.get("role") content = (message.get("content") or "").strip() if not content: continue if assistant_text is None: if role == "assistant": assistant_text = content continue if role == "user": return content, assistant_text return None def set_reasoning(self, enabled: bool) -> None: """Enable or disable reasoning mode for the next turn only.""" self._reasoning = enabled def wait_for_memory(self) -> None: """Block until all enqueued memory writes have been persisted.""" self._mem_queue.join() # ── internal ────────────────────────────────────────────────────────────── def _try_tool_call(self, messages: list[dict], system: str) -> tuple[str | None, str | None]: """ Ask the LLM whether a tool should be called for this turn. Sends the conversation + tool schemas with tool_choice="auto" and a short max_tokens budget (this is a routing decision, not the final answer). If the model returns tool_calls, dispatch the first one to its Python implementation and return (tag, result) for context injection. Returns (None, None) if no tool call was made, the tool name/args were invalid, or the request failed for any reason — in which case the caller falls back to regex-based intent detection. """ from core.tools import TOOL_SCHEMAS, TOOL_DISPATCH try: response = self._client.post( "/", json={ "messages": [{"role": "system", "content": system}] + messages, "tools": TOOL_SCHEMAS, "tool_choice": "auto", "stream": False, "temperature": 0.2, "max_tokens": 150, }, ) data = response.json() msg = data.get("choices", [{}])[0].get("message", {}) calls = msg.get("tool_calls") if not calls: return None, None call = calls[0] name = call.get("function", {}).get("name") args_raw = call.get("function", {}).get("arguments", "{}") if name not in TOOL_DISPATCH: log.warning("LLM requested unknown tool: %s", name) return None, None try: args = json.loads(args_raw) if isinstance(args_raw, str) else (args_raw or {}) except (ValueError, TypeError) as exc: log.warning("Failed to parse tool args for %s: %s (%r)", name, exc, args_raw) return None, None tag, fn = TOOL_DISPATCH[name] if self._token_callback: self._token_callback(f"__TOOL__:Calling {name}...") try: result = fn(**args) except Exception as exc: log.warning("Tool execution failed for %s: %s", name, exc) return None, None return tag, result except Exception as exc: log.warning("Tool-call attempt failed: %s", exc) return None, None def _stream_response(self, messages: list[dict], system: str = "") -> tuple[str, None]: num_predict = _BASE_PREDICT * _REASONING_SCALE if self._reasoning else _BASE_PREDICT try: response = self._client.post( "/", json={ "messages": ([{"role": "system", "content": system}] + messages) if system else messages, "stream": False, "temperature": float(os.getenv("LLAMA_TEMPERATURE", 0.75)), "max_tokens": num_predict, "top_p": float(os.getenv("LLAMA_TOP_P", 0.90)), "top_k": int(os.getenv("LLAMA_TOP_K", 40)), "repeat_penalty": float(os.getenv("LLAMA_REPEAT_PENALTY", 1.18)), "stop": ["<|im_end|>", "", "[INST]"], }, ) data = response.json() full_text = data.get("choices", [{}])[0].get("message", {}).get("content", "") or "" clean_text = re.sub(r"\[?SEARCH:\s*.+?\]?", "", full_text, flags=re.IGNORECASE).strip() if self._token_callback and clean_text: self._token_callback(clean_text) except Exception as exc: msg = f"Stream failed: {exc}" log.error(msg) if self._token_callback: self._token_callback(f"[think] {msg}") return "", None return full_text, None def _sanitize_history(self, messages: list[dict]) -> list[dict]: """Enforce strict user/assistant alternation.""" if not messages: return [] sanitized = [messages[0]] for msg in messages[1:]: if msg["role"] == sanitized[-1]["role"]: sanitized[-1] = msg else: sanitized.append(msg) while sanitized and sanitized[0]["role"] != "user": sanitized.pop(0) return sanitized def _store_async(self, user_input: str, response_text: str, user_id: str) -> None: self._mem_queue.put((user_input, response_text, user_id)) def _mem_write_loop(self) -> None: """Serial background worker that drains the memory write queue.""" while True: user_input, response_text, user_id = self._mem_queue.get() try: if self._memorize: self._memorize.add( [ {"role": "user", "content": user_input[:500]}, {"role": "assistant", "content": response_text[:800]}, ], user_id=user_id, ) except Exception as exc: log.error("Async memory write failed: %s", exc) finally: self._mem_queue.task_done()