diff --git a/ai_commands.py b/ai_commands.py index df54bfe..0f3883c 100644 --- a/ai_commands.py +++ b/ai_commands.py @@ -1,4 +1,5 @@ # ai command cogs +import asyncio import logging import re import sys @@ -17,7 +18,7 @@ from communication_subroutine import AI_QUERY_Q import ai_functions from constants import ( - ASSISTANTS, + OLLAMA_WARM_MINUTES, DATA, GRAPHICS_PATH, INITIAL_TIME_WAIT, @@ -101,74 +102,45 @@ class Events(commands.Cog): text = text[1900:] async def cog_load(self): - # The AI query worker must run regardless of the OpenAI guard below - it - # answers via handle_response, which works on Claude too. Start it first. + # The AI query worker answers via handle_response, so it works on every + # backend. Start it first. if not self.ai_query_worker.is_running(): self.ai_query_worker.start() - self.logger.info("Starting personal assistants") - # Personal assistants use the OpenAI Assistants API (threads/runs), which - # has no Anthropic equivalent - skip cleanly when OpenAI isn't wired up - # (e.g. a Claude-only deployment) instead of crashing the cog load. - if OPENAICLIENT is None: - self.logger.warning( - "OPENAICLIENT niedostępny - osobiści asystenci (OpenAI Assistants API) wyłączeni" - ) - return - # The bootstrap below must NEVER take the cog down with it. It calls the - # OpenAI Assistants API (beta threads/runs), which is a legacy surface - - # it now answers 404, and that exception propagated out of cog_load, - # failed the whole extension, and took EVERY AI command with it - # ($gadaj_teraz, $modele_ai, the conversation handler). Personal - # assistants are one optional feature; losing them must not disable the - # AI cog, which otherwise works fine on Claude and Ollama. + # Keeps a self-hosted model resident; it no-ops on any other provider. + if not self.ollama_warm_loop.is_running(): + self.ollama_warm_loop.start() + # NOTE: there is no OpenAI-Assistants bootstrap any more. It called a + # sunset API (beta threads), 404'd, and failed the WHOLE extension - + # taking every AI command with it. Personal assistants now ride + # handle_response with per-user memory (ai_functions), so they work on + # Claude and Ollama too and nothing has to be created at startup. + self.logger.info("Osobiści asystenci: pamięć per-user, aktywny backend AI") + + @tasks.loop(minutes=OLLAMA_WARM_MINUTES) + async def ollama_warm_loop(self): + """Keep a self-hosted model resident so users don't pay the load wait. + + Loading is the slow part on a GPU shared with other users, so we + re-assert Ollama's keep_alive well inside its window. This preloads + WITHOUT generating - no tokens, no cost. + + Hard guard: it does nothing unless the ACTIVE backend is Ollama. Firing + warm-ups at a metered API would burn tokens and money for nothing. + """ try: - await self._start_personal_assistants() + if ai_functions.active_provider() != "ollama": + return + await ai_functions.warm_active_model() except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning( - "Osobiści asystenci (OpenAI Assistants API) wyłączeni - %s: %s. " - "Reszta AI (rozmowy, $gadaj_teraz, $modele_ai) działa normalnie.", - type(exc).__name__, exc, - ) + self.logger.info("Rozgrzewanie Ollamy nieudane (nieszkodliwe): %s", exc) - async def _start_personal_assistants(self): - """Bootstrap the per-user OpenAI Assistants threads. Optional feature: - callers must treat a failure here as non-fatal (see cog_load).""" - for superfryta_id, superfryta in SPECJALNE_ZIEMNIACZKI.items(): - - if superfryta[4] != "": - self.logger.info( - "Personal assistant for user: %s, exists id: %s,name: %s, owner: %s, special instructions: %s assistant id: %s ", - superfryta_id, - superfryta[0], - superfryta[1], - superfryta[2], - superfryta[3], - superfryta[4], - ) - thread = await OPENAICLIENT.beta.threads.create() - self.logger.info("Thread id: %s", thread.id) - ASSISTANTS[superfryta[1]] = ( - superfryta[2], - superfryta[4], - superfryta[0], - thread, - ) - else: - self.logger.info( - "Creating personal assistant for user: %s, id: %s,name: %s, owner: %s, special instructions: %s", - superfryta_id, - superfryta[0], - superfryta[1], - superfryta[2], - superfryta[3], - ) - await ai_functions.create_chat_assistant( - superfryta_id, superfryta[0], superfryta[1], superfryta[2], superfryta[3] - ) - self.logger.info("Started personal assistants") + @ollama_warm_loop.before_loop + async def before_ollama_warm_loop(self): + await self.bot.wait_until_ready() async def cog_unload(self): self.ai_query_worker.cancel() + self.ollama_warm_loop.cancel() @commands.hybrid_command( name="switch_dm_mode", @@ -297,6 +269,15 @@ class Events(commands.Cog): f"Teraz gadam przez **{nazwa_konfigu}** — " f"{cfg.get('provider')} / {cfg.get('latest_model')}." ) + if cfg.get("provider") == "ollama": + # Pay the (slow, shared-GPU) load cost NOW, in the background, + # so it lands on the operator switching backends rather than on + # whoever sends the first message. Not awaited: loading can take + # minutes and the command must answer immediately. + asyncio.create_task(ai_functions.warm_active_model()) + message += ( + "\nRozgrzewam model w tle — pierwsza odpowiedź może chwilę potrwać." + ) # Switched without pinning a model: show what else is on offer. if not model: try: @@ -430,8 +411,10 @@ class Events(commands.Cog): if message.author.id == superfryta[0]: self.logger.info("Specjalny ziemniak") if self.armia[message.author.id] == Dm_Mode.SPECJALNY_ZIEMNIACZEK: - #await self.bot.process_commands(message) - await ai_functions.chat_with_assistant(message, superfryta[1]) + # superfryta = [discord_id, assistant_name, owner, instructions, legacy_assistant_id] + await ai_functions.chat_with_personal_assistant( + message, superfryta[2], superfryta[3] + ) return elif self.armia[message.author.id] == Dm_Mode.ECHO_ECHO: await ai_functions.echo(message) diff --git a/ai_functions.py b/ai_functions.py index b629c36..21f33f5 100644 --- a/ai_functions.py +++ b/ai_functions.py @@ -1,16 +1,21 @@ import asyncio import json import logging +import os import random +import tempfile import openai import tiktoken import time from other_functions import discord_friendly_send +import requests + from constants import ( AI_CONFIGS, AI_TIMEOUT_SECONDS, - ASSISTANTS, + ASSISTANT_MEMORY_FILE, + ASSISTANT_MEMORY_TURNS, CLAUDECLIENT, CYCLIC_WORDS, DEFAULT_AI_CONFIG, @@ -21,6 +26,9 @@ from constants import ( MESSAGE_TABLE, MESSAGE_TABLE_MUZYKA, OLLAMACLIENT, + OLLAMA_KEEP_ALIVE, + OLLAMA_PRELOAD_TIMEOUT, + OLLAMA_URL, OPENAICLIENT, SYSTEM_GPT_SETTINGS, WORD_REACTIONS, @@ -290,6 +298,52 @@ async def _ollama_call(messages, model, cfg): return (resp.choices[0].message.content or "").strip() +def _ollama_preload(model, keep_alive=None) -> bool: + """Load ``model`` into Ollama and keep it resident, generating NOTHING. + + Ollama's /api/generate with a model and no prompt is the documented preload: + it pays the (slow, GPU-shared) load cost once and returns, producing no + tokens. Used to warm up on switch and to re-assert keep_alive periodically. + + Blocking on purpose - callers wrap it in asyncio.to_thread. + """ + if not OLLAMA_URL: + return False + logger = logging.getLogger("discord") + try: + resp = requests.post( + f"{OLLAMA_URL}/api/generate", + json={"model": model, "keep_alive": keep_alive or OLLAMA_KEEP_ALIVE}, + timeout=OLLAMA_PRELOAD_TIMEOUT, + ) + ok = resp.status_code == 200 + logger.info("Ollama preload %s -> HTTP %s", model, resp.status_code) + return ok + except requests.exceptions.RequestException as exc: + logger.info("Ollama preload %s failed: %s", model, exc) + return False + + +def active_provider() -> str: + """Provider of the active config - the guard every warm-up must check. + + Preloading only makes sense for a self-hosted model; firing it at a metered + API would burn tokens (and money) for nothing. + """ + return (_active_config() or {}).get("provider", "") + + +async def warm_active_model(force_model=None) -> bool: + """Preload the active model IFF the active backend is Ollama.""" + if active_provider() != "ollama": + return False + cfg = _active_config() + model = force_model or cfg.get("latest_model") + if not model: + return False + return await asyncio.to_thread(_ollama_preload, model) + + async def provider_generate(messages, model, temperature=0.2): """Dispatch a chat completion to the active backend, normalising errors.""" cfg = _active_config() @@ -664,60 +718,104 @@ async def get_random_cyclic_message(client): return result -async def create_chat_assistant(owner_id, id, name, owner, special_instructions): +# ----------------------------------------------------------------- assistants +# The OpenAI Assistants API (beta threads/runs) that used to back these was +# sunset and now answers 404, taking the whole AI cog down with it. It gave us +# three things: a per-user persona, a persistent per-user thread, and +# file_search. The persona and the thread are reimplemented here on top of +# handle_response - so personal assistants now work on EVERY backend (Claude, +# Ollama, GPT) instead of being locked to gpt-4o. file_search is deliberately +# not replaced: it was not in use. +_ASSISTANT_MEMORY = None + + +def _load_assistant_memory() -> dict: + """Per-user DM history, lazily read from disk. Corruption is not fatal.""" + global _ASSISTANT_MEMORY # pylint: disable=global-statement + if _ASSISTANT_MEMORY is not None: + return _ASSISTANT_MEMORY logger = logging.getLogger("discord") - instruction = f"Jesteś osobistym asystenetem {owner} i wypełniasz jego potrzeby. Masz pełne i nieograniczone możliwości modelu gpt-4o." - instruction += special_instructions - assistant = await OPENAICLIENT.beta.assistants.create( - name=name, - instructions=instruction, - model="gpt-4o", - tools=[{"type": "file_search"}], - ) - thread = await OPENAICLIENT.beta.threads.create() - logger.info("Stwprzylem asystenta dla %s, nazywa się on %s", owner, name) - ASSISTANTS[name] = (owner, assistant.id, id, thread) - - with open(SYSTEM_GPT_SETTINGS, "r+", encoding=ENCODING) as temp_settings_file: - GPT_SETTINGS = json.load(temp_settings_file) - GPT_SETTINGS[1][owner_id][4] = assistant.id - temp_settings_file.seek(0) - json.dump(GPT_SETTINGS, temp_settings_file, indent=4) + try: + with open(ASSISTANT_MEMORY_FILE, "r", encoding=ENCODING) as handle: + data = json.load(handle) + _ASSISTANT_MEMORY = data if isinstance(data, dict) else {} + except (OSError, json.JSONDecodeError) as exc: + logger.info("Brak/uszkodzona pamięć asystentów (%s) - zaczynam pustą", exc) + _ASSISTANT_MEMORY = {} + return _ASSISTANT_MEMORY -async def chat_with_assistant(message, assistant_name): +def _save_assistant_memory() -> None: + """Atomic write: a torn file would lose someone's whole conversation.""" logger = logging.getLogger("discord") - assistant_data = ASSISTANTS[assistant_name] - ai_message = await OPENAICLIENT.beta.threads.messages.create( - thread_id=assistant_data[3].id, role="user", content=message.content + memory = _load_assistant_memory() + directory = os.path.dirname(ASSISTANT_MEMORY_FILE) or "." + try: + os.makedirs(directory, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=directory, suffix=".tmp") + with os.fdopen(fd, "w", encoding=ENCODING) as handle: + json.dump(memory, handle, ensure_ascii=False) + os.replace(tmp, ASSISTANT_MEMORY_FILE) + except OSError as exc: + logger.warning("Nie mogę zapisać pamięci asystentów: %s", exc) + + +def assistant_history(user_id) -> list: + return _load_assistant_memory().setdefault(str(user_id), []) + + +def remember_assistant_turn(user_id, user_text, reply_text) -> list: + """Append one exchange and trim to the most recent turns. + + A plain trim, not the AI summarisation used for the bar's shared memory: + these are private DMs and must not end up in a public 'legend'. + """ + history = assistant_history(user_id) + history.append({"role": "user", "content": user_text}) + history.append({"role": "assistant", "content": reply_text}) + if len(history) > ASSISTANT_MEMORY_TURNS: + del history[: len(history) - ASSISTANT_MEMORY_TURNS] + _save_assistant_memory() + return history + + +def build_assistant_messages(user_id, owner, special_instructions, prompt) -> list: + """System persona + this user's own history + the new turn.""" + system = ( + f"Jesteś osobistym asystentem {owner} i wypełniasz jego potrzeby. " + f"{special_instructions or ''}" + ).strip() + return ( + [{"role": "system", "content": system}] + + list(assistant_history(user_id)) + + [{"role": "user", "content": prompt}] ) - logger.info(ai_message) - run = await OPENAICLIENT.beta.threads.runs.create_and_poll( - thread_id=assistant_data[3].id, - assistant_id=assistant_data[1], - instructions=f"Pisze do Ciebie {assistant_data[0]} udziel mu wszelkiej pomocy", + + +async def chat_with_personal_assistant(message, owner, special_instructions): + """Answer a DM as this user's personal assistant, on the active backend. + + request_type="NONE" with an explicit message list keeps this OUT of the + bar's shared memory - the conversation is carried by the per-user history + built above and stored separately. + """ + logger = logging.getLogger("discord") + user_id = message.author.id + prompt = message.content + messages = build_assistant_messages(user_id, owner, special_instructions, prompt) + result, _table = await handle_response( + prompt, + False, + False, + [], + str(owner), + "NONE", + none_request=messages, ) - done = False - while not done: - if run.status == "completed": - messsages = await OPENAICLIENT.beta.threads.messages.list( - thread_id=assistant_data[3].id - ) - logger.info(messsages) - reply_content = messsages.data[0].content - logger.info(reply_content) - chat_response = "" - for block in reply_content: - logger.info(block.text.value) - chat_response += block.text.value - await discord_friendly_send(message.channel, chat_response) - # await message.channel.send(chat_response) - done = True - elif run.status == "cancelled": - await discord_friendly_send(message.channel, "Cos sie wywaliło") - else: - logger.info(run.status) - asyncio.sleep(5) + remember_assistant_turn(user_id, prompt, result) + logger.info("Asystent odpowiedział %s (%d znaków)", owner, len(result or "")) + await discord_friendly_send(message.channel, result) + return result async def echo(message): diff --git a/constants.py b/constants.py index 4741d43..86acdcc 100644 --- a/constants.py +++ b/constants.py @@ -245,6 +245,16 @@ DELIVERED_DIR = os.getenv( ) DELIVERED_MAX = int(os.getenv("CONJURER_DELIVERED_MAX", "10000")) +# Personal DM assistants. Replaces the OpenAI Assistants API (threads/runs), +# which was sunset and answers 404: the persona now rides handle_response, so it +# works on EVERY backend, and the conversation lives here instead of on OpenAI's +# server. Kept per user so private DMs never bleed into the bar's shared memory, +# and trimmed to the most recent turns so it cannot grow without bound. +ASSISTANT_MEMORY_FILE = os.getenv( + "CONJURER_ASSISTANT_MEMORY", os.path.join(_STATE_ROOT, "assistant_memory.json") +) +ASSISTANT_MEMORY_TURNS = int(os.getenv("CONJURER_ASSISTANT_MEMORY_TURNS", "40")) + FILE_SERVICE_ADDRESS = os.getenv("CONJURER_FILE_SERVICE", "http://192.168.1.15:5000") RADIO_HARBOR_ADDRESS = os.getenv("CONJURER_RADIO_HARBOR", "http://192.168.1.15:54321") # Betoniarka (radio-operator service colocated with Liquidsoap). Falls back to @@ -411,6 +421,19 @@ else: AI_TIMEOUT_SECONDS = int(os.getenv("CONJURER_AI_TIMEOUT_SECONDS", "120")) OLLAMA_URL = os.getenv("CONJURER_OLLAMA_URL", "").rstrip("/") +# Keeping a self-hosted model resident. Loading it is the slow part (it is +# offloaded to a GPU shared with other users), so we preload it - Ollama's +# /api/generate with a model and NO prompt loads it and generates nothing, which +# costs no tokens and no money. KEEP_ALIVE is how long Ollama should then hold +# it; the warm loop re-asserts that well inside the window. +# STRICTLY Ollama-only: doing this against a paid API would burn tokens for +# nothing, so every caller checks the active provider first. +OLLAMA_KEEP_ALIVE = os.getenv("CONJURER_OLLAMA_KEEP_ALIVE", "30m") +OLLAMA_WARM_MINUTES = float(os.getenv("CONJURER_OLLAMA_WARM_MINUTES", "10")) +# A preload waits for the model to finish loading, which on a shared GPU is the +# slow path we are trying to move off the user's first message. +OLLAMA_PRELOAD_TIMEOUT = int(os.getenv("CONJURER_OLLAMA_PRELOAD_TIMEOUT", "600")) + OLLAMA_LATEST_MODEL = os.getenv("CONJURER_OLLAMA_MODEL", "llama3.1:8b") OLLAMA_CHEAP_MODEL = os.getenv("CONJURER_OLLAMA_CHEAP_MODEL", OLLAMA_LATEST_MODEL) if openai and OLLAMA_URL: diff --git a/tests/unit/test_ai_provider_switch.py b/tests/unit/test_ai_provider_switch.py index 0d30069..b98a876 100644 --- a/tests/unit/test_ai_provider_switch.py +++ b/tests/unit/test_ai_provider_switch.py @@ -394,3 +394,143 @@ def test_persist_survives_an_unreadable_settings_file(tmp_path, monkeypatch): broken.write_text("{ not json", encoding="utf-8") monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(broken)) ai_functions._persist_active_ai_config("gpt", model_for="gpt") # must not raise + + +# ----------------------------------------------- keep-warm (Ollama ONLY) ---- +# The money guard: preloading a self-hosted model is free, but firing the same +# thing at a metered API would burn tokens for nothing. These pin that it can +# only ever happen for Ollama. + + +def test_warm_active_model_is_a_noop_for_paid_providers(monkeypatch): + called = [] + monkeypatch.setattr( + ai_functions, "_ollama_preload", lambda *a, **k: called.append(a) or True + ) + for paid in ("gpt", "claude"): + _reset_active(paid) + try: + assert asyncio.run(ai_functions.warm_active_model()) is False + finally: + _reset_active("gpt") + assert called == [], "a paid backend must never be preloaded" + + +def test_warm_active_model_preloads_when_ollama_is_active(monkeypatch): + monkeypatch.setitem( + ai_functions.AI_CONFIGS, + "ollama", + {"provider": "ollama", "latest_model": "qwen2.5:7b", "cheap_model": "c"}, + ) + seen = {} + monkeypatch.setattr( + ai_functions, "_ollama_preload", lambda model, *a, **k: seen.update(model=model) or True + ) + _reset_active("ollama") + try: + assert asyncio.run(ai_functions.warm_active_model()) is True + finally: + _reset_active("gpt") + assert seen["model"] == "qwen2.5:7b" + + +def test_active_provider_reports_the_switch(): + _reset_active("gpt") + assert ai_functions.active_provider() == "openai" + _reset_active("claude") + try: + assert ai_functions.active_provider() == "anthropic" + finally: + _reset_active("gpt") + + +def test_preload_sends_no_prompt_so_it_generates_nothing(monkeypatch): + # Ollama's documented preload: a model and keep_alive, and NO prompt. If a + # prompt ever crept in, every warm-up would silently generate tokens. + sent = {} + + class _Resp: + status_code = 200 + + monkeypatch.setattr(ai_functions, "OLLAMA_URL", "http://ollama:11434") + monkeypatch.setattr( + ai_functions.requests, "post", + lambda url, json=None, timeout=None: sent.update(url=url, body=json) or _Resp(), + ) + assert ai_functions._ollama_preload("qwen2.5:7b") is True + assert sent["url"].endswith("/api/generate") + assert sent["body"]["model"] == "qwen2.5:7b" + assert "keep_alive" in sent["body"] + assert "prompt" not in sent["body"], "a preload must not generate" + + +def test_preload_without_endpoint_is_a_noop(monkeypatch): + monkeypatch.setattr(ai_functions, "OLLAMA_URL", "") + assert ai_functions._ollama_preload("x") is False + + +# ------------------------------------------- personal assistants (per user) -- +# Replaces the sunset OpenAI Assistants API. The two properties that matter: +# each user's DM history is ISOLATED (private DMs must not leak into another +# user's context or the bar's shared memory), and it stays BOUNDED. + + +def _fresh_assistant_memory(tmp_path, monkeypatch, turns=40): + monkeypatch.setattr( + ai_functions, "ASSISTANT_MEMORY_FILE", str(tmp_path / "assistant_memory.json") + ) + monkeypatch.setattr(ai_functions, "ASSISTANT_MEMORY_TURNS", turns) + monkeypatch.setattr(ai_functions, "_ASSISTANT_MEMORY", None) + + +def test_assistant_history_is_isolated_per_user(tmp_path, monkeypatch): + _fresh_assistant_memory(tmp_path, monkeypatch) + ai_functions.remember_assistant_turn(111, "sekret Anny", "ok Anna") + ai_functions.remember_assistant_turn(222, "sekret Bartka", "ok Bartek") + + anna = ai_functions.assistant_history(111) + bartek = ai_functions.assistant_history(222) + assert [m["content"] for m in anna] == ["sekret Anny", "ok Anna"] + assert [m["content"] for m in bartek] == ["sekret Bartka", "ok Bartek"] + assert "sekret Anny" not in str(bartek) # no cross-user bleed + + +def test_assistant_history_is_trimmed_to_the_bound(tmp_path, monkeypatch): + _fresh_assistant_memory(tmp_path, monkeypatch, turns=4) + for i in range(10): + ai_functions.remember_assistant_turn(1, f"u{i}", f"a{i}") + history = ai_functions.assistant_history(1) + assert len(history) == 4 # bounded + assert history[-1]["content"] == "a9" # newest kept + assert all("u0" != m["content"] for m in history) # oldest dropped + + +def test_assistant_history_survives_a_restart(tmp_path, monkeypatch): + _fresh_assistant_memory(tmp_path, monkeypatch) + ai_functions.remember_assistant_turn(7, "pamietaj", "pamietam") + # Simulate a restart: drop the in-memory cache, re-read from disk. + monkeypatch.setattr(ai_functions, "_ASSISTANT_MEMORY", None) + assert [m["content"] for m in ai_functions.assistant_history(7)] == [ + "pamietaj", + "pamietam", + ] + + +def test_assistant_messages_carry_persona_history_and_new_turn(tmp_path, monkeypatch): + _fresh_assistant_memory(tmp_path, monkeypatch) + ai_functions.remember_assistant_turn(5, "wczoraj", "odpowiedz") + msgs = ai_functions.build_assistant_messages( + 5, "Towarzysz Młotek", "Mówisz po polsku.", "dzisiaj" + ) + assert msgs[0]["role"] == "system" + assert "Towarzysz Młotek" in msgs[0]["content"] + assert "Mówisz po polsku." in msgs[0]["content"] + assert [m["content"] for m in msgs[1:]] == ["wczoraj", "odpowiedz", "dzisiaj"] + + +def test_corrupt_assistant_memory_starts_empty_instead_of_crashing(tmp_path, monkeypatch): + path = tmp_path / "assistant_memory.json" + path.write_text("{ not json", encoding="utf-8") + monkeypatch.setattr(ai_functions, "ASSISTANT_MEMORY_FILE", str(path)) + monkeypatch.setattr(ai_functions, "_ASSISTANT_MEMORY", None) + assert ai_functions.assistant_history(1) == []