Compare commits
1 Commits
main
..
c36d6d3fc2
| Author | SHA1 | Date | |
|---|---|---|---|
| c36d6d3fc2 |
@@ -22,23 +22,6 @@ jobs:
|
||||
docker build -f docker/Dockerfile.bot -t gitea.czernobog.pl/gitea/conjurer-bot:$TAG .
|
||||
docker push gitea.czernobog.pl/gitea/conjurer-bot:$TAG
|
||||
|
||||
# DEPLOY channel: the production (deploy) bot tracks a SEPARATE image,
|
||||
# conjurer-bot-deploy, which only gets a new tag when the commit message
|
||||
# contains the trigger "[deploy]". So the test bot & librarian update on
|
||||
# every build, but the deploy bot only on commits you explicitly promote.
|
||||
# Same bytes as conjurer-bot:$TAG - just re-tagged, no rebuild.
|
||||
- name: Promote bot to deploy channel (only on [deploy] in commit message)
|
||||
run: |
|
||||
TAG=${GITHUB_SHA::8}
|
||||
MSG="$(git log -1 --pretty=%B)"
|
||||
if echo "$MSG" | grep -qiF '[deploy]'; then
|
||||
echo "Commit message contains [deploy] -> promoting conjurer-bot-deploy:$TAG"
|
||||
docker tag gitea.czernobog.pl/gitea/conjurer-bot:$TAG gitea.czernobog.pl/gitea/conjurer-bot-deploy:$TAG
|
||||
docker push gitea.czernobog.pl/gitea/conjurer-bot-deploy:$TAG
|
||||
else
|
||||
echo "No [deploy] trigger in commit message -> deploy bot image left unchanged"
|
||||
fi
|
||||
|
||||
# Docker-owe: musician i share lekkie, radio CIĘŻKIE (opam+OCaml+liquidsoap)
|
||||
- name: Build docker-stack images (musician, share)
|
||||
run: |
|
||||
|
||||
@@ -202,8 +202,3 @@ not_in_db.json
|
||||
rr_results.json
|
||||
s_results.json
|
||||
*.bak.DS_Store
|
||||
|
||||
# Durable result-delivery spool (runtime, per-deploy)
|
||||
result_inbox/
|
||||
delivered_uuids/
|
||||
outbox/
|
||||
|
||||
@@ -171,17 +171,6 @@ class AdministrationModule(commands.Cog):
|
||||
"""
|
||||
# logger.info("Heartbeat of cleanup proc")
|
||||
channel = self.bot.get_channel(1062047367337095268)
|
||||
if channel is None:
|
||||
# get_channel() returns None before the gateway cache is populated
|
||||
# (the before_loop below normally prevents that) or when the bot
|
||||
# cannot see the channel at all (wrong id / not in the guild /
|
||||
# missing permission). Skip this tick instead of crashing - an
|
||||
# unhandled exception here stops the whole loop, killing log
|
||||
# rollover and the spontaneous messages with it.
|
||||
self.logger.warning(
|
||||
"check_self: channel 1062047367337095268 unavailable - skipping tick"
|
||||
)
|
||||
return
|
||||
messages = [message async for message in channel.history(limit=1)]
|
||||
for mess in messages:
|
||||
channel = mess.channel
|
||||
@@ -265,13 +254,6 @@ class AdministrationModule(commands.Cog):
|
||||
self.logger.info(message)
|
||||
await channel.send(message)
|
||||
|
||||
@check_self.before_loop
|
||||
async def before_check_self(self):
|
||||
# Don't fire the first tick until the gateway is READY and the channel
|
||||
# cache is populated - get_channel() returns None before that, which is
|
||||
# exactly what used to crash check_self at startup.
|
||||
await self.bot.wait_until_ready()
|
||||
|
||||
|
||||
async def setup(bot):
|
||||
logger = logging.getLogger("discord")
|
||||
|
||||
+55
-153
@@ -1,5 +1,4 @@
|
||||
# ai command cogs
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
import sys
|
||||
@@ -18,7 +17,7 @@ from communication_subroutine import AI_QUERY_Q
|
||||
|
||||
import ai_functions
|
||||
from constants import (
|
||||
OLLAMA_WARM_MINUTES,
|
||||
ASSISTANTS,
|
||||
DATA,
|
||||
GRAPHICS_PATH,
|
||||
INITIAL_TIME_WAIT,
|
||||
@@ -102,45 +101,55 @@ class Events(commands.Cog):
|
||||
text = text[1900:]
|
||||
|
||||
async def cog_load(self):
|
||||
# The AI query worker answers via handle_response, so it works on every
|
||||
# backend. Start it first.
|
||||
# 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.
|
||||
if not self.ai_query_worker.is_running():
|
||||
self.ai_query_worker.start()
|
||||
# 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:
|
||||
if ai_functions.active_provider() != "ollama":
|
||||
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
|
||||
await ai_functions.warm_active_model()
|
||||
except Exception as exc: # pylint: disable=broad-exception-caught
|
||||
self.logger.info("Rozgrzewanie Ollamy nieudane (nieszkodliwe): %s", exc)
|
||||
for superfryta_id, superfryta in SPECJALNE_ZIEMNIACZKI.items():
|
||||
|
||||
@ollama_warm_loop.before_loop
|
||||
async def before_ollama_warm_loop(self):
|
||||
await self.bot.wait_until_ready()
|
||||
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")
|
||||
|
||||
async def cog_unload(self):
|
||||
self.ai_query_worker.cancel()
|
||||
self.ollama_warm_loop.cancel()
|
||||
|
||||
@commands.hybrid_command(
|
||||
name="switch_dm_mode",
|
||||
@@ -164,58 +173,20 @@ class Events(commands.Cog):
|
||||
else:
|
||||
await ctx.reply("Nope. Nie wiesz jak użyć")
|
||||
|
||||
@commands.hybrid_command(
|
||||
name="modele_ai",
|
||||
description="Pokaż modele dostępne dla danego backendu (Ollamę pyta na żywo).",
|
||||
)
|
||||
async def modele_ai(self, ctx, nazwa_konfigu: Optional[str] = None):
|
||||
"""Read-only model listing. For Ollama this queries the server, so it
|
||||
shows exactly what is pulled on the box right now."""
|
||||
async with ctx.channel.typing():
|
||||
target = nazwa_konfigu or ai_functions.get_active_ai_config()
|
||||
if target not in ai_functions.list_ai_configs():
|
||||
await discord_friendly_reply(
|
||||
ctx,
|
||||
f"Nie znam configu '{target}'. Dostępne: "
|
||||
f"{', '.join(ai_functions.list_ai_configs())}",
|
||||
)
|
||||
return
|
||||
try:
|
||||
models = await ai_functions.list_provider_models(target)
|
||||
except ai_functions.AIError as exc:
|
||||
await discord_friendly_reply(
|
||||
ctx, f"Nie mogę pobrać modeli dla '{target}': {exc}"
|
||||
)
|
||||
return
|
||||
if not models:
|
||||
await discord_friendly_reply(ctx, f"Brak modeli dla '{target}'.")
|
||||
return
|
||||
await discord_friendly_reply(
|
||||
ctx,
|
||||
f"Modele dla **{target}**: {', '.join(models)}\n"
|
||||
f"Wepniesz przez `$gadaj_teraz {target} <model>` (tylko Vykidailo).",
|
||||
)
|
||||
|
||||
@commands.hybrid_command(
|
||||
name="gadaj_teraz",
|
||||
description="Pokaż/przełącz backend AI i model (bez argumentu = status). Przełączanie: Vykidailo.",
|
||||
description="Pokaż/przełącz backend AI (bez argumentu = status). Przełączanie: Vykidailo.",
|
||||
)
|
||||
async def gadaj_teraz(
|
||||
self, ctx, nazwa_konfigu: Optional[str] = None, model: Optional[str] = None
|
||||
):
|
||||
async def gadaj_teraz(self, ctx, nazwa_konfigu: Optional[str] = None):
|
||||
async with ctx.channel.typing():
|
||||
available = ai_functions.list_ai_configs()
|
||||
# No argument -> report the active backend (read-only, open to all).
|
||||
if not nazwa_konfigu:
|
||||
active = ai_functions.get_active_ai_config()
|
||||
active_cfg = ai_functions.AI_CONFIGS.get(active, {})
|
||||
await discord_friendly_reply(
|
||||
ctx,
|
||||
f"Teraz gadam przez **{active}** "
|
||||
f"({active_cfg.get('provider')} / {active_cfg.get('latest_model')}). "
|
||||
f"Dostępne: {', '.join(available)}. "
|
||||
"Przełączysz przez `$gadaj_teraz <config> [model]` (tylko Vykidailo), "
|
||||
"modele zobaczysz przez `$modele_ai`.",
|
||||
f"Teraz gadam przez **{active}**. Dostępne: {', '.join(available)}. "
|
||||
"Przełączysz przez `$gadaj_teraz <config>` (tylko Vykidailo).",
|
||||
)
|
||||
return
|
||||
is_admin = isinstance(ctx.author, discord.Member) and any(
|
||||
@@ -237,69 +208,13 @@ class Events(commands.Cog):
|
||||
ctx, f"Nie mogę przełączyć na '{nazwa_konfigu}': {exc}"
|
||||
)
|
||||
return
|
||||
|
||||
# Optional second argument pins the model. For a backend we can
|
||||
# enumerate (Ollama), reject an unknown id up front with the list -
|
||||
# otherwise the typo only surfaces later as a failed reply.
|
||||
if model:
|
||||
try:
|
||||
known = await ai_functions.list_provider_models(nazwa_konfigu)
|
||||
except ai_functions.AIError:
|
||||
known = [] # cannot enumerate -> accept verbatim
|
||||
if known and cfg.get("provider") == "ollama" and model not in known:
|
||||
self.logger.info(
|
||||
"Przełączono AI na config %s (%s)", nazwa_konfigu, cfg.get("provider")
|
||||
)
|
||||
await discord_friendly_reply(
|
||||
ctx,
|
||||
f"Model '{model}' nie jest wgrany na Ollamie. "
|
||||
f"Dostępne: {', '.join(known)}",
|
||||
f"Teraz gadam przez **{nazwa_konfigu}** — {cfg.get('provider')} / {cfg.get('latest_model')}.",
|
||||
)
|
||||
return
|
||||
try:
|
||||
cfg = ai_functions.set_active_model(model, nazwa_konfigu)
|
||||
except (KeyError, ValueError) as exc:
|
||||
await discord_friendly_reply(
|
||||
ctx, f"Nie mogę wpiąć modelu '{model}': {exc}"
|
||||
)
|
||||
return
|
||||
|
||||
self.logger.info(
|
||||
"Przełączono AI na config %s (%s / %s)",
|
||||
nazwa_konfigu, cfg.get("provider"), cfg.get("latest_model"),
|
||||
)
|
||||
message = (
|
||||
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:
|
||||
others = await ai_functions.list_provider_models(nazwa_konfigu)
|
||||
except ai_functions.AIError:
|
||||
others = []
|
||||
current = cfg.get("latest_model")
|
||||
if others and current not in others:
|
||||
# The configured/pinned model is not on the server: every
|
||||
# reply would fail with "model not found" and nothing would
|
||||
# say why. Flag it here, where the list is already in hand.
|
||||
message += (
|
||||
f"\n⚠ Uwaga: '{current}' nie jest wgrany na serwerze. "
|
||||
f"Dostępne: {', '.join(others)} "
|
||||
f"(`$gadaj_teraz {nazwa_konfigu} <model>`)."
|
||||
)
|
||||
elif len(others) > 1:
|
||||
message += (
|
||||
f"\nDostępne modele: {', '.join(others)} "
|
||||
f"(`$gadaj_teraz {nazwa_konfigu} <model>`)."
|
||||
)
|
||||
await discord_friendly_reply(ctx, message)
|
||||
|
||||
@commands.hybrid_command(
|
||||
name="armia_hammera",
|
||||
@@ -411,10 +326,8 @@ 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:
|
||||
# superfryta = [discord_id, assistant_name, owner, instructions, legacy_assistant_id]
|
||||
await ai_functions.chat_with_personal_assistant(
|
||||
message, superfryta[2], superfryta[3]
|
||||
)
|
||||
#await self.bot.process_commands(message)
|
||||
await ai_functions.chat_with_assistant(message, superfryta[1])
|
||||
return
|
||||
elif self.armia[message.author.id] == Dm_Mode.ECHO_ECHO:
|
||||
await ai_functions.echo(message)
|
||||
@@ -514,11 +427,6 @@ class Events(commands.Cog):
|
||||
return
|
||||
message_content_lower = message_content_lower.replace("imaginuje sobie: ", "")
|
||||
self.logger.debug("Wywolanie obrazka: %s", message_content_lower)
|
||||
# Every error branch below must RETURN: otherwise control falls
|
||||
# through to `if response:` with `response` unbound (the call
|
||||
# raised) -> UnboundLocalError, crashing the handler right after
|
||||
# the friendly message was already sent.
|
||||
response = None
|
||||
try:
|
||||
response = await OPENAICLIENT.images.generate(
|
||||
model="dall-e-3",
|
||||
@@ -532,12 +440,10 @@ class Events(commands.Cog):
|
||||
await discord_friendly_reply(
|
||||
message, f"*Kondziu patrzy na terminal, czeka, czeka, czeka,.... Jeszcze chwile czeka Przypierdala w niego pięścią....* Nie mogę się połączyć z Openai spróbuj od nowa. *Na ekranie pojawia się*: {e}"
|
||||
)
|
||||
return
|
||||
except openai.APIConnectionError as e:
|
||||
await discord_friendly_reply(
|
||||
message, f"*Kondziu patrzy na terminal, chwile się zastanawia. Przypierdala w niego pięścią....* Nie mogę się połączyć z Openai. *Na ekranie pojawia się*: {e}"
|
||||
)
|
||||
return
|
||||
except openai.BadRequestError as e:
|
||||
# Handle invalid request error, e.g. validate parameters or log
|
||||
if message.author.nick:
|
||||
@@ -555,31 +461,27 @@ class Events(commands.Cog):
|
||||
await discord_friendly_reply(
|
||||
message, f"Sorki, cenzura: {resp}. Jak chcesz to są kanały na nudle #sexy-foteczky i #kanal-do-fapania *Na ekranie pojawia się: {e}"
|
||||
)
|
||||
return
|
||||
except openai.AuthenticationError as e:
|
||||
# Handle authentication error, e.g. check credentials or log
|
||||
await discord_friendly_reply(
|
||||
message, f"*Kondziu patrzy na terminal, chwile się zastanawia. Przypierdala w niego pięścią....* Wołaj szefa - coś się z hasłem zjebało. *Na terminalu pojawia się:* {e}"
|
||||
)
|
||||
return
|
||||
except openai.PermissionDeniedError as e:
|
||||
# Handle permission error, e.g. check scope or log
|
||||
# (was accidentally passing a (message, text) TUPLE as one arg)
|
||||
await discord_friendly_reply(
|
||||
(
|
||||
message, f"*Kondziu patrzy na terminal, chwile się zastanawia. Przypierdala w niego pięścią....* Wołaj szefa - coś się z uprawnieniami zjebało. *Na terminalu pojawia się:* {e}"
|
||||
)
|
||||
return
|
||||
)
|
||||
except openai.RateLimitError as e:
|
||||
await discord_friendly_reply(
|
||||
message, f"*Kondziu patrzy na terminal* Wołaj szefa. Zapłacić rachunki za AI trzeba. Jak chcesz to się na #zebranie dorzuć. {e}"
|
||||
)
|
||||
return
|
||||
except openai.APIError as e:
|
||||
# Handle API error, e.g. retry or log
|
||||
await discord_friendly_reply(
|
||||
message, f"*Kondziu nurkuje za bar, terminal wybucha. Przed tobą ląduje pergamin zapisany pięknym gotykiem a na nim*: {e}"
|
||||
)
|
||||
return
|
||||
if response:
|
||||
self.logger.info(response)
|
||||
image_url = response.data[0].url
|
||||
|
||||
+54
-232
@@ -1,21 +1,15 @@
|
||||
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,
|
||||
ASSISTANT_MEMORY_FILE,
|
||||
ASSISTANT_MEMORY_TURNS,
|
||||
ASSISTANTS,
|
||||
CLAUDECLIENT,
|
||||
CYCLIC_WORDS,
|
||||
DEFAULT_AI_CONFIG,
|
||||
@@ -25,10 +19,6 @@ from constants import (
|
||||
MEMORY_FIVE_SIARA,
|
||||
MESSAGE_TABLE,
|
||||
MESSAGE_TABLE_MUZYKA,
|
||||
OLLAMACLIENT,
|
||||
OLLAMA_KEEP_ALIVE,
|
||||
OLLAMA_PRELOAD_TIMEOUT,
|
||||
OLLAMA_URL,
|
||||
OPENAICLIENT,
|
||||
SYSTEM_GPT_SETTINGS,
|
||||
WORD_REACTIONS,
|
||||
@@ -102,54 +92,12 @@ def set_active_ai_config(name: str) -> dict:
|
||||
raise RuntimeError("klient Anthropic nie jest skonfigurowany (brak ANTHROPIC_API_KEY)")
|
||||
if provider == "openai" and OPENAICLIENT is None:
|
||||
raise RuntimeError("klient OpenAI nie jest skonfigurowany (brak OPENAI_API_KEY)")
|
||||
if provider == "ollama" and OLLAMACLIENT is None:
|
||||
raise RuntimeError("Ollama nie jest skonfigurowana (ustaw CONJURER_OLLAMA_URL)")
|
||||
_ACTIVE_CONFIG_NAME = name
|
||||
_persist_active_ai_config(name)
|
||||
return cfg
|
||||
|
||||
|
||||
async def list_provider_models(name: str = None):
|
||||
"""Model ids selectable for a config.
|
||||
|
||||
For Ollama this ASKS THE SERVER (its OpenAI-compatible /v1/models), so the
|
||||
picker always reflects what is actually pulled on the box rather than a
|
||||
hardcoded list. Hosted providers are not enumerated - we only report what
|
||||
the config is wired to.
|
||||
"""
|
||||
cfg = AI_CONFIGS.get(name or _ACTIVE_CONFIG_NAME) or _active_config()
|
||||
if cfg.get("provider") == "ollama":
|
||||
if OLLAMACLIENT is None:
|
||||
raise AIError(
|
||||
"auth",
|
||||
RuntimeError("Ollama nie jest skonfigurowana (ustaw CONJURER_OLLAMA_URL)"),
|
||||
)
|
||||
try:
|
||||
resp = await OLLAMACLIENT.models.list()
|
||||
except Exception as exc: # pylint: disable=broad-except
|
||||
raise _map_openai_error(exc)
|
||||
return sorted({item.id for item in resp.data})
|
||||
return [m for m in (cfg.get("latest_model"), cfg.get("cheap_model")) if m]
|
||||
|
||||
|
||||
def set_active_model(model: str, name: str = None) -> dict:
|
||||
"""Pin the model a config uses for normal replies, and persist it.
|
||||
|
||||
Only ``latest_model`` is changed; ``cheap_model`` stays as configured so the
|
||||
MUSIC path keeps its cheaper backend.
|
||||
"""
|
||||
cfg_name = name or _ACTIVE_CONFIG_NAME
|
||||
if cfg_name not in AI_CONFIGS:
|
||||
raise KeyError(cfg_name)
|
||||
if not model or not model.strip():
|
||||
raise ValueError("pusta nazwa modelu")
|
||||
cfg = AI_CONFIGS[cfg_name]
|
||||
cfg["latest_model"] = model.strip()
|
||||
_persist_active_ai_config(_ACTIVE_CONFIG_NAME, model_for=cfg_name)
|
||||
return cfg
|
||||
|
||||
|
||||
def _persist_active_ai_config(name: str, model_for: str = None) -> None:
|
||||
def _persist_active_ai_config(name: str) -> None:
|
||||
"""Best-effort write of the active-config choice into system_gpt_settings.json.
|
||||
|
||||
Keeps the historical two-element structure intact: updates index 2 if it
|
||||
@@ -169,16 +117,7 @@ def _persist_active_ai_config(name: str, model_for: str = None) -> None:
|
||||
return
|
||||
if len(data) > 2 and isinstance(data[2], dict):
|
||||
data[2]["active"] = name
|
||||
# Write back ONLY what this process actually changed. Assigning the whole
|
||||
# in-memory AI_CONFIGS here would clobber operator hand-edits (the only
|
||||
# way to change cheap_model/temperature/max_tokens) and re-seed configs
|
||||
# deliberately deleted from the file, because AI_CONFIGS is the built-in
|
||||
# defaults merged under the file. setdefault alone is not enough either:
|
||||
# it would drop a model pinned via set_active_model, hence model_for.
|
||||
block = data[2].setdefault("configs", AI_CONFIGS)
|
||||
if model_for and model_for in AI_CONFIGS:
|
||||
entry = block.setdefault(model_for, dict(AI_CONFIGS[model_for]))
|
||||
entry["latest_model"] = AI_CONFIGS[model_for]["latest_model"]
|
||||
data[2].setdefault("configs", AI_CONFIGS)
|
||||
else:
|
||||
data = data[:2] + [{"active": name, "configs": AI_CONFIGS}]
|
||||
try:
|
||||
@@ -275,83 +214,12 @@ async def _anthropic_call(messages, model, cfg):
|
||||
return text.strip()
|
||||
|
||||
|
||||
async def _ollama_call(messages, model, cfg):
|
||||
"""Self-hosted counterpart of openai_call. Returns a plain string.
|
||||
|
||||
Ollama exposes an OpenAI-compatible /v1 surface, so the same message format
|
||||
and the same error mapping apply - only the base_url and the model ids
|
||||
differ. Chat Completions (not the Responses API) is what Ollama implements.
|
||||
"""
|
||||
if OLLAMACLIENT is None:
|
||||
raise AIError(
|
||||
"auth",
|
||||
RuntimeError("Ollama nie jest skonfigurowana (ustaw CONJURER_OLLAMA_URL)"),
|
||||
)
|
||||
try:
|
||||
resp = await OLLAMACLIENT.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=float(cfg.get("temperature", 0.2)),
|
||||
)
|
||||
except Exception as exc: # pylint: disable=broad-except
|
||||
raise _map_openai_error(exc)
|
||||
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()
|
||||
try:
|
||||
if cfg.get("provider") == "anthropic":
|
||||
return await _anthropic_call(messages, model, cfg)
|
||||
if cfg.get("provider") == "ollama":
|
||||
return await _ollama_call(messages, model, cfg)
|
||||
return await openai_call(messages, model, temperature)
|
||||
except AIError:
|
||||
raise
|
||||
@@ -591,7 +459,7 @@ async def handle_response(
|
||||
try:
|
||||
# ...przygotowanie messages/system prompt/itp. jak masz...
|
||||
# retry/backoff + deadline (zachowuje Twoją semantykę logowania)
|
||||
timeout_sec = AI_TIMEOUT_SECONDS
|
||||
timeout_sec = 120
|
||||
deadline = time.time() + timeout_sec
|
||||
response = await asyncio.wait_for(
|
||||
provider_generate(messages=history_msgs, model=model_to_use),
|
||||
@@ -673,12 +541,10 @@ async def get_random_cyclic_message(client):
|
||||
# trunk-ignore(bandit/B311)
|
||||
ai_check = random.randint(0, 10)
|
||||
logger.info("Losowa wypowiedź")
|
||||
if ai_check < 2 and CYCLIC_WORDS:
|
||||
if ai_check < 2:
|
||||
logger.info("Predefiniowana")
|
||||
# randrange(n) is 0..n-1; randint(0, n) was inclusive and could return n
|
||||
# -> list(...)[n] IndexError. Guarded on empty CYCLIC_WORDS above.
|
||||
# trunk-ignore(bandit/B311)
|
||||
messnum = random.randrange(len(CYCLIC_WORDS))
|
||||
messnum = random.randint(0, len(CYCLIC_WORDS))
|
||||
logger.debug(messnum)
|
||||
logger.debug(len(CYCLIC_WORDS))
|
||||
mess_key = list(CYCLIC_WORDS.keys())[messnum]
|
||||
@@ -718,104 +584,60 @@ async def get_random_cyclic_message(client):
|
||||
return result
|
||||
|
||||
|
||||
# ----------------------------------------------------------------- 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
|
||||
async def create_chat_assistant(owner_id, id, name, owner, special_instructions):
|
||||
logger = logging.getLogger("discord")
|
||||
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
|
||||
|
||||
|
||||
def _save_assistant_memory() -> None:
|
||||
"""Atomic write: a torn file would lose someone's whole conversation."""
|
||||
logger = logging.getLogger("discord")
|
||||
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}]
|
||||
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)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
async def chat_with_assistant(message, assistant_name):
|
||||
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,
|
||||
assistant_data = ASSISTANTS[assistant_name]
|
||||
ai_message = await OPENAICLIENT.beta.threads.messages.create(
|
||||
thread_id=assistant_data[3].id, role="user", content=message.content
|
||||
)
|
||||
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
|
||||
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",
|
||||
)
|
||||
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)
|
||||
|
||||
|
||||
async def echo(message):
|
||||
|
||||
@@ -39,7 +39,6 @@ from constants import (
|
||||
LIBRARIAN_SERVICE_ADDRESS,
|
||||
LOGFILE,
|
||||
RADIO_SERVICE_ADDRESS,
|
||||
SELF_CALLBACK,
|
||||
TOKEN,
|
||||
service_headers,
|
||||
)
|
||||
@@ -129,15 +128,13 @@ SERVICE_EXTENSION_GROUPS = {
|
||||
SERVICE_RECHECK_SECONDS = 300
|
||||
|
||||
|
||||
def _service_health(url: str):
|
||||
"""Return None when the service answers HTTP at all (any status counts),
|
||||
otherwise the connection error explaining WHY it's unreachable (refused vs
|
||||
timeout vs DNS - the difference points straight at the cause)."""
|
||||
def _service_alive(url: str) -> bool:
|
||||
"""True when the service answers HTTP at all (any status code counts)."""
|
||||
try:
|
||||
requests.get(url, timeout=3)
|
||||
return None
|
||||
except requests.exceptions.RequestException as exc:
|
||||
return exc
|
||||
return True
|
||||
except requests.exceptions.RequestException:
|
||||
return False
|
||||
|
||||
|
||||
async def _load_extension_safe(name: str) -> bool:
|
||||
@@ -174,24 +171,14 @@ async def _load_service_groups() -> bool:
|
||||
LIBRARIAN_PING,
|
||||
service_headers(),
|
||||
LIBRARIAN_PING_TIMEOUT,
|
||||
SELF_CALLBACK,
|
||||
)
|
||||
else:
|
||||
alive = await asyncio.to_thread(_service_alive, group["health_url"])
|
||||
if not alive:
|
||||
logger.warning(
|
||||
"Service 'librarian' ping round-trip failed (%s) - cogs stay disabled: %s",
|
||||
group["health_url"],
|
||||
", ".join(missing),
|
||||
)
|
||||
continue
|
||||
else:
|
||||
err = await asyncio.to_thread(_service_health, group["health_url"])
|
||||
if err is not None:
|
||||
logger.warning(
|
||||
"Service '%s' unreachable (%s) [%s: %s] - cogs stay disabled: %s",
|
||||
"Service '%s' unreachable (%s) - cogs stay disabled: %s",
|
||||
service,
|
||||
group["health_url"],
|
||||
type(err).__name__,
|
||||
err,
|
||||
", ".join(missing),
|
||||
)
|
||||
continue
|
||||
@@ -238,16 +225,6 @@ async def on_ready():
|
||||
for extension in CORE_EXTENSIONS:
|
||||
await _load_extension_safe(extension)
|
||||
|
||||
# Log the ACTUALLY-resolved service addresses. When one shows the built-in
|
||||
# default (192.168.1.15:5000) it means the matching CONJURER_* env var never
|
||||
# reached the process - the single most common cause of "service unreachable"
|
||||
# confusion. Printing them makes env-vs-default obvious at a glance.
|
||||
logger.info(
|
||||
"Resolved service addresses -> musician(file): %s | librarian: %s | radio: %s",
|
||||
FILE_SERVICE_ADDRESS,
|
||||
LIBRARIAN_SERVICE_ADDRESS,
|
||||
RADIO_SERVICE_ADDRESS,
|
||||
)
|
||||
await _load_service_groups()
|
||||
logger.info("Sensors: online")
|
||||
|
||||
|
||||
+10
-86
@@ -13,9 +13,6 @@ import requests
|
||||
from flask import Flask, abort, jsonify, request
|
||||
from waitress import serve
|
||||
|
||||
from constants import DELIVERED_DIR, DELIVERED_MAX, RESULT_INBOX_DIR
|
||||
from durable_queue import DiskQueue
|
||||
|
||||
HOST_ADDRESS = os.getenv("CONJURER_DISCORD_HOST", "192.168.1.92")
|
||||
PORT_ADDRESS = int(os.getenv("CONJURER_DISCORD_PORT", "5000"))
|
||||
ICECAST_ADDRESS = os.getenv("CONJURER_ICECAST", "http://192.168.1.12:8000")
|
||||
@@ -36,43 +33,9 @@ incoming_q = Queue()
|
||||
# records older than this. Kept well above the ping timeout so a slow-but-alive
|
||||
# round-trip is never swept out from under a waiter.
|
||||
PING_TTL_SECONDS = 30
|
||||
|
||||
# Durable result delivery. Every incoming search result is persisted to _inbox
|
||||
# before we ack the librarian, and only removed once it has actually been
|
||||
# rendered to the user (its uuid recorded in _delivered). This makes /conjurer
|
||||
# idempotent (the librarian can safely resend until acked) and lets an accepted
|
||||
# result survive a bot restart mid-flight (replayed from _inbox on startup).
|
||||
_inbox = DiskQueue(RESULT_INBOX_DIR)
|
||||
_delivered = DiskQueue(DELIVERED_DIR)
|
||||
app = Flask(__name__)
|
||||
|
||||
|
||||
def mark_delivered(query_uuid) -> None:
|
||||
"""Record that a result was rendered to the user.
|
||||
|
||||
After this, the librarian's resends of that uuid are dropped as duplicates
|
||||
and it is never replayed from the inbox again. Called by the cog once it has
|
||||
actually posted the result to Discord."""
|
||||
_delivered.put(query_uuid, {})
|
||||
_delivered.prune(DELIVERED_MAX)
|
||||
_inbox.remove(query_uuid)
|
||||
|
||||
|
||||
def replay_inbox() -> None:
|
||||
"""Re-queue INBOX results that were accepted but not yet rendered.
|
||||
|
||||
Recovers an expensive result that reached the bot (and was acked to the
|
||||
librarian, so it won't be resent) but whose render was lost to a bot restart.
|
||||
"""
|
||||
logger = logging.getLogger("discord")
|
||||
for query_uuid, payload, _ts in _inbox.items():
|
||||
if _delivered.contains(query_uuid):
|
||||
_inbox.remove(query_uuid)
|
||||
continue
|
||||
logger.info("Replaying un-rendered result %s from INBOX", query_uuid)
|
||||
incoming_q.put(payload)
|
||||
|
||||
|
||||
def _authorize_request() -> None:
|
||||
"""Reject inbound calls lacking the shared key (no-op if key is unset)."""
|
||||
if API_KEY and request.headers.get("X-Conjurer-Api-Key") != API_KEY:
|
||||
@@ -150,31 +113,12 @@ def answer_external_command():
|
||||
"""
|
||||
_authorize_request()
|
||||
logger = logging.getLogger("discord")
|
||||
logger.info(request)
|
||||
record = json.loads(request.data)
|
||||
logger.info("DATA RECEIVED: %s", record)
|
||||
# Health-check pongs are ephemeral - never persisted or deduped.
|
||||
if isinstance(record, dict) and "__pong__" in record:
|
||||
logger.info(record)
|
||||
logger.info("DATA RECEIVED")
|
||||
incoming_q.put(record)
|
||||
return jsonify("SUCCESS")
|
||||
# Search results: idempotent, durable intake. Persist each uuid to the inbox
|
||||
# before acking, and queue only a uuid we have NOT already delivered or
|
||||
# accepted. This lets the librarian's resender retry safely (a duplicate is
|
||||
# dropped, never double-rendered) and lets an accepted-but-unrendered result
|
||||
# be replayed after a bot restart.
|
||||
if isinstance(record, dict):
|
||||
for query_uuid in list(record.keys()):
|
||||
if _delivered.contains(query_uuid):
|
||||
logger.info("Result %s already delivered - dropping duplicate", query_uuid)
|
||||
continue
|
||||
if _inbox.contains(query_uuid):
|
||||
logger.info("Result %s already pending - dropping duplicate", query_uuid)
|
||||
continue
|
||||
single = {query_uuid: record[query_uuid]}
|
||||
_inbox.put(query_uuid, single)
|
||||
incoming_q.put(single)
|
||||
else:
|
||||
incoming_q.put(record) # unexpected shape - preserve old behaviour
|
||||
return jsonify("SUCCESS")
|
||||
|
||||
|
||||
def submit_ai_query(prompt, channel_id=None, request_type="NONE", username="conjurer", query_uuid=None):
|
||||
@@ -325,20 +269,14 @@ def scan_incoming(stop_event: Optional[threading.Event] = None):
|
||||
awaiting_q.remove(record)
|
||||
logger.info("PONG matched for %s", pong_uuid)
|
||||
continue
|
||||
# Collect matched records and drop them from awaiting_q afterwards -
|
||||
# they used to stay forever (awaiting_q only ever grew), leaking
|
||||
# memory over the bot's uptime and letting a reused UUID re-match a
|
||||
# stale record.
|
||||
matched = []
|
||||
record_stored = False
|
||||
for record in awaiting_q:
|
||||
if record.uuid in answer.keys():
|
||||
record_stored = True
|
||||
record.stop = True
|
||||
record.entries = answer[record.uuid]
|
||||
IN_COMM_Q.put(record)
|
||||
matched.append(record)
|
||||
for record in matched:
|
||||
awaiting_q.remove(record)
|
||||
if not matched:
|
||||
if not record_stored:
|
||||
for key in answer.keys():
|
||||
record = QueryControl("Orphaned", key, "Orphan", None)
|
||||
record.stop = True
|
||||
@@ -349,13 +287,9 @@ def scan_incoming(stop_event: Optional[threading.Event] = None):
|
||||
|
||||
|
||||
def librarian_ping(address: str, endpoint: str, headers: Optional[dict] = None,
|
||||
timeout: float = 3.0, callback: str = "") -> bool:
|
||||
timeout: float = 3.0) -> bool:
|
||||
"""Health-check the librarian by round-tripping a ping through the FULL path.
|
||||
|
||||
``callback`` is where THIS bot wants the pong sent back to, so a shared
|
||||
librarian pongs each bot at its own address (else a second bot's ping would
|
||||
be ponged to the first and its health check would always time out).
|
||||
|
||||
The ping is a pseudo-query that exercises exactly the same machinery a real
|
||||
search does, on BOTH sides:
|
||||
|
||||
@@ -384,7 +318,7 @@ def librarian_ping(address: str, endpoint: str, headers: Optional[dict] = None,
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{address}{endpoint}",
|
||||
json={"UUID": ping_uuid, "callback": callback},
|
||||
json={"UUID": ping_uuid},
|
||||
headers=headers or {},
|
||||
timeout=timeout,
|
||||
)
|
||||
@@ -419,13 +353,10 @@ def id3(url: str) -> dict:
|
||||
resp.read(
|
||||
metaint
|
||||
) # this isn't seekable so, arbitrarily read to the point we want
|
||||
# Guard the headers: an Icecast stream that omits icy-name / icy-genre
|
||||
# (e.g. while the radio is down) made `.title()` raise AttributeError on
|
||||
# None, 500-ing the /prepped_tracks "next" handler that calls this.
|
||||
tagdata = dict(
|
||||
site_url=resp.headers.get("icy-url"),
|
||||
name=(resp.headers.get("icy-name") or "").title(),
|
||||
genre=(resp.headers.get("icy-genre") or "").title(),
|
||||
name=resp.headers.get("icy-name").title(),
|
||||
genre=resp.headers.get("icy-genre").title(),
|
||||
title=get_stream_title(resp.read(255)),
|
||||
)
|
||||
return tagdata
|
||||
@@ -465,13 +396,6 @@ def comm_subroutine(stop_event: Optional[threading.Event] = None):
|
||||
for worker in threads:
|
||||
worker.start()
|
||||
|
||||
# Recover any result that was accepted before a previous shutdown but never
|
||||
# rendered - re-queue it now that scan_incoming is running.
|
||||
try:
|
||||
replay_inbox()
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
logger.exception("INBOX replay on startup failed")
|
||||
|
||||
try:
|
||||
while any(thread.is_alive() for thread in threads):
|
||||
if stop_event and stop_event.is_set():
|
||||
|
||||
@@ -203,12 +203,6 @@ def wyszukaj(word_list, how_many, _logger=None, write_to=None):
|
||||
# ---------------------------------------------------------------- tailer
|
||||
def scan_tracks():
|
||||
"""Tail the radio logs and forward play events to the bot."""
|
||||
# On a fresh deploy Liquidsoap may not have written its logs yet; wait for
|
||||
# them instead of dying with FileNotFoundError, which used to silently kill
|
||||
# the now-playing forwarder until the container was restarted.
|
||||
while not (RADIOLOG_PATH.exists() and PERSISTENCE_PATH.exists()):
|
||||
logger.info("Waiting for radio logs (%s, %s)...", RADIOLOG_PATH, PERSISTENCE_PATH)
|
||||
time.sleep(5)
|
||||
with open(RADIOLOG_PATH, "r", encoding=ENCODING) as log_file:
|
||||
log_file.seek(os.stat(RADIOLOG_PATH).st_size)
|
||||
prev_size = os.stat(PERSISTENCE_PATH).st_size
|
||||
|
||||
@@ -14,17 +14,14 @@ Functions:
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
from json.decoder import JSONDecodeError
|
||||
from logging import handlers
|
||||
from pathlib import Path
|
||||
from queue import Empty, Queue
|
||||
from queue import Queue
|
||||
from typing import Dict, Optional
|
||||
|
||||
import requests
|
||||
@@ -32,7 +29,6 @@ import lib_paths
|
||||
import scrape_bot
|
||||
import search_bot
|
||||
# import search_bot2 as search_bot
|
||||
from durable_queue import DiskQueue
|
||||
from flask import Flask, jsonify, request, abort
|
||||
from habanero import Crossref
|
||||
from waitress import serve
|
||||
@@ -68,207 +64,11 @@ LOGFILE_PATH = _env_path(
|
||||
"CONJURER_LIBRARIAN_LOG", str(BASE_DIR / "librarian.log")
|
||||
)
|
||||
|
||||
# Durable OUTBOX for finished results. A result is expensive (hours of compute),
|
||||
# so it is written here and only removed once the bot ACKs it (HTTP 200). Lives
|
||||
# on the librarian's persistent state volume, so it survives a librarian restart
|
||||
# and a transient bot outage; the resender thread keeps retrying until delivered.
|
||||
OUTBOX_DIR = _env("CONJURER_LIBRARIAN_OUTBOX", os.path.join(lib_paths.STATE_DIR, "outbox"))
|
||||
RESULT_SEND_ATTEMPTS = int(_env("CONJURER_RESULT_SEND_ATTEMPTS", "3"))
|
||||
RESULT_SEND_BACKOFF = float(_env("CONJURER_RESULT_SEND_BACKOFF", "2"))
|
||||
OUTBOX_RESEND_SECONDS = int(_env("CONJURER_OUTBOX_RESEND_SECONDS", "60"))
|
||||
_outbox = DiskQueue(OUTBOX_DIR)
|
||||
|
||||
# cr_results/rr_results/s_results.json are write-only debug dumps (nothing reads
|
||||
# them). They used to accumulate EVERY search forever AND json.load the whole
|
||||
# growing file on each write - unbounded RAM + disk, and for a deep search the
|
||||
# raw dump is hundreds of MB. Off by default now; when explicitly enabled they
|
||||
# are overwritten with just the latest search (no load, no accumulation).
|
||||
DEBUG_DUMPS = _env("CONJURER_LIBRARIAN_DEBUG_DUMPS", "0").lower() in ("1", "true", "yes")
|
||||
# Log level: INFO keeps normal runs readable (the desktop-era per-line/per-file
|
||||
# chatter is now DEBUG); set DEBUG to get the full verbosity back.
|
||||
LOG_LEVEL = _env("CONJURER_LIBRARIAN_LOG_LEVEL", "INFO").upper()
|
||||
|
||||
|
||||
def _dump_debug(path, uuid, data) -> None:
|
||||
"""Optionally dump the latest search's data for debugging.
|
||||
|
||||
Overwrites (never accumulates) and does nothing unless DEBUG_DUMPS is on, so
|
||||
it can't grow RAM or the state volume in normal operation."""
|
||||
if not DEBUG_DUMPS:
|
||||
return
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
json.dump({uuid: data}, handle)
|
||||
except OSError as exc:
|
||||
logging.getLogger("conjurer_librarian").warning("Debug dump to %s failed: %s", path, exc)
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
librarian_queue = Queue()
|
||||
librarian_list = []
|
||||
|
||||
# Lifecycle of every real search uuid: "queued" (accepted, sitting in
|
||||
# librarian_queue) -> "processing" (worker pulled it) -> removed (worker
|
||||
# finished AND attempted to send the result). The bot's per-query watchdog polls
|
||||
# /query_status against this: a uuid that VANISHES from here without its result
|
||||
# reaching the bot is a lost result (finished-but-never-delivered) and gets
|
||||
# flagged in chat.
|
||||
active_queries: Dict[str, str] = {}
|
||||
_active_lock = threading.Lock()
|
||||
# Set while the worker is grinding a real search. A ping arriving during this
|
||||
# pongs back immediately WITHOUT queueing - being busy is healthy (you can keep
|
||||
# piling searches on), so "busy" must never look like "dead" to the health check.
|
||||
worker_busy = threading.Event()
|
||||
|
||||
# ---- Graceful shutdown + resumable search state ----------------------------
|
||||
# SHUTDOWN_EVENT is set by the SIGTERM/SIGINT handler; the running search checks
|
||||
# it (via search_bot) and checkpoints itself. SHUTDOWN_DONE is set by the worker
|
||||
# once it has stopped cleanly, so the main thread can exit promptly - bounded by
|
||||
# GRACEFUL_TIMEOUT so we never become an un-killable zombie pod.
|
||||
SHUTDOWN_EVENT = threading.Event()
|
||||
SHUTDOWN_DONE = threading.Event()
|
||||
GRACEFUL_TIMEOUT = float(_env("CONJURER_LIBRARIAN_GRACEFUL_TIMEOUT", "45"))
|
||||
# Persisted, per-uuid: accepted-but-unfinished search REQUESTS (so a restart
|
||||
# re-runs them) and in-progress CHECKPOINTS (found-so-far + per-file resume
|
||||
# offset, so a restart CONTINUES a long scan instead of restarting it).
|
||||
REQUESTS_DIR = _env("CONJURER_LIBRARIAN_REQUESTS", os.path.join(lib_paths.STATE_DIR, "requests"))
|
||||
CHECKPOINT_DIR = _env("CONJURER_LIBRARIAN_CHECKPOINTS", os.path.join(lib_paths.STATE_DIR, "checkpoints"))
|
||||
_requests = DiskQueue(REQUESTS_DIR)
|
||||
_checkpoints = DiskQueue(CHECKPOINT_DIR)
|
||||
|
||||
# Simple result cache: a repeat of the same query (normalised) skips the whole
|
||||
# Crossref + DB-scan and returns the stored hits. Disk-backed so it survives a
|
||||
# restart, TTL'd, and size-bounded. Set CONJURER_LIBRARIAN_CACHE_TTL=0 to disable.
|
||||
CACHE_DIR = _env("CONJURER_LIBRARIAN_CACHE", os.path.join(lib_paths.STATE_DIR, "cache"))
|
||||
CACHE_TTL_SECONDS = int(_env("CONJURER_LIBRARIAN_CACHE_TTL", str(7 * 24 * 3600)))
|
||||
CACHE_MAX_ENTRIES = int(_env("CONJURER_LIBRARIAN_CACHE_MAX", "500"))
|
||||
_cache = DiskQueue(CACHE_DIR)
|
||||
|
||||
|
||||
def _cache_key(query, deep_search) -> str:
|
||||
"""Stable key for a query: whitespace-normalised, case-insensitive, and
|
||||
scoped by deep vs shallow (they return different result sets)."""
|
||||
normalised = " ".join(str(query).lower().split())
|
||||
return hashlib.sha256(f"{int(bool(deep_search))}:{normalised}".encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _cache_get(query, deep_search):
|
||||
"""Return the cached final_result for this query, or None on miss/expiry."""
|
||||
if CACHE_TTL_SECONDS <= 0:
|
||||
return None
|
||||
entry = _cache.get(_cache_key(query, deep_search))
|
||||
if not entry or entry.get("expires", 0) < time.time():
|
||||
return None
|
||||
return entry.get("final_result")
|
||||
|
||||
|
||||
def _cache_put(query, deep_search, final_result) -> None:
|
||||
"""Store a completed search's hits, with a TTL, and bound the cache size."""
|
||||
if CACHE_TTL_SECONDS <= 0:
|
||||
return
|
||||
_cache.put(
|
||||
_cache_key(query, deep_search),
|
||||
{"query": str(query), "final_result": final_result, "expires": time.time() + CACHE_TTL_SECONDS},
|
||||
)
|
||||
_cache.prune(CACHE_MAX_ENTRIES)
|
||||
|
||||
|
||||
# ---- "Still alive" heartbeat for the running search ------------------------
|
||||
# A deep scan runs for hours with nothing in the log between start and finish.
|
||||
# Every HEARTBEAT_SECONDS the running search says it is still going, with its
|
||||
# uuid, the phrase, and a ROUGH how-far-along. The estimate is deliberately
|
||||
# cheap: producers already record a byte offset per chunk file, and the total
|
||||
# size is stat()'d once at search start - so it costs a sum over ~40 ints.
|
||||
HEARTBEAT_SECONDS = int(_env("CONJURER_LIBRARIAN_HEARTBEAT_SECONDS", "1200")) # 20 min
|
||||
_current_search: Dict[str, object] = {}
|
||||
_current_lock = threading.Lock()
|
||||
|
||||
|
||||
def _set_current_search(uuid, query, progress, live_results) -> None:
|
||||
with _current_lock:
|
||||
_current_search.clear()
|
||||
_current_search.update({
|
||||
"uuid": str(uuid), "query": str(query), "started": time.monotonic(),
|
||||
"progress": progress, "live": live_results,
|
||||
})
|
||||
|
||||
|
||||
def _clear_current_search() -> None:
|
||||
with _current_lock:
|
||||
_current_search.clear()
|
||||
|
||||
|
||||
def _progress_summary(progress):
|
||||
"""(done_bytes, total_bytes, percent) from a live progress dict. Cheap: a
|
||||
sum over one int per chunk file. Percent is 0.0 when the total is unknown."""
|
||||
progress = progress or {}
|
||||
positions = progress.get("positions") or {}
|
||||
total = progress.get("total_bytes") or 0
|
||||
done = sum(positions.values())
|
||||
if total > 0:
|
||||
done = min(done, total) # a partially-buffered tail can nudge past 100%
|
||||
return done, total, 100.0 * done / total
|
||||
return done, total, 0.0
|
||||
|
||||
|
||||
def search_heartbeat(app_logger) -> None:
|
||||
"""Log a 'still searching' line every HEARTBEAT_SECONDS while one runs."""
|
||||
while not SHUTDOWN_EVENT.wait(HEARTBEAT_SECONDS):
|
||||
try:
|
||||
with _current_lock:
|
||||
snapshot = dict(_current_search) if _current_search else None
|
||||
if not snapshot:
|
||||
continue # nothing running - stay quiet
|
||||
done, total, percent = _progress_summary(snapshot.get("progress"))
|
||||
app_logger.info(
|
||||
"SEARCH ALIVE %s | '%s' | ~%.1f%% przeskanowane (%.2f/%.2f GB, "
|
||||
"~%.2f GB do końca) | %d trafień | %.0f min",
|
||||
snapshot["uuid"], snapshot["query"], percent,
|
||||
done / 1e9, total / 1e9, max(0, total - done) / 1e9,
|
||||
len(snapshot.get("live") or []),
|
||||
(time.monotonic() - snapshot["started"]) / 60.0,
|
||||
)
|
||||
except Exception as exc: # pylint: disable=broad-exception-caught
|
||||
app_logger.warning("Heartbeat tick failed: %s", exc)
|
||||
|
||||
|
||||
# Crossref is a public service that times out / rate-limits under load. A single
|
||||
# transient ReadTimeout used to blow up the whole (expensive) search, so every
|
||||
# habanero call is retried with backoff, and a search that still fails is retried
|
||||
# as a whole a few times before being given up on.
|
||||
CROSSREF_ATTEMPTS = int(_env("CONJURER_CROSSREF_ATTEMPTS", "4"))
|
||||
CROSSREF_BACKOFF = float(_env("CONJURER_CROSSREF_BACKOFF", "5"))
|
||||
SEARCH_MAX_ATTEMPTS = int(_env("CONJURER_SEARCH_MAX_ATTEMPTS", "3"))
|
||||
|
||||
|
||||
def _crossref_call(app_logger, what, func, *args, **kwargs):
|
||||
"""Run one habanero call, retrying transient failures with linear backoff.
|
||||
|
||||
habanero wraps httpx errors (ReadTimeout, connection resets, 5xx) in a plain
|
||||
RuntimeError, so we cannot filter narrowly - we retry a BOUNDED number of
|
||||
times on any failure and re-raise the last error if none succeed. Blocking
|
||||
on purpose: callers invoke it via asyncio.to_thread, which also keeps the
|
||||
worker's event loop free while Crossref is slow."""
|
||||
last_exc = None
|
||||
for attempt in range(1, max(1, CROSSREF_ATTEMPTS) + 1):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except Exception as exc: # pylint: disable=broad-exception-caught
|
||||
last_exc = exc
|
||||
app_logger.warning(
|
||||
"Crossref %s failed (attempt %d/%d): %s", what, attempt, CROSSREF_ATTEMPTS, exc
|
||||
)
|
||||
if attempt < CROSSREF_ATTEMPTS:
|
||||
time.sleep(CROSSREF_BACKOFF * attempt)
|
||||
raise last_exc
|
||||
|
||||
|
||||
def _forget_search(uuid) -> None:
|
||||
"""A search is fully done (or abandoned): drop its persisted request and any
|
||||
checkpoint so it is never replayed or resumed again."""
|
||||
_requests.remove(uuid)
|
||||
_checkpoints.remove(uuid)
|
||||
|
||||
|
||||
def _service_headers() -> Dict[str, str]:
|
||||
if API_KEY:
|
||||
@@ -281,130 +81,13 @@ def _authorize_request() -> None:
|
||||
abort(401)
|
||||
|
||||
|
||||
def _post_pong(app_logger, ping_uuid, callback="") -> None:
|
||||
"""POST a pong for ``ping_uuid`` back to the pinging bot. Non-fatal.
|
||||
|
||||
``callback`` is the address of the bot that sent the ping, so a librarian
|
||||
shared by several bots pongs each at its OWN address (empty => the static
|
||||
MAIN_BOT_ADDRESS). Same return path a real result takes (bot's /conjurer),
|
||||
so a delivered pong proves the librarian->that-bot leg works."""
|
||||
target = f"{callback or MAIN_BOT_ADDRESS}{SEND_RESULTS}"
|
||||
try:
|
||||
requests.post(
|
||||
target, json={"__pong__": ping_uuid}, headers=_service_headers(), timeout=5
|
||||
)
|
||||
except requests.exceptions.RequestException as exc:
|
||||
app_logger.warning("PING pong send failed for %s: %s", ping_uuid, exc)
|
||||
|
||||
|
||||
def _outbox_target_payload(entry):
|
||||
"""Unpack an OUTBOX entry into (target, payload).
|
||||
|
||||
New shape: {"target": <bot address>, "payload": {uuid: result}}. Old shape
|
||||
(from before per-origin callbacks) is the raw payload - delivered to the
|
||||
default bot - so an upgrade doesn't strand results already on disk."""
|
||||
if isinstance(entry, dict) and "target" in entry and "payload" in entry:
|
||||
return entry["target"], entry["payload"]
|
||||
return "", entry
|
||||
|
||||
|
||||
def _deliver_result(target, uuid, payload, app_logger, attempts=RESULT_SEND_ATTEMPTS) -> bool:
|
||||
"""POST one result to ``target``'s /conjurer, retrying with backoff. True only
|
||||
on HTTP 200. ``target`` is the origin bot's address so a shared librarian
|
||||
answers each bot where its query came from (empty => MAIN_BOT_ADDRESS). The
|
||||
bot's /conjurer is idempotent, so re-POSTing what it already has is safe."""
|
||||
url = f"{target or MAIN_BOT_ADDRESS}{SEND_RESULTS}"
|
||||
for attempt in range(1, max(1, attempts) + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
url, json=payload, headers=_service_headers(), timeout=60
|
||||
)
|
||||
if response.status_code == 200:
|
||||
app_logger.info("Result %s delivered to %s (HTTP 200) attempt %d", uuid, url, attempt)
|
||||
return True
|
||||
app_logger.warning(
|
||||
"Result %s: %s returned HTTP %s (attempt %d/%d): %s",
|
||||
uuid, url, response.status_code, attempt, attempts, response.text[:300],
|
||||
)
|
||||
except requests.exceptions.RequestException as exc:
|
||||
app_logger.warning(
|
||||
"Result %s delivery to %s failed (attempt %d/%d): %s",
|
||||
uuid, url, attempt, attempts, exc,
|
||||
)
|
||||
if attempt < attempts:
|
||||
time.sleep(RESULT_SEND_BACKOFF * attempt)
|
||||
return False
|
||||
|
||||
|
||||
def _resend_once(app_logger) -> None:
|
||||
"""One sweep of the OUTBOX: try to deliver every un-acked result, once each,
|
||||
to the ORIGIN bot recorded with it. Removes each entry only after a positive
|
||||
ACK. Corrupt/unreadable entries are skipped by DiskQueue.items()."""
|
||||
for uuid, entry, _ts in _outbox.items():
|
||||
target, payload = _outbox_target_payload(entry)
|
||||
if _deliver_result(target, uuid, payload, app_logger, attempts=1):
|
||||
_outbox.remove(uuid)
|
||||
|
||||
|
||||
def outbox_resender(app_logger) -> None:
|
||||
"""Background loop: periodically flush the OUTBOX until the bot is reachable.
|
||||
|
||||
This is what makes an expensive result survive a transient bot outage or a
|
||||
librarian restart - on restart the persisted OUTBOX is simply resent."""
|
||||
pending = len(_outbox)
|
||||
if pending:
|
||||
app_logger.info("OUTBOX has %d un-acked result(s) on startup - will resend", pending)
|
||||
while True:
|
||||
try:
|
||||
_resend_once(app_logger)
|
||||
except Exception as exc: # pylint: disable=broad-exception-caught
|
||||
app_logger.exception("OUTBOX resend sweep failed: %s", exc)
|
||||
time.sleep(OUTBOX_RESEND_SECONDS)
|
||||
|
||||
|
||||
def replay_requests(app_logger) -> None:
|
||||
"""Re-enqueue accepted-but-unfinished searches after a restart.
|
||||
|
||||
Requests persisted by /query but never completed are put back on the internal
|
||||
queue. Those with a checkpoint resume mid-scan (answer_query loads it); the
|
||||
rest simply re-run. Marked 'queued' so the bot's watchdog sees them as known
|
||||
again."""
|
||||
pending = _requests.items()
|
||||
if not pending:
|
||||
return
|
||||
app_logger.info("Replaying %d unfinished search request(s) after restart", len(pending))
|
||||
for uuid, payload, _ts in pending:
|
||||
try:
|
||||
cl = Librarian(
|
||||
app, payload["query"], uuid,
|
||||
payload.get("deep_search", False), payload.get("callback", ""),
|
||||
)
|
||||
except Exception as exc: # pylint: disable=broad-exception-caught
|
||||
app_logger.warning("Cannot replay request %s (dropping): %s", uuid, exc)
|
||||
_forget_search(uuid)
|
||||
continue
|
||||
with _active_lock:
|
||||
active_queries[str(uuid)] = "queued"
|
||||
librarian_queue.put(cl)
|
||||
librarian_list.append(cl)
|
||||
|
||||
|
||||
def _handle_shutdown(signum, _frame) -> None:
|
||||
"""SIGTERM/SIGINT: ask the running search to checkpoint and stop. The main
|
||||
thread then waits (bounded) for it to finish - see __main__."""
|
||||
logging.getLogger("conjurer_librarian").warning(
|
||||
"Signal %s received - beginning graceful shutdown", signum
|
||||
)
|
||||
SHUTDOWN_EVENT.set()
|
||||
|
||||
|
||||
# trunk-ignore(pylint/R0902)
|
||||
class Librarian(object):
|
||||
"""
|
||||
Represents a librarian object that performs search and refinement operations on queries.
|
||||
"""
|
||||
|
||||
def __init__(self, _app, query, uuid, _deep_search, callback="") -> None:
|
||||
def __init__(self, _app, query, uuid, _deep_search) -> None:
|
||||
"""
|
||||
Initializes a Librarian object.
|
||||
|
||||
@@ -428,9 +111,6 @@ class Librarian(object):
|
||||
- search_result_from_cr: A dictionary to store the search results from Crossref.
|
||||
- done: A flag indicating if the search is done.
|
||||
"""
|
||||
# Crossref only needs a contact mailto. It can come from
|
||||
# CONJURER_CROSSREF_MAILTO (the usual container setup) OR from a
|
||||
# "crossref" entry in the netrc; netrc takes precedence when present.
|
||||
mailto_contact: Optional[str] = os.getenv("CONJURER_CROSSREF_MAILTO")
|
||||
if netrc:
|
||||
try:
|
||||
@@ -438,21 +118,9 @@ class Librarian(object):
|
||||
auth_tokens = netrc_mod.authenticators("crossref")
|
||||
if auth_tokens:
|
||||
mailto_contact = auth_tokens[0]
|
||||
except (FileNotFoundError, netrc.NetrcParseError) as exc:
|
||||
# A missing/unreadable netrc is NORMAL when the mailto is set via
|
||||
# env - don't cry wolf on every single search. Only warn when we
|
||||
# genuinely have no contact from either source.
|
||||
_log = logging.getLogger("conjurer_librarian")
|
||||
if mailto_contact:
|
||||
_log.debug(
|
||||
"netrc %s not used (%s) - using CONJURER_CROSSREF_MAILTO",
|
||||
NETRC_FILE, exc,
|
||||
)
|
||||
else:
|
||||
_log.warning(
|
||||
"Crossref contact not configured: netrc %s unreadable (%s) "
|
||||
"and CONJURER_CROSSREF_MAILTO unset",
|
||||
NETRC_FILE, exc,
|
||||
except (FileNotFoundError, netrc.NetrcParseError):
|
||||
logging.getLogger("conjurer_librarian").warning(
|
||||
"Crossref credentials missing in netrc %s", NETRC_FILE
|
||||
)
|
||||
if not mailto_contact:
|
||||
raise RuntimeError(
|
||||
@@ -475,13 +143,6 @@ class Librarian(object):
|
||||
self.search_result_from_cr = {}
|
||||
self.done = False
|
||||
self.deep_search = _deep_search
|
||||
# Where to send this search's result back to (the bot that asked). Lets
|
||||
# one librarian serve several bots; empty => static MAIN_BOT_ADDRESS.
|
||||
self.callback = callback or ""
|
||||
# Set True when a graceful shutdown interrupts this search mid-scan; the
|
||||
# worker then leaves the request + checkpoint in place instead of
|
||||
# delivering, so a restart resumes it.
|
||||
self.interrupted = False
|
||||
|
||||
async def search_crossref(self, query, deep_search=False):
|
||||
"""
|
||||
@@ -501,20 +162,14 @@ class Librarian(object):
|
||||
|
||||
if not deep_search:
|
||||
query_limit = MAX_CR_RESULTS if MAX_CR_RESULTS < 1000 else 1000
|
||||
cr_result = await asyncio.to_thread(
|
||||
_crossref_call, self.app.logger, "works",
|
||||
self.cr.works, query=query, limit=query_limit,
|
||||
)
|
||||
cr_result = self.cr.works(query=query, limit=query_limit)
|
||||
self.search_result_from_cr.update(cr_result)
|
||||
self.total = cr_result["message"]["total-results"]
|
||||
self.fetched += len(cr_result["message"]["items"])
|
||||
self.app.logger.info(self.total)
|
||||
self.app.logger.info(self.fetched)
|
||||
while self.total > self.fetched and self.limit > self.fetched:
|
||||
tmp_result = await asyncio.to_thread(
|
||||
_crossref_call, self.app.logger, "works(offset)",
|
||||
self.cr.works, query=query, limit=query_limit, offset=self.fetched,
|
||||
)
|
||||
tmp_result = self.cr.works(query=query, limit=query_limit, offset=self.fetched)
|
||||
cr_result["message"]["items"].extend(tmp_result["message"]["items"])
|
||||
self.total = tmp_result["message"]["total-results"]
|
||||
self.fetched = len(cr_result["message"]["items"])
|
||||
@@ -523,10 +178,7 @@ class Librarian(object):
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
else:
|
||||
cr_result = await asyncio.to_thread(
|
||||
_crossref_call, self.app.logger, "works(deep cursor)",
|
||||
self.cr.works, query=query, cursor_max=15000, cursor='*', progress_bar=True,
|
||||
)
|
||||
cr_result = self.cr.works(query=query, cursor_max=15000, cursor='*', progress_bar = True)
|
||||
result = cr_result[0]
|
||||
for item in cr_result[1:]:
|
||||
result["message"]["items"].extend(item["message"]["items"])
|
||||
@@ -543,7 +195,20 @@ class Librarian(object):
|
||||
self.app.logger.info("CROSSREF DONE")
|
||||
|
||||
self.app.logger.info("CROSSREF DONE")
|
||||
_dump_debug(lib_paths.CR_RESULTS, self.uuid, self.search_result_from_cr)
|
||||
with open(lib_paths.CR_RESULTS, "r+", encoding="utf-8") as data_file:
|
||||
# First we load existing data into a dict.
|
||||
try:
|
||||
file_data = json.load(data_file)
|
||||
except JSONDecodeError:
|
||||
file_data = {}
|
||||
data_file.truncate(0)
|
||||
data_file.seek(0)
|
||||
tmp = {self.uuid : self.search_result_from_cr}
|
||||
if file_data:
|
||||
file_data.update(tmp)
|
||||
else:
|
||||
file_data = tmp
|
||||
json.dump(file_data, data_file, indent=4)
|
||||
return cr_result
|
||||
|
||||
|
||||
@@ -594,10 +259,23 @@ class Librarian(object):
|
||||
|
||||
for item in temp:
|
||||
refined_result[item["DOI"]]= item
|
||||
_dump_debug(lib_paths.RR_RESULTS, self.uuid, refined_result)
|
||||
with open(lib_paths.RR_RESULTS, "r+", encoding="utf-8") as data_file:
|
||||
# First we load existing data into a dict.
|
||||
try:
|
||||
file_data = json.load(data_file)
|
||||
except JSONDecodeError:
|
||||
file_data = {}
|
||||
data_file.truncate(0)
|
||||
data_file.seek(0)
|
||||
tmp = {self.uuid: refined_result}
|
||||
if file_data:
|
||||
file_data.update(tmp)
|
||||
else:
|
||||
file_data = tmp
|
||||
json.dump(file_data, data_file, indent=4)
|
||||
return refined_result
|
||||
|
||||
async def check_if_exists(self, refined_result, resume=None):
|
||||
async def check_if_exists(self, refined_result):
|
||||
"""
|
||||
Checks if the given DOI exists.
|
||||
|
||||
@@ -611,36 +289,15 @@ class Librarian(object):
|
||||
Raises:
|
||||
- None.
|
||||
"""
|
||||
result = {}
|
||||
self.app.logger.info("REFINE: Running search in the backend app")
|
||||
dois = []
|
||||
for item, value in refined_result.items():
|
||||
dois.append([item, value])
|
||||
# Publish this scan as "the running search" so the heartbeat can report
|
||||
# it; cleared in finally so a finished/crashed scan never lingers there.
|
||||
progress = {}
|
||||
_set_current_search(self.uuid, self.query, progress, self.live_results)
|
||||
try:
|
||||
result, positions, interrupted = await asyncio.to_thread(
|
||||
search_bot.search_for_doi,
|
||||
dois, self.live_results, self.app.logger, SHUTDOWN_EVENT, resume,
|
||||
progress,
|
||||
coro = asyncio.to_thread(
|
||||
search_bot.search_for_doi, dois, self.live_results, self.app.logger
|
||||
)
|
||||
finally:
|
||||
_clear_current_search()
|
||||
if interrupted:
|
||||
# Graceful shutdown hit mid-scan: checkpoint found-so-far + per-file
|
||||
# resume offsets + the DOI list, so a restart continues instead of
|
||||
# restarting. The worker sees self.interrupted and does NOT deliver.
|
||||
found = [item["DOI"] for item in result if item["exists"]]
|
||||
_checkpoints.put(
|
||||
self.uuid,
|
||||
{"dois": refined_result, "found": found, "positions": positions},
|
||||
)
|
||||
self.interrupted = True
|
||||
self.app.logger.info(
|
||||
"Search %s checkpointed (%d found so far) for resume", self.uuid, len(found)
|
||||
)
|
||||
return [], []
|
||||
result = await coro
|
||||
result_list = []
|
||||
result_no_db = []
|
||||
for item in result:
|
||||
@@ -665,48 +322,21 @@ class Librarian(object):
|
||||
Raises:
|
||||
- None.
|
||||
"""
|
||||
checkpoint = _checkpoints.get(self.uuid)
|
||||
if checkpoint is not None:
|
||||
# Resume a search interrupted by a previous shutdown: the expensive
|
||||
# Crossref + refine work is already captured in the checkpoint, so go
|
||||
# straight to the DB scan with the saved offsets + found-so-far.
|
||||
self.app.logger.info(
|
||||
"Resuming search %s from checkpoint (%d found so far)",
|
||||
self.uuid, len(checkpoint.get("found", [])),
|
||||
)
|
||||
refined_result = checkpoint["dois"]
|
||||
resume = {
|
||||
"found": checkpoint.get("found", []),
|
||||
"positions": checkpoint.get("positions", {}),
|
||||
}
|
||||
else:
|
||||
# Cache hit: an identical (normalised) query ran recently - return its
|
||||
# stored hits and skip Crossref + the whole DB scan entirely.
|
||||
cached = _cache_get(self.query, deep_search)
|
||||
if cached is not None:
|
||||
self.final_result = cached
|
||||
self.app.logger.info(
|
||||
"Cache HIT for %s (%d hits): %s", self.uuid, len(cached), self.query
|
||||
)
|
||||
return self.final_result
|
||||
self.app.logger.info(f"Search started {self.uuid}")
|
||||
cr_result = await self.search_crossref(query=self.query, deep_search=deep_search)
|
||||
refined_result = await self.refine_search(cr_result)
|
||||
resume = None
|
||||
answer, negative_answer = await self.check_if_exists(refined_result)
|
||||
|
||||
answer, negative_answer = await self.check_if_exists(refined_result, resume=resume)
|
||||
if self.interrupted:
|
||||
# Graceful shutdown mid-scan: checkpoint is written, request stays.
|
||||
# Signal the worker (None) NOT to deliver - a restart resumes this.
|
||||
return None
|
||||
self.app.logger.info("Returning result")
|
||||
self.app.logger.info(answer)
|
||||
self.app.logger.info(negative_answer)
|
||||
|
||||
for item in answer:
|
||||
self.final_result[item["DOI"]] = {"Title": item["data"]["title"], "type": item["data"]["type"]}
|
||||
for item in negative_answer:
|
||||
self.not_in_db[item["DOI"]] = {"Title": item["data"]["title"], "type": item["data"]["type"]}
|
||||
self.app.logger.info("Search %s produced %d hits", self.uuid, len(self.final_result))
|
||||
# Cache the completed result so a repeat of this query is instant.
|
||||
_cache_put(self.query, deep_search, self.final_result)
|
||||
self.app.logger.info("Returning result case2")
|
||||
self.app.logger.info(self.final_result)
|
||||
return self.final_result
|
||||
|
||||
# ============================= FLASK INTERNALS===============================
|
||||
@@ -778,15 +408,10 @@ class BackgroundTaskSearch(threading.Thread):
|
||||
|
||||
The search task continues running indefinitely until the thread is stopped.
|
||||
"""
|
||||
while not SHUTDOWN_EVENT.is_set():
|
||||
while True:
|
||||
database = None
|
||||
ndb_database = None
|
||||
# Bounded get so the loop can observe SHUTDOWN_EVENT while idle
|
||||
# (blocked on a plain get() it would never notice a shutdown).
|
||||
try:
|
||||
item = librarian_queue.get(timeout=1)
|
||||
except Empty:
|
||||
continue
|
||||
item = librarian_queue.get()
|
||||
# Health-check ping: it has flowed through the internal queue and is
|
||||
# now pulled off it - that is the whole point. Pong it straight back
|
||||
# with the same uuid and DO NOT run a search.
|
||||
@@ -796,32 +421,24 @@ class BackgroundTaskSearch(threading.Thread):
|
||||
"PING %s pulled off internal queue - ponging back (no search)",
|
||||
ping_uuid,
|
||||
)
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
_post_pong, self.app.logger, ping_uuid, item.get("callback", "")
|
||||
requests.post,
|
||||
f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}",
|
||||
json={"__pong__": ping_uuid},
|
||||
headers=_service_headers(),
|
||||
timeout=5,
|
||||
)
|
||||
except requests.exceptions.RequestException as exc:
|
||||
self.app.logger.warning(
|
||||
"PING pong send failed for %s: %s", ping_uuid, exc
|
||||
)
|
||||
continue
|
||||
librarian = item
|
||||
# Mark busy + processing for the whole search, and ALWAYS clear both
|
||||
# (even on a crash) in finally: worker_busy so a ping doesn't wait
|
||||
# behind us, and active_queries so the bot's watchdog can tell a
|
||||
# finished-and-gone query from one still in flight.
|
||||
worker_busy.set()
|
||||
requeued = False # set when a crash schedules another attempt
|
||||
with _active_lock:
|
||||
active_queries[str(librarian.uuid)] = "processing"
|
||||
try:
|
||||
self.app.logger.info("Processing search %s", librarian.uuid)
|
||||
self.app.logger.info("STARTED")
|
||||
result = await librarian.answer_query(librarian.deep_search)
|
||||
if result is None:
|
||||
# Graceful shutdown interrupted this search mid-scan. Its
|
||||
# checkpoint + persisted request stay in place, so a restart
|
||||
# picks it up and RESUMES from where it stopped.
|
||||
self.app.logger.info(
|
||||
"Search %s interrupted by shutdown - will resume on restart",
|
||||
librarian.uuid,
|
||||
)
|
||||
break
|
||||
result = {librarian.uuid: result}
|
||||
self.app.logger.info("Saving to file")
|
||||
|
||||
# Save results to "not_in_db.json" file
|
||||
with open(lib_paths.NOT_IN_DB, "r+", encoding="utf-8") as ndb_file:
|
||||
@@ -838,81 +455,68 @@ class BackgroundTaskSearch(threading.Thread):
|
||||
ndb_file.seek(0)
|
||||
json.dump(ndb_database, ndb_file)
|
||||
|
||||
# Optional debug dump of the final result (off by default).
|
||||
_dump_debug(lib_paths.S_RESULTS, librarian.uuid, result[librarian.uuid])
|
||||
self.app.logger.info("Search %s finished", librarian.uuid)
|
||||
# Save results to "s_results.json" file
|
||||
with open(lib_paths.S_RESULTS, "r+", encoding="utf-8") as s_file:
|
||||
database = {}
|
||||
try:
|
||||
database = json.load(s_file)
|
||||
except JSONDecodeError:
|
||||
pass
|
||||
if database:
|
||||
self.app.logger.info(database)
|
||||
self.app.logger.info(result)
|
||||
database.update(result)
|
||||
else:
|
||||
database = result
|
||||
self.app.logger.info("DUMPING DATA")
|
||||
s_file.truncate(0)
|
||||
s_file.seek(0)
|
||||
json.dump(database, s_file)
|
||||
self.app.logger.info("FINISHED")
|
||||
|
||||
# Persist the result to the durable OUTBOX FIRST, then try to
|
||||
# deliver it. Writing to disk before sending is the whole point:
|
||||
# an expensive (hours-long) result now survives a failed send, a
|
||||
# bot outage, or a librarian restart - the resender keeps
|
||||
# retrying until the bot ACKs, and only then is it removed.
|
||||
# Send the result back to the bot. Log EXACTLY what goes out (target,
|
||||
# uuid, how many DOIs and which) so the librarian log makes it plain a
|
||||
# result was sent and what was in it.
|
||||
payload = result # shape: {uuid: {DOI: {"Title": ..., "type": ...}}}
|
||||
uuid = str(librarian.uuid)
|
||||
target = librarian.callback # answer the bot that asked
|
||||
hits = payload.get(librarian.uuid, {}) if isinstance(payload, dict) else {}
|
||||
# OUTBOX entry carries the origin bot's address so the resender
|
||||
# delivers it to the right bot even after a librarian restart.
|
||||
_outbox.put(uuid, {"target": target, "payload": payload})
|
||||
target = f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}"
|
||||
self.app.logger.info(
|
||||
"SENDING result for %s to %s: %d DOI(s): %s (queued to OUTBOX)",
|
||||
uuid, target or "default", len(hits), list(hits.keys()),
|
||||
"SENDING result for %s to %s: %d DOI(s): %s",
|
||||
librarian.uuid,
|
||||
target,
|
||||
len(hits),
|
||||
list(hits.keys()),
|
||||
)
|
||||
if await asyncio.to_thread(_deliver_result, target, uuid, payload, self.app.logger):
|
||||
_outbox.remove(uuid)
|
||||
else:
|
||||
self.app.logger.warning(
|
||||
"Result %s not acked yet - left in OUTBOX for the resender", uuid
|
||||
# A failed send must NOT kill this worker - otherwise a bot that is
|
||||
# momentarily down stalls every future query until the librarian is
|
||||
# restarted. Log and carry on to the next queued search.
|
||||
try:
|
||||
response = await asyncio.to_thread(
|
||||
requests.post,
|
||||
target,
|
||||
json=payload,
|
||||
headers=_service_headers(),
|
||||
timeout=360,
|
||||
)
|
||||
# Computed + handed to the durable OUTBOX: the search is done, so
|
||||
# forget its request + checkpoint (never replay/resume it again).
|
||||
_forget_search(uuid)
|
||||
except Exception as exc: # pylint: disable=broad-exception-caught
|
||||
# A crashing search must not kill the worker thread (that would
|
||||
# freeze the whole queue). It also must not silently vanish just
|
||||
# because Crossref timed out once: retry the whole search a
|
||||
# bounded number of times (attempt count persisted with the
|
||||
# request, so it can't loop forever), keeping any checkpoint so a
|
||||
# crashed DB scan resumes rather than restarts. Only after
|
||||
# SEARCH_MAX_ATTEMPTS do we give up and let the bot's watchdog
|
||||
# tell the user it vanished.
|
||||
uuid = str(librarian.uuid)
|
||||
self.app.logger.exception("Search %s crashed: %s", uuid, exc)
|
||||
stored = _requests.get(uuid) or {}
|
||||
attempts = int(stored.get("attempts", 0)) + 1
|
||||
if attempts < max(1, SEARCH_MAX_ATTEMPTS):
|
||||
stored.update({
|
||||
"query": librarian.query,
|
||||
"deep_search": librarian.deep_search,
|
||||
"callback": librarian.callback,
|
||||
"attempts": attempts,
|
||||
})
|
||||
_requests.put(uuid, stored)
|
||||
librarian_queue.put(
|
||||
Librarian(self.app, librarian.query, uuid,
|
||||
librarian.deep_search, librarian.callback)
|
||||
)
|
||||
requeued = True
|
||||
self.app.logger.warning(
|
||||
"Search %s requeued after crash (attempt %d/%d)",
|
||||
uuid, attempts, SEARCH_MAX_ATTEMPTS,
|
||||
if response.status_code == 200:
|
||||
self.app.logger.info(
|
||||
"SENT result for %s -> HTTP 200 (bot accepted)", librarian.uuid
|
||||
)
|
||||
else:
|
||||
self.app.logger.warning(
|
||||
"SENT result for %s but bot returned HTTP %s: %s",
|
||||
librarian.uuid,
|
||||
response.status_code,
|
||||
response.text[:500],
|
||||
)
|
||||
except requests.exceptions.RequestException as exc:
|
||||
self.app.logger.error(
|
||||
"Search %s failed %d times - giving up", uuid, attempts
|
||||
"FAILED to send result for %s to %s: %s",
|
||||
librarian.uuid,
|
||||
target,
|
||||
exc,
|
||||
)
|
||||
_forget_search(uuid)
|
||||
finally:
|
||||
worker_busy.clear()
|
||||
with _active_lock:
|
||||
if requeued:
|
||||
# Still known to the bot's watchdog - it's going round again.
|
||||
active_queries[str(librarian.uuid)] = "queued"
|
||||
else:
|
||||
active_queries.pop(str(librarian.uuid), None)
|
||||
await asyncio.sleep(1)
|
||||
SHUTDOWN_DONE.set()
|
||||
self.app.logger.info("Search worker stopped cleanly")
|
||||
|
||||
|
||||
# ==================================SERVER ROUTES==========================================
|
||||
@@ -933,25 +537,14 @@ async def query_database():
|
||||
tuple: A tuple containing a JSON response and a status code.
|
||||
"""
|
||||
record = json.loads(request.data)
|
||||
app.logger.info(record)
|
||||
app.logger.info(record["query"])
|
||||
app.logger.info(record["UUID"])
|
||||
uuid = record["UUID"]
|
||||
deep_search = record["deep_search"]
|
||||
# Where to answer THIS query - the bot that sent it. Persisted with the
|
||||
# request so a replay after restart still answers the right bot.
|
||||
callback = record.get("callback", "")
|
||||
app.logger.info("Query accepted %s (callback %s): %s", uuid, callback or "default", record["query"])
|
||||
# Persist the request BEFORE enqueuing, so an accepted search survives a
|
||||
# restart (it is replayed on startup) - not just an in-progress one.
|
||||
_requests.put(
|
||||
str(uuid),
|
||||
{"query": record["query"], "deep_search": deep_search, "callback": callback},
|
||||
)
|
||||
cl = Librarian(app, record["query"], uuid, deep_search, callback)
|
||||
cl = Librarian(app, record["query"], uuid, deep_search)
|
||||
librarian_queue.put(cl)
|
||||
librarian_list.append(cl)
|
||||
# The bot's per-query watchdog polls /query_status for this uuid; mark it
|
||||
# "queued" now so it counts as known the moment we accept it.
|
||||
with _active_lock:
|
||||
active_queries[str(uuid)] = "queued"
|
||||
answer_data = (record["query"], record["UUID"], librarian_queue.qsize())
|
||||
return_data = (
|
||||
jsonify(isError=False, message="Success", statusCode=200, data=answer_data),
|
||||
@@ -961,68 +554,28 @@ async def query_database():
|
||||
|
||||
|
||||
@app.route("/ping", methods=["POST"])
|
||||
def ping_roundtrip():
|
||||
async def ping_roundtrip():
|
||||
_authorize_request()
|
||||
"""
|
||||
Health-check round-trip.
|
||||
|
||||
Two cases, one guarantee - the pong always comes back over the librarian->bot
|
||||
return path (the only thing the ping must prove):
|
||||
|
||||
* IDLE: put a ping marker onto the SAME internal ``librarian_queue`` real
|
||||
searches use and return 200. The worker pulls it off and pongs it back,
|
||||
so a successful pong proves the whole pipeline flows (queue + worker + the
|
||||
return leg), not just that Flask is up.
|
||||
* BUSY (a search is grinding): DO NOT queue - the ping would just wait behind
|
||||
a possibly hours-long search and time out, making a perfectly healthy busy
|
||||
librarian look dead. Pong back immediately instead. Being busy is fine; you
|
||||
can keep piling searches on. The ping only needs to catch a BROKEN return
|
||||
path, and the direct pong exercises exactly that.
|
||||
Puts a lightweight ping marker onto the SAME internal ``librarian_queue``
|
||||
that real searches go through and returns 200 immediately. The background
|
||||
worker pulls it off the queue and pongs it back to the bot with the same
|
||||
uuid, WITHOUT running any Crossref/DOI search. A successful pong therefore
|
||||
proves the whole pipeline (HTTP in -> internal queue -> worker -> HTTP out)
|
||||
is flowing, not just that Flask is up.
|
||||
"""
|
||||
record = json.loads(request.data)
|
||||
ping_uuid = record["UUID"]
|
||||
# Pong goes back to the bot that pinged (carried through the queue when idle),
|
||||
# so a shared librarian health-checks correctly for every bot.
|
||||
callback = record.get("callback", "")
|
||||
if worker_busy.is_set():
|
||||
app.logger.info("PING %s while busy grinding - direct pong (skip queue)", ping_uuid)
|
||||
_post_pong(app.logger, ping_uuid, callback)
|
||||
else:
|
||||
app.logger.info("PING received %s - queued for round-trip", ping_uuid)
|
||||
librarian_queue.put({"__ping__": ping_uuid, "callback": callback})
|
||||
librarian_queue.put({"__ping__": ping_uuid})
|
||||
return (
|
||||
jsonify(isError=False, message="ping-queued", statusCode=200, data=ping_uuid),
|
||||
200,
|
||||
)
|
||||
|
||||
|
||||
@app.route("/query_status", methods=["POST"])
|
||||
def query_status():
|
||||
_authorize_request()
|
||||
"""
|
||||
Per-query watchdog probe.
|
||||
|
||||
Returns whether ``UUID`` is still known to the librarian (queued or being
|
||||
processed). The bot polls this after dispatching a search: while the uuid is
|
||||
known the search is progressing; once it VANISHES here without the result
|
||||
ever reaching the bot, the result was lost in transit and the bot tells the
|
||||
user. A busy/queued search is never mistaken for a lost one.
|
||||
"""
|
||||
record = json.loads(request.data)
|
||||
uuid = str(record["UUID"])
|
||||
with _active_lock:
|
||||
state = active_queries.get(uuid, "unknown")
|
||||
return (
|
||||
jsonify(
|
||||
isError=False,
|
||||
message="Success",
|
||||
statusCode=200,
|
||||
data={"uuid": uuid, "known": state != "unknown", "state": state},
|
||||
),
|
||||
200,
|
||||
)
|
||||
|
||||
|
||||
@app.route("/get_partial_result", methods=["POST"])
|
||||
async def get_partial():
|
||||
_authorize_request()
|
||||
@@ -1048,10 +601,7 @@ async def get_partial():
|
||||
|
||||
# =======================================MAIN===================================================
|
||||
if __name__ == "__main__":
|
||||
# Default INFO (readable). Set CONJURER_LIBRARIAN_LOG_LEVEL=DEBUG for the
|
||||
# full per-file / per-line search chatter.
|
||||
app.logger.setLevel(LOG_LEVEL)
|
||||
_fmt = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
app.logger.setLevel(logging.DEBUG)
|
||||
LOGFILE_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
h1 = handlers.RotatingFileHandler(
|
||||
filename=str(LOGFILE_PATH),
|
||||
@@ -1060,20 +610,8 @@ if __name__ == "__main__":
|
||||
maxBytes=6 * 1024 * 1024,
|
||||
backupCount=6,
|
||||
)
|
||||
h1.setFormatter(_fmt)
|
||||
|
||||
app.logger.addHandler(h1)
|
||||
# Console handler so `kubectl logs` shows what's happening (k8s reads stdout);
|
||||
# the search internals no longer print() straight to stdout.
|
||||
_console = logging.StreamHandler()
|
||||
_console.setFormatter(_fmt)
|
||||
app.logger.addHandler(_console)
|
||||
|
||||
# Graceful shutdown: on SIGTERM (k8s) / SIGINT the running search checkpoints
|
||||
# itself and the worker stops; the main thread then exits within a bounded
|
||||
# window so we never linger as an un-killable zombie pod.
|
||||
signal.signal(signal.SIGTERM, _handle_shutdown)
|
||||
signal.signal(signal.SIGINT, _handle_shutdown)
|
||||
|
||||
threads = []
|
||||
threads.append(threading.Thread(target=waitress_run, daemon=True))
|
||||
# threads.append(threading.Thread(target=flask_debug))
|
||||
@@ -1086,31 +624,13 @@ if __name__ == "__main__":
|
||||
target=scrape_bot.scraper, args=(app.logger,), daemon=True
|
||||
)
|
||||
)
|
||||
# Durable delivery: keep flushing the OUTBOX so any result not yet acked by
|
||||
# the bot (transient outage, or left over from before a restart) is resent.
|
||||
threads.append(
|
||||
threading.Thread(target=outbox_resender, args=(app.logger,), daemon=True)
|
||||
)
|
||||
# "Still searching" heartbeat, so an hours-long scan isn't radio silence.
|
||||
threads.append(
|
||||
threading.Thread(target=search_heartbeat, args=(app.logger,), daemon=True)
|
||||
)
|
||||
i = 0
|
||||
try:
|
||||
for worker in threads:
|
||||
app.logger.info("App number: %s", i)
|
||||
i += 1
|
||||
worker.start()
|
||||
# Re-enqueue searches that were accepted/in-progress before the last stop.
|
||||
replay_requests(app.logger)
|
||||
app.logger.info("Librarian ready (graceful-shutdown timeout %ss)", GRACEFUL_TIMEOUT)
|
||||
|
||||
# Main thread parks until a shutdown signal, then gives the worker a BOUNDED
|
||||
# window to checkpoint. sleep() (not Event.wait) so the signal is delivered
|
||||
# promptly to this thread on every platform.
|
||||
while not SHUTDOWN_EVENT.is_set():
|
||||
time.sleep(0.5)
|
||||
app.logger.info("Waiting up to %ss for the search to checkpoint...", GRACEFUL_TIMEOUT)
|
||||
if SHUTDOWN_DONE.wait(GRACEFUL_TIMEOUT):
|
||||
app.logger.info("Graceful shutdown complete - state saved")
|
||||
else:
|
||||
app.logger.warning(
|
||||
"Graceful shutdown TIMED OUT after %ss - exiting anyway (no zombie)",
|
||||
GRACEFUL_TIMEOUT,
|
||||
)
|
||||
for worker in threads:
|
||||
worker.join()
|
||||
except KeyboardInterrupt:
|
||||
app.logger.info("Shutdown requested - exiting librarian service")
|
||||
|
||||
@@ -106,10 +106,8 @@ def check_if_exists_brute_force(logger):
|
||||
):
|
||||
pass
|
||||
if blocked:
|
||||
# Expected, routine sci-hub behaviour (we back off an hour and carry
|
||||
# on) - WARNING, not ERROR, so it stops masquerading as a fault when
|
||||
# you're scanning the log for real problems.
|
||||
logger.warning("Got blocked. Fuck. Backing off an hour: %s", item[0])
|
||||
logger.info(item)
|
||||
logger.error("Got blocked. Fuck.")
|
||||
time.sleep(60 * 60)
|
||||
# trunk-ignore(bandit/B311)
|
||||
rand = random.randint(1, 60)
|
||||
|
||||
@@ -21,7 +21,7 @@ Global Variables:
|
||||
# TODO: Wpiemdolić to wszystko w klasę z loggerem przysłanym z góry
|
||||
import os
|
||||
import re
|
||||
from queue import Empty, Full, Queue
|
||||
from queue import Empty, Queue
|
||||
from threading import Thread
|
||||
import time
|
||||
q = Queue()
|
||||
@@ -44,13 +44,7 @@ CHUNK = os.getenv("CONJURER_LIBRARIAN_CHUNK", "_chunk.txt")
|
||||
MAXTHREADS = int(os.getenv("CONJURER_LIBRARIAN_MAXTHREADS", "0"))
|
||||
|
||||
_sentinel = object()
|
||||
# BOUNDED work queue. The producers stream the WHOLE DOI database (potentially
|
||||
# tens of millions of lines across chunks) into this queue; the previous cap of
|
||||
# 35_500_000 items was effectively unbounded (~3.5 GB of buffered lines), which
|
||||
# OOM-killed the 1 GiB container mid-search. A small bound makes the producers
|
||||
# backpressure to the consumers, keeping RAM to a few MB. The producer put below
|
||||
# stays responsive to the TERM sentinel so a full queue can never deadlock it.
|
||||
WORK_Q_SIZE = int(os.getenv("CONJURER_LIBRARIAN_WORKQ_SIZE", "100000"))
|
||||
WORK_Q_SIZE = 35500000
|
||||
# Idle backstop: after this many consecutive empty seconds a consumer assumes
|
||||
# the producers are done (or dead) and exits, so the search can never hang even
|
||||
# if a sentinel were somehow lost. The primary, correct termination is still the
|
||||
@@ -85,17 +79,15 @@ def discover_chunk_files(_logger):
|
||||
return ordered
|
||||
|
||||
|
||||
def producer(out_q, control_q, filename, _logger, stop_event=None, positions=None,
|
||||
start_offsets=None):
|
||||
"""Stream a chunk file's lines onto the work queue, resumably.
|
||||
def producer(out_q, control_q, filename, _logger):
|
||||
"""
|
||||
Produces items from the output queue and puts them into the control queue.
|
||||
|
||||
``start_offsets[filename]`` (a tell() cookie) is where to RESUME reading from
|
||||
- so a search continued after a restart skips the part already scanned.
|
||||
``positions[filename]`` is updated to the tell() cookie just PAST each line
|
||||
successfully enqueued; because search_for_doi drains the queue before it
|
||||
returns, that cookie is a safe "everything up to here is processed" watermark
|
||||
to checkpoint. ``stop_event`` (graceful shutdown) makes the producer stop
|
||||
reading and record its watermark, mirroring the early-TERM path.
|
||||
Args:
|
||||
out_q (Queue): Output queue.
|
||||
control_q (Queue): Control queue.
|
||||
filename (str): Name of the file.
|
||||
_logger: Logger object for logging.
|
||||
"""
|
||||
try:
|
||||
# errors="replace" so a stray non-UTF-8 byte in a chunk (they happen in
|
||||
@@ -104,68 +96,28 @@ def producer(out_q, control_q, filename, _logger, stop_event=None, positions=Non
|
||||
# producer partway and leaving every DOI after the bad byte unsearched.
|
||||
# DOIs are ASCII, so a replaced byte can only affect junk, never a match.
|
||||
with open(DATABASE_PATH + filename, "r", encoding=ENCODING, errors="replace") as operated_file:
|
||||
if start_offsets and filename in start_offsets:
|
||||
operated_file.seek(start_offsets[filename])
|
||||
_logger.debug("Producer %s: resuming at offset %s", filename, start_offsets[filename])
|
||||
else:
|
||||
_logger.debug("Producer started: %s", filename)
|
||||
print(f"Worker {filename} ")
|
||||
line_no = 0
|
||||
while True:
|
||||
if stop_event is not None and stop_event.is_set():
|
||||
_logger.debug("Producer %s: stop requested (graceful)", filename)
|
||||
break
|
||||
line = operated_file.readline()
|
||||
line_no += 1
|
||||
# Coarse progress at DEBUG only - the old per-line carriage-return
|
||||
# print flooded stdout / the log file with millions of lines.
|
||||
if line_no % 500000 == 0:
|
||||
_logger.debug("Producer %s: %d lines read", filename, line_no)
|
||||
print(f"\t \t \t \t \t \t W{filename}{line_no}\r", end="")
|
||||
|
||||
if not line:
|
||||
# EOF: record the end offset so a resume seeks here and stops
|
||||
# immediately (the file is fully scanned).
|
||||
if positions is not None:
|
||||
positions[filename] = operated_file.tell()
|
||||
_logger.debug("Producer %s: EOF at %d lines", filename, line_no)
|
||||
print(f"EOF {filename}")
|
||||
break
|
||||
|
||||
# Backpressure-safe put onto the BOUNDED queue: wait for room, but
|
||||
# keep polling stop_event / the TERM sentinel so a full queue whose
|
||||
# consumers have already finished can never deadlock us here.
|
||||
stopped = False
|
||||
while True:
|
||||
try:
|
||||
out_q.put(line, timeout=1)
|
||||
break
|
||||
except Full:
|
||||
if stop_event is not None and stop_event.is_set():
|
||||
stopped = True
|
||||
break
|
||||
try:
|
||||
if control_q.get(block=False) is _sentinel:
|
||||
control_q.put(_sentinel)
|
||||
stopped = True
|
||||
break
|
||||
except Empty:
|
||||
pass
|
||||
if stopped:
|
||||
_logger.debug("Producer %s: stop while enqueuing", filename)
|
||||
break
|
||||
# Watermark AFTER a successful enqueue: safe to resume past here
|
||||
# once the queue drains (which it does before search_for_doi ends).
|
||||
if positions is not None:
|
||||
positions[filename] = operated_file.tell()
|
||||
|
||||
out_q.put(line)
|
||||
try:
|
||||
check = control_q.get(block=False)
|
||||
except Empty:
|
||||
check = False
|
||||
|
||||
if check is _sentinel:
|
||||
_logger.debug("Producer %s: TERM signal", filename)
|
||||
print("TERM signal received")
|
||||
control_q.put(check)
|
||||
break
|
||||
_logger.debug("Producer finished: %s", filename)
|
||||
print(f"Worker finished: {filename}")
|
||||
except Exception as exc: # pylint: disable=broad-except
|
||||
# No per-file error (missing/unreadable chunk, a decode edge case that
|
||||
# slips past errors="replace", anything unforeseen) may take the whole
|
||||
@@ -173,6 +125,7 @@ def producer(out_q, control_q, filename, _logger, stop_event=None, positions=Non
|
||||
# the sentinel below still fires (finally), so the consumers' count stays
|
||||
# correct and nothing deadlocks or silently loses a producer.
|
||||
_logger.warning("Chunk %s failed, skipping rest of it: %s", filename, exc)
|
||||
print(f"Worker {filename} failed: {exc}")
|
||||
finally:
|
||||
# ALWAYS emit exactly one sentinel per producer, on every exit path (EOF,
|
||||
# early TERM, or crash). This is what lets the consumers count producers
|
||||
@@ -192,102 +145,70 @@ def consumer(in_q, control_q, doi, live_results, result_list, control_dict, expe
|
||||
live_results (list): List to store the search results.
|
||||
_logger: Logger object for logging.
|
||||
"""
|
||||
_logger.debug("Consumer %s started", no)
|
||||
print(f"Consumer thread started: {no} no")
|
||||
empty_counter = 0
|
||||
alive_no = 0
|
||||
# DOI -> result item, so a line is matched with one O(1) dict lookup instead
|
||||
# of scanning every queried DOI. Items are shared with result_list, so
|
||||
# setting exists here is seen by everyone.
|
||||
doi_index = {item["DOI"]: item for item in result_list}
|
||||
while True:
|
||||
done_check = True
|
||||
try:
|
||||
data = in_q.get(block=True, timeout = 1)
|
||||
if data is _sentinel:
|
||||
control_dict["sentinels"] += 1
|
||||
_logger.debug(
|
||||
"Consumer %s: producer done (%d/%d)",
|
||||
no, control_dict["sentinels"], expected_sentinels,
|
||||
)
|
||||
print(f"Workers finished: {control_dict['sentinels']} reported by consumer {no}")
|
||||
|
||||
else:
|
||||
empty_counter = 0
|
||||
alive_no += 1
|
||||
print(f"C{no}__{alive_no}\r", end="")
|
||||
|
||||
# Each DB line is a DOI (optionally followed by metadata). Match
|
||||
# the WHOLE first token exactly - the old `item["DOI"] in data`
|
||||
# was a substring test, so a DOI that is a prefix of a longer one
|
||||
# (10.1/1 vs 10.1/12) produced a false 'exists' hit.
|
||||
parts = data.split()
|
||||
line_doi = parts[0] if parts else ""
|
||||
item = doi_index.get(line_doi)
|
||||
if item is not None and not item["exists"]:
|
||||
# HIT can fire thousands of times for a deep search -> DEBUG.
|
||||
_logger.debug("HIT %s (consumer %s)", line_doi, no)
|
||||
for item in result_list:
|
||||
if item["DOI"] in data and not item["exists"]:
|
||||
print(f"HIT in {no} content {data[0]} line {data[1]} file {data[2]} {item['exists']}")
|
||||
_logger.info(data)
|
||||
_logger.info("HIT")
|
||||
item["exists"] = True
|
||||
live_results.append(item)
|
||||
# All found? Signal producers to stop early (rare -> cheap).
|
||||
if all(it["exists"] for it in result_list):
|
||||
done_check = done_check and item["exists"]
|
||||
if done_check:
|
||||
control_q.put(_sentinel)
|
||||
except Empty:
|
||||
empty_counter += 1
|
||||
time.sleep(1)
|
||||
print(f"Consumer {no} empty")
|
||||
# Order matters: the >EMPTY_LIMIT break must be checked BEFORE the
|
||||
# lesser threshold, otherwise (as in the original) the first branch
|
||||
# always wins and the break is dead code, leaving the sentinel count
|
||||
# as the only exit - which is exactly what used to hang the search.
|
||||
if empty_counter > EMPTY_LIMIT:
|
||||
_logger.debug("Consumer %s finished (idle backstop)", no)
|
||||
print(f"Consumer thread finished {no} (idle backstop)")
|
||||
break
|
||||
if empty_counter > 5:
|
||||
print(f"Consumer {no} empty lvl 2")
|
||||
time.sleep(2)
|
||||
|
||||
if control_dict["sentinels"] >= expected_sentinels:
|
||||
_logger.debug("Consumer %s: all producers finished", no)
|
||||
_logger.info(f"All workers finished {no}")
|
||||
break
|
||||
|
||||
|
||||
|
||||
def search_for_doi(doi, live_results, _logger, stop_event=None, resume=None,
|
||||
progress=None):
|
||||
"""Search for DOI in live_results, resumably.
|
||||
def search_for_doi(doi, live_results, _logger):
|
||||
"""
|
||||
Search for DOI in live_results using _logger for logging.
|
||||
|
||||
Returns ``(result_list, positions, interrupted)``:
|
||||
* ``result_list`` - the queried DOIs with their ``exists`` flag,
|
||||
* ``positions`` - ``{filename: tell()-cookie}`` safe-to-resume watermarks
|
||||
(the queue is drained before return, so everything up to each cookie is
|
||||
processed),
|
||||
* ``interrupted`` - True if ``stop_event`` fired (the scan is PARTIAL; check
|
||||
point ``positions`` + the found DOIs and call again with ``resume=`` to
|
||||
continue where it left off).
|
||||
|
||||
``resume`` is ``{"positions": {...}, "found": [doi, ...]}`` from a previous
|
||||
interrupted run: already-found DOIs are pre-marked and each producer seeks to
|
||||
its saved offset, so no already-scanned line is read twice.
|
||||
|
||||
``progress``, if given, is a dict this fills with ``positions`` (the LIVE
|
||||
dict, updated as producers read) and ``total_bytes`` (summed once, up front).
|
||||
That makes a rough "how far along" reading free: sum the offsets, divide by
|
||||
the total - no counting, no extra work in the read loop.
|
||||
Args:
|
||||
doi (list): List of DOI to search for.
|
||||
live_results (list): List to store the search results.
|
||||
_logger: Logger object for logging.
|
||||
"""
|
||||
control_dict = {"sentinels":0}
|
||||
result_list = []
|
||||
threads = []
|
||||
work_q = Queue(maxsize=WORK_Q_SIZE)
|
||||
control_q = Queue()
|
||||
positions = {}
|
||||
|
||||
resume = resume or {}
|
||||
already_found = set(resume.get("found", []))
|
||||
start_offsets = resume.get("positions", {})
|
||||
|
||||
for item in doi:
|
||||
entry = {"DOI": item[0], "exists": False, "data": item[1]}
|
||||
if item[0] in already_found:
|
||||
# Pre-mark hits from the previous (interrupted) run so we neither
|
||||
# re-scan for them nor drop them from live_results.
|
||||
entry["exists"] = True
|
||||
live_results.append(entry)
|
||||
result_list.append(entry)
|
||||
result_list.append({"DOI": item[0], "exists": False, "data": item[1]})
|
||||
|
||||
# One producer per chunk file that actually exists; the sentinel threshold is
|
||||
# that same count, so the two can never drift apart the way MAXTHREADS did.
|
||||
@@ -297,42 +218,25 @@ def search_for_doi(doi, live_results, _logger, stop_event=None, resume=None,
|
||||
_logger.error(
|
||||
"No '<n>%s' chunk files in %s - DOI search cannot run", CHUNK, DATABASE_PATH
|
||||
)
|
||||
return result_list, positions, bool(stop_event and stop_event.is_set())
|
||||
|
||||
if progress is not None:
|
||||
# One stat() per chunk file, ONCE - then progress is just sum(positions)
|
||||
# / total_bytes, with nothing extra happening per line.
|
||||
total_bytes = 0
|
||||
for name in chunk_files:
|
||||
try:
|
||||
total_bytes += os.path.getsize(DATABASE_PATH + name)
|
||||
except OSError:
|
||||
pass
|
||||
progress["positions"] = positions # live dict, updated by the producers
|
||||
progress["total_bytes"] = total_bytes
|
||||
progress["chunk_files"] = expected
|
||||
return result_list
|
||||
|
||||
for i in range (0, (len(doi)//1000)+2):
|
||||
t_cons = Thread(
|
||||
target=consumer,
|
||||
args=(work_q, control_q, doi, live_results, result_list, control_dict, expected, i, _logger),
|
||||
)
|
||||
_logger.debug("Consumer thread created")
|
||||
_logger.info("Consumer thread created")
|
||||
threads.append(t_cons)
|
||||
for filename in chunk_files:
|
||||
_logger.debug("Creating worker thread for %s", filename)
|
||||
_logger.info("Creating worker thread for %s", filename)
|
||||
threads.append(
|
||||
Thread(
|
||||
target=producer,
|
||||
args=(work_q, control_q, filename, _logger, stop_event, positions, start_offsets),
|
||||
)
|
||||
Thread(target=producer, args=(work_q, control_q, filename, _logger))
|
||||
)
|
||||
for worker in threads:
|
||||
worker.start()
|
||||
for worker in threads:
|
||||
worker.join()
|
||||
interrupted = bool(stop_event and stop_event.is_set())
|
||||
return result_list, positions, interrupted
|
||||
return result_list
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -155,13 +155,7 @@ interactive.persistent("/srv/betoniarka/data/script.params")
|
||||
# Configure output formats and destinations
|
||||
|
||||
output.icecast(%mp3, host="localhost", port=8000, password=credentials.password, icy_metadata="true", mount="mp3-stream", radio)
|
||||
# Local monitor output. fallible=true so a dead or missing pulse daemon degrades
|
||||
# to "no monitor" instead of failing its clock and taking the whole radio down
|
||||
# with it ("Shutdown started!"). The stream that actually matters is the Icecast
|
||||
# one above, which needs no sound device at all.
|
||||
# NOTE: input.pulseaudio() (the mic, further up) is still a HARD dependency - on
|
||||
# a headless container with no capture device, comment BOTH of them out.
|
||||
output.pulseaudio(fallible=true, radio)
|
||||
output.pulseaudio(radio)
|
||||
#output.file.hls("/tmp/hls", [("mp3-low", %mp3(bitrate=96)), ("mp3-hi", %mp3(bitrate=160))], radio)
|
||||
# Uncomment the following lines to enable additional output formats
|
||||
# output.icecast(%opus, host="retropie", port=8000, password=credentials.password, icy_metadata="true", mount="opus-stream", radio)
|
||||
|
||||
+1
-76
@@ -93,9 +93,6 @@ SEND_QUERY = "/query"
|
||||
# Health-check round-trip: a pseudo-query that the librarian must pull off its
|
||||
# own internal queue and answer (same uuid) WITHOUT running a real search.
|
||||
LIBRARIAN_PING = "/ping"
|
||||
# Per-query watchdog probe: "do you still know this uuid?" (queued/processing).
|
||||
# A uuid that vanishes here without its result reaching the bot was lost.
|
||||
QUERY_STATUS = "/query_status"
|
||||
TIME_BETWEEN_CALLS = 100000
|
||||
LAST_SPONTANEOUS_CALL = datetime.now()
|
||||
|
||||
@@ -230,31 +227,6 @@ MEMORY_COMPACT_THRESHOLD = int(os.getenv("CONJURER_MEMORY_COMPACT_THRESHOLD", "4
|
||||
MEMORY_KEEP_RECENT = int(os.getenv("CONJURER_MEMORY_KEEP_RECENT", "200"))
|
||||
MEMORY_COMPACT_HOURS = float(os.getenv("CONJURER_MEMORY_COMPACT_HOURS", "6"))
|
||||
|
||||
# Durable result-delivery spool (librarian -> bot). The bot persists every
|
||||
# incoming search result to RESULT_INBOX_DIR before acking and only forgets it
|
||||
# once rendered (uuid recorded in DELIVERED_DIR), so an expensive (hours-long)
|
||||
# result survives a bot restart mid-flight and duplicate resends are idempotent.
|
||||
# Rooted under CONJURER_DATA_DIR when set (a mounted volume), else next to the
|
||||
# log file. DELIVERED_MAX bounds the remembered-uuid set.
|
||||
_STATE_ROOT = _DATA_DIR or (os.path.dirname(LOGFILE) or ".")
|
||||
RESULT_INBOX_DIR = os.getenv(
|
||||
"CONJURER_RESULT_INBOX", os.path.join(_STATE_ROOT, "result_inbox")
|
||||
)
|
||||
DELIVERED_DIR = os.getenv(
|
||||
"CONJURER_DELIVERED_DIR", os.path.join(_STATE_ROOT, "delivered_uuids")
|
||||
)
|
||||
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
|
||||
@@ -264,10 +236,6 @@ SKIP_TRACK = os.getenv("CONJURER_SKIP_ENDPOINT", "/skip")
|
||||
LIBRARIAN_SERVICE_ADDRESS = os.getenv(
|
||||
"CONJURER_LIBRARIAN_SERVICE", "http://192.168.1.192:5001"
|
||||
)
|
||||
# The address the librarian (and its pongs) should send results BACK to for THIS
|
||||
# bot - so one librarian can serve several bots (test + deploy), each getting its
|
||||
# own answers. Empty => the librarian falls back to its static CONJURER_MAIN_BOT.
|
||||
SELF_CALLBACK = os.getenv("CONJURER_SELF_CALLBACK", "")
|
||||
HOST_ADDRESS = os.getenv("CONJURER_DISCORD_HOST", "192.168.1.191")
|
||||
PORT_ADDRESS = int(os.getenv("CONJURER_DISCORD_PORT", "5000"))
|
||||
|
||||
@@ -410,38 +378,6 @@ if anthropic and ANTHROPIC_API_KEY:
|
||||
else:
|
||||
CLAUDECLIENT = None
|
||||
|
||||
# Ollama (self-hosted models). There is no API key - the endpoint IS the whole
|
||||
# configuration, so the feature stays dormant until CONJURER_OLLAMA_URL is set
|
||||
# (same pattern as the Conan bridge). We talk to Ollama's OpenAI-COMPATIBLE
|
||||
# surface (/v1) with the openai SDK we already depend on, which means the
|
||||
# existing message format and _map_openai_error handling work unchanged.
|
||||
# How long handle_response waits for ANY backend before giving up. 120s was
|
||||
# hardcoded and is fine for hosted APIs, but a self-hosted model on a modest GPU
|
||||
# can legitimately take longer, so it is now tunable.
|
||||
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:
|
||||
# api_key is required by the SDK but ignored by Ollama.
|
||||
OLLAMACLIENT = openai.AsyncOpenAI(base_url=f"{OLLAMA_URL}/v1", api_key="ollama")
|
||||
else:
|
||||
OLLAMACLIENT = None
|
||||
|
||||
TOKEN = _resolve_token("discord", "DISCORD_TOKEN")
|
||||
|
||||
# Voice recognition (AssemblyAI). None = the voice cog reports and disables.
|
||||
@@ -531,12 +467,6 @@ def _default_ai_configs():
|
||||
# sent for Claude (Opus 4.8 / Sonnet 5 reject sampling params).
|
||||
"max_tokens": 2048,
|
||||
},
|
||||
"ollama": {
|
||||
"provider": "ollama",
|
||||
"latest_model": OLLAMA_LATEST_MODEL,
|
||||
"cheap_model": OLLAMA_CHEAP_MODEL,
|
||||
"temperature": 0.2,
|
||||
},
|
||||
# Template for wiring further providers. Copy it, rename the key, point
|
||||
# "provider" at a backend ai_functions.provider_generate implements, and
|
||||
# fill in the model ids. Keys starting with "_" are treated as inert
|
||||
@@ -556,12 +486,7 @@ _ai_block = (
|
||||
if isinstance(GPT_SETTINGS, list) and len(GPT_SETTINGS) > 2 and isinstance(GPT_SETTINGS[2], dict)
|
||||
else {}
|
||||
)
|
||||
# Built-in defaults FIRST, then whatever the settings file defines on top. The
|
||||
# file cannot simply win outright: every provider switch persists a "configs"
|
||||
# block, so a file written by an older build would permanently hide providers
|
||||
# added later (ollama) from the picker.
|
||||
AI_CONFIGS = _default_ai_configs()
|
||||
AI_CONFIGS.update(_ai_block.get("configs") or {})
|
||||
AI_CONFIGS = _ai_block.get("configs") or _default_ai_configs()
|
||||
# Single switch: env var wins, then the settings-file "active" key, then "gpt".
|
||||
DEFAULT_AI_CONFIG = (
|
||||
os.getenv("CONJURER_AI_CONFIG")
|
||||
|
||||
@@ -14,9 +14,6 @@ RUN pip install --no-cache-dir --upgrade pip \
|
||||
&& pip install --no-cache-dir -r requirements_librarian.txt requests
|
||||
|
||||
COPY conjurer_librarian/ ./
|
||||
# durable_queue lives at the repo root and is shared with the bot; the librarian
|
||||
# imports it for the durable result OUTBOX.
|
||||
COPY durable_queue.py ./
|
||||
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
CONJURER_LIBRARIAN_HOST=0.0.0.0 \
|
||||
|
||||
@@ -11,16 +11,6 @@ CREDS="$SECRETS/icecast_credentials.json"
|
||||
mkdir -p "$DATA" "$MUSIC" "$SECRETS"
|
||||
|
||||
# Seed the script + persistent interactive params from the image on first run.
|
||||
# NOTE: the live script is deliberately never overwritten, so hand edits win -
|
||||
# but that also means image fixes NEVER reach a volume seeded long ago. Set
|
||||
# RADIO_FORCE_SCRIPT=1 to take the image's version (the old one is kept as
|
||||
# radio_conjurer.liq.bak so nothing hand-written is lost).
|
||||
if [ -e "$DATA/radio_conjurer.liq" ] && [ "${RADIO_FORCE_SCRIPT:-0}" = "1" ]; then
|
||||
cp "$DATA/radio_conjurer.liq" "$DATA/radio_conjurer.liq.bak"
|
||||
cp /app/radio_conjurer.liq "$DATA/"
|
||||
echo "RADIO_FORCE_SCRIPT=1: reseeded radio_conjurer.liq from the image" >&2
|
||||
echo " (previous version saved as radio_conjurer.liq.bak)" >&2
|
||||
fi
|
||||
[ -e "$DATA/radio_conjurer.liq" ] || cp /app/radio_conjurer.liq "$DATA/"
|
||||
if [ ! -e "$DATA/script.params" ]; then
|
||||
if [ -e /app/script.params ]; then cp /app/script.params "$DATA/"; else : > "$DATA/script.params"; fi
|
||||
@@ -83,36 +73,14 @@ fi
|
||||
# none - you edited the script to drop pulse in/out.
|
||||
case "${PULSE_MODE:-internal}" in
|
||||
internal)
|
||||
# Clear stale runtime state FIRST. `docker restart` - and the crash-loop
|
||||
# that restart:unless-stopped produces - reuses the container's writable
|
||||
# layer, so /run/pulse/pid left by a killed daemon survives and the next
|
||||
# start dies with "Daemon startup failed"; that kills liquidsoap, which
|
||||
# restarts the container, forever. Removing the pid/socket of a daemon
|
||||
# that is demonstrably not running breaks the loop.
|
||||
if ! pidof pulseaudio >/dev/null 2>&1; then
|
||||
rm -f /run/pulse/pid /var/run/pulse/pid \
|
||||
/run/pulse/native /var/run/pulse/native 2>/dev/null || true
|
||||
fi
|
||||
# --disallow-module-loading: modules from system.pa still load at
|
||||
# startup; this only blocks later client-requested loads (and
|
||||
# silences the system-mode warning). The "forcibly disabling SHM"
|
||||
# notice is inherent to system mode and harmless.
|
||||
if pulseaudio --system --daemonize=yes --disallow-exit \
|
||||
--disallow-module-loading --exit-idle-time=-1; then
|
||||
pulseaudio --system --daemonize=yes --disallow-exit \
|
||||
--disallow-module-loading --exit-idle-time=-1 \
|
||||
|| echo "WARNING: internal pulseaudio failed to start" >&2
|
||||
export PULSE_SERVER="${PULSE_SERVER:-unix:/var/run/pulse/native}"
|
||||
else
|
||||
# Be loud: with pulse dead, input.pulseaudio()/output.pulseaudio()
|
||||
# fail to start, liquidsoap tears down the whole clock ("Shutdown
|
||||
# started!") and the container crash-loops. The stream itself only
|
||||
# needs Icecast, so the way out is dropping the pulse tor.
|
||||
echo "ERROR: internal pulseaudio failed to start." >&2
|
||||
echo " Liquidsoap will crash-loop while the script still uses" >&2
|
||||
echo " input.pulseaudio()/output.pulseaudio(). The Icecast" >&2
|
||||
echo " output does NOT need pulse: comment those out in" >&2
|
||||
echo " $DATA/radio_conjurer.liq (or set RADIO_FORCE_SCRIPT=1" >&2
|
||||
echo " to re-seed the script from the image) and restart." >&2
|
||||
echo " Diagnose with: pulseaudio --system --daemonize=no -vvvv" >&2
|
||||
fi
|
||||
;;
|
||||
host)
|
||||
[ -n "$PULSE_SERVER" ] || echo "WARNING: PULSE_MODE=host but PULSE_SERVER is unset" >&2
|
||||
|
||||
Vendored
+2
-21
@@ -13,29 +13,10 @@ CONJURER_NETRC_FILE=/secrets/.netrc
|
||||
|
||||
# --- AI backend switch --------------------------------------------------
|
||||
# Which AI config from system_gpt_settings.json is active at startup
|
||||
# (e.g. "gpt", "claude" or "ollama"). Runtime switch:
|
||||
# $gadaj_teraz <config> [model]. Unset = whatever the settings file's "active"
|
||||
# key says, falling back to "gpt".
|
||||
# (e.g. "gpt" or "claude"). Runtime switch: $gadaj_teraz <config>. Unset =
|
||||
# whatever the settings file's "active" key says, falling back to "gpt".
|
||||
# CONJURER_AI_CONFIG=gpt
|
||||
|
||||
# --- Ollama (self-hosted models) ----------------------------------------
|
||||
# The endpoint IS the whole configuration - no API key. Leave unset and the
|
||||
# "ollama" backend simply refuses to be selected. In-cluster, use the Service
|
||||
# DNS name; from outside, host:port. Port 11434 is Ollama's default.
|
||||
# CONJURER_OLLAMA_URL=http://ollama.ollama.svc.cluster.local:11434
|
||||
# CONJURER_OLLAMA_URL=
|
||||
# DEFAULT model for normal replies. $modele_ai lists what the server actually
|
||||
# has pulled. Pinning one at runtime with `$gadaj_teraz ollama <model>` is
|
||||
# persisted into system_gpt_settings.json and from then on WINS over this
|
||||
# variable - the pin is the more recent, more explicit choice. Clear the
|
||||
# "latest_model" of the ollama entry in that file to fall back to this default.
|
||||
# CONJURER_OLLAMA_MODEL=llama3.1:8b
|
||||
# Model used for the cheaper MUSIC path; defaults to CONJURER_OLLAMA_MODEL.
|
||||
# CONJURER_OLLAMA_CHEAP_MODEL=
|
||||
# How long to wait for ANY backend to answer. 120s suits hosted APIs; a
|
||||
# self-hosted model on a modest GPU may need more.
|
||||
# CONJURER_AI_TIMEOUT_SECONDS=120
|
||||
|
||||
# --- Data ---------------------------------------------------------------
|
||||
# Single mounted volume; all writable state is rooted here.
|
||||
CONJURER_DATA_DIR=/data
|
||||
|
||||
@@ -10,9 +10,8 @@ Runbook for running Conjurer as Docker containers across Proxmox VMs:
|
||||
| **Share** (optional) | VM-musician or its own | `conjurer-share` | 8081 → 80 | `docker/Dockerfile.share` |
|
||||
|
||||
The share service publishes short-lived file links over Apache and is documented
|
||||
separately: [SHARE_NODE_SETUP.md](SHARE_NODE_SETUP.md) to stand a node up from
|
||||
zero, [FILE_SHARING.md](FILE_SHARING.md) for how the feature works. It shares two
|
||||
volumes with the musician, so set it up after the musician is running.
|
||||
separately in [FILE_SHARING.md](FILE_SHARING.md) — it shares two volumes with the
|
||||
musician, so set it up after the musician is running.
|
||||
|
||||
The three talk to each other over HTTP on the Proxmox LAN. Direction of calls:
|
||||
|
||||
|
||||
@@ -7,10 +7,6 @@ Historically only the Python half of this lived in the repo; the Apache config
|
||||
and the cron entries were placed on the host by hand. This document plus
|
||||
`docker/Dockerfile.share` close that gap.
|
||||
|
||||
> Standing it up on a fresh node (directories, permissions, reverse proxy,
|
||||
> verification)? Start with **[SHARE_NODE_SETUP.md](SHARE_NODE_SETUP.md)** — this
|
||||
> document explains how the feature works and how it is secured.
|
||||
|
||||
## What it actually does
|
||||
|
||||
1. A **scanner** walks the media library and writes a JSON index of every path.
|
||||
|
||||
@@ -1,149 +0,0 @@
|
||||
# Share node setup — from zero
|
||||
|
||||
Everything the share service needs is **inside the image**: Apache, the vhost,
|
||||
the index scanner and the link revoker. The node needs no Apache, no cron, no
|
||||
Python and no copies of the old hand-placed scripts.
|
||||
|
||||
What the node actually provides is four directories, the media library, and a
|
||||
way in from the internet.
|
||||
|
||||
> Feature docs (how links work, security model, TTL semantics) live in
|
||||
> [FILE_SHARING.md](FILE_SHARING.md). This file is only "how to stand it up on a
|
||||
> fresh box".
|
||||
|
||||
## What runs where
|
||||
|
||||
| Piece | Where it lives | Notes |
|
||||
|---|---|---|
|
||||
| Apache + vhost | in the image | rendered from `share-vhost.conf.tpl` at start |
|
||||
| `scan_shares.py` (index) | in the image | sleep loop, not cron |
|
||||
| `revoke_shares.py` (expiry) | in the image | sleep loop, not cron |
|
||||
| the symlinks | host `/srv/share/links` | **created by the musician**, served here |
|
||||
| the index | host `/srv/share/db` | written here, **read by the musician** |
|
||||
| TLS / public name | your reverse proxy | container speaks plain HTTP |
|
||||
|
||||
The single most important fact: **the musician creates the links, this service
|
||||
serves them.** They must see the same directories, at the same paths.
|
||||
|
||||
## 1. Directories
|
||||
|
||||
```bash
|
||||
sudo mkdir -p /srv/share/{media,links,db,logs}
|
||||
```
|
||||
|
||||
| Path | Contents | Mounted as |
|
||||
|---|---|---|
|
||||
| `/srv/share/media` | the media library | `/mnt/shares` (read-only) |
|
||||
| `/srv/share/links` | published symlinks | `/var/www/html/share` |
|
||||
| `/srv/share/db` | `share_scan.json` | `/srv/share/db` |
|
||||
| `/srv/share/logs` | Apache logs incl. `share_access.log` | `/var/log/apache2` |
|
||||
|
||||
## 2. Media library and permissions
|
||||
|
||||
Put the library at `/srv/share/media` (bind mount, NFS mount, whatever — it is
|
||||
only ever read). Apache serves as **`www-data`, uid 33 inside the container**, so
|
||||
that uid must be able to traverse and read it:
|
||||
|
||||
```bash
|
||||
sudo chmod -R o+rX /srv/share/media # simplest; or use ACLs/group instead
|
||||
sudo -u '#33' test -r /srv/share/media/<some-file> && echo "readable by www-data"
|
||||
```
|
||||
|
||||
A library that root can read but uid 33 cannot is the classic "every link
|
||||
404s / 403s" cause.
|
||||
|
||||
## 3. Deploy the stack
|
||||
|
||||
The musician and the share service share a filesystem, so the supported layout
|
||||
is **both on the same node**, from one stack:
|
||||
|
||||
- Portainer → **Stacks → Add stack** → paste `docker/compose.musician-share.stack.yaml`
|
||||
- or CLI: `docker compose -f docker/compose.musician-share.stack.yaml up -d`
|
||||
|
||||
That stack already wires the four mounts on both containers. Set
|
||||
`CONJURER_SHARE_SERVER_NAME` to your public hostname.
|
||||
|
||||
Share-only node (musician elsewhere): use `docker/compose.share.yaml` and put
|
||||
`/srv/share/links` + `/srv/share/db` on storage **both** hosts mount — otherwise
|
||||
the musician cheerfully creates links this container cannot see.
|
||||
|
||||
## 4. Way in from the internet
|
||||
|
||||
The container listens on **8081 → 80**, plain HTTP by design; TLS stays on the
|
||||
reverse proxy you already run. Forward the `/share/` prefix **unchanged** — the
|
||||
links are `https://<host>/share/<token>`:
|
||||
|
||||
nginx:
|
||||
```nginx
|
||||
location /share/ {
|
||||
proxy_pass http://<share-node-ip>:8081/share/;
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
}
|
||||
```
|
||||
|
||||
Apache (as reverse proxy):
|
||||
```apache
|
||||
ProxyPass /share/ http://<share-node-ip>:8081/share/
|
||||
ProxyPassReverse /share/ http://<share-node-ip>:8081/share/
|
||||
```
|
||||
|
||||
Then: DNS for the public name points at the proxy, and the node's firewall lets
|
||||
the proxy reach 8081 (nothing else needs to).
|
||||
|
||||
## 5. Tell the musician where to publish
|
||||
|
||||
The musician builds the URLs it posts to Discord. In its env:
|
||||
|
||||
```ini
|
||||
CONJURER_SHARE_DIR=/var/www/html/share
|
||||
CONJURER_SHARE_DB=/srv/share/db/share_scan.json
|
||||
CONJURER_SHARE_BASE_URL=https://czernobog.pl/share
|
||||
```
|
||||
|
||||
`CONJURER_SHARE_BASE_URL` is **not** set in the bundled stack file — it falls back
|
||||
to `https://czernobog.pl/share`. If your public name differs, set it explicitly
|
||||
or every posted link points at the wrong host.
|
||||
|
||||
## The three couplings that break it
|
||||
|
||||
1. **Same container path for the media.** The index records absolute paths and
|
||||
the symlinks are absolute. Both containers must mount the library at
|
||||
`/mnt/shares`. Mount it elsewhere on one side and every link dangles.
|
||||
2. **Same link dir and index dir** for musician and share (same host, or shared
|
||||
storage).
|
||||
3. **`CONJURER_SHARE_BASE_URL` must equal your real public `/share` URL.**
|
||||
|
||||
## Verification
|
||||
|
||||
```bash
|
||||
# index built (should be non-trivial JSON)
|
||||
sudo head -c 200 /srv/share/db/share_scan.json; echo
|
||||
|
||||
# the two jobs and Apache are alive
|
||||
docker logs conjurer-share | tail -20 # "[share] serving ... scan every ...s"
|
||||
|
||||
# directory listing MUST fail (403) - it would leak every live token
|
||||
curl -sI http://<share-node-ip>:8081/share/ | head -1
|
||||
|
||||
# revoker state files must NOT be served
|
||||
curl -sI http://<share-node-ip>:8081/share/.downloads.json | head -1
|
||||
|
||||
# a real link: publish one from Discord, then
|
||||
curl -sI https://<public-host>/share/<token> | head -1 # 200
|
||||
```
|
||||
|
||||
After the first download of a link, `revoke_shares.py` removes it once
|
||||
`CONJURER_SHARE_TTL_SECONDS` (default 1h) has passed. A link nobody downloads is
|
||||
never revoked by that job.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
| Symptom | Cause | Fix |
|
||||
|---|---|---|
|
||||
| every link 404 | media mounted at a different container path than `/mnt/shares`, or the symlink target is gone | align the mounts on both containers |
|
||||
| every link 403 | media not readable by uid 33 | step 2 |
|
||||
| links created but not served | musician and share not sharing `/srv/share/links` | step 3 |
|
||||
| `/get_share_list` empty | index missing/not shared | check `/srv/share/db/share_scan.json` and the musician's `CONJURER_SHARE_DB` |
|
||||
| links never expire | revoker cannot see the access log, or nobody downloaded them | check `/srv/share/logs/share_access.log` exists and grows |
|
||||
| wrong host in posted links | `CONJURER_SHARE_BASE_URL` | step 5 |
|
||||
@@ -1,112 +0,0 @@
|
||||
"""Dependency-free, disk-backed queue for durable message delivery.
|
||||
|
||||
One JSON file per key under a directory. Used on both sides of the
|
||||
librarian <-> bot result path so an expensive (hours-long) search result is
|
||||
never lost to a transient network failure or a restart:
|
||||
|
||||
* the librarian keeps a result in its OUTBOX until the bot acks it,
|
||||
* the bot keeps a result in its INBOX until it is actually rendered, and
|
||||
remembers delivered uuids so duplicate resends are idempotent.
|
||||
|
||||
Only ``json`` + ``os`` are imported, so the logic is unit-testable without
|
||||
flask, discord, or the network. Writes are atomic (temp file + ``os.replace``)
|
||||
so a crash mid-write can never leave a half-written record that poisons replay.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
|
||||
def _safe_name(key: str) -> str:
|
||||
"""Filesystem-safe file stem for a key (uuids are safe; be defensive)."""
|
||||
stem = "".join(c for c in str(key) if c.isalnum() or c in "-_.")
|
||||
return stem or "_"
|
||||
|
||||
|
||||
class DiskQueue:
|
||||
"""A directory of ``<key>.json`` records, each ``{key, payload, ts}``."""
|
||||
|
||||
def __init__(self, directory: str):
|
||||
# No disk touch here on purpose: constructing a DiskQueue at import time
|
||||
# must not create directories (tests, read-only default paths). The
|
||||
# directory is created lazily on the first put().
|
||||
self.directory = directory
|
||||
|
||||
def _path(self, key) -> str:
|
||||
return os.path.join(self.directory, _safe_name(key) + ".json")
|
||||
|
||||
def put(self, key, payload) -> None:
|
||||
"""Atomically write (overwrite) the record for ``key``."""
|
||||
os.makedirs(self.directory, exist_ok=True)
|
||||
record = {"key": str(key), "payload": payload, "ts": time.time()}
|
||||
fd, tmp = tempfile.mkstemp(dir=self.directory, suffix=".tmp")
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
json.dump(record, handle)
|
||||
os.replace(tmp, self._path(key)) # atomic on POSIX
|
||||
finally:
|
||||
if os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
def contains(self, key) -> bool:
|
||||
return os.path.exists(self._path(key))
|
||||
|
||||
def get(self, key):
|
||||
"""Return the payload stored for ``key``, or None if absent/unreadable."""
|
||||
try:
|
||||
with open(self._path(key), encoding="utf-8") as handle:
|
||||
return json.load(handle)["payload"]
|
||||
except (OSError, ValueError, KeyError, TypeError):
|
||||
return None
|
||||
|
||||
def remove(self, key) -> None:
|
||||
try:
|
||||
os.remove(self._path(key))
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
def items(self):
|
||||
"""Return ``[(key, payload, ts), ...]`` oldest-first.
|
||||
|
||||
Unreadable / half-written / corrupt files are skipped (never raise),
|
||||
so one bad file can't stall replay of the rest.
|
||||
"""
|
||||
out = []
|
||||
try:
|
||||
names = os.listdir(self.directory)
|
||||
except FileNotFoundError:
|
||||
return out
|
||||
for name in names:
|
||||
if not name.endswith(".json"):
|
||||
continue
|
||||
try:
|
||||
with open(os.path.join(self.directory, name), encoding="utf-8") as handle:
|
||||
record = json.load(handle)
|
||||
out.append((record["key"], record["payload"], record.get("ts", 0)))
|
||||
except (OSError, ValueError, KeyError, TypeError):
|
||||
continue
|
||||
out.sort(key=lambda triple: triple[2])
|
||||
return out
|
||||
|
||||
def keys(self):
|
||||
return [key for key, _payload, _ts in self.items()]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items())
|
||||
|
||||
def prune(self, max_entries: int) -> int:
|
||||
"""Keep only the newest ``max_entries`` (by ts); drop the rest.
|
||||
|
||||
Used for the delivered-uuid set so it cannot grow without bound.
|
||||
Returns how many were dropped.
|
||||
"""
|
||||
if max_entries < 0:
|
||||
return 0
|
||||
entries = self.items() # oldest first
|
||||
excess = len(entries) - max_entries
|
||||
dropped = 0
|
||||
for key, _payload, _ts in entries[: max(0, excess)]:
|
||||
self.remove(key)
|
||||
dropped += 1
|
||||
return dropped
|
||||
+4
-129
@@ -3,7 +3,6 @@ import io
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from queue import Empty
|
||||
|
||||
@@ -15,48 +14,16 @@ import requests
|
||||
from discord.ext import commands, tasks
|
||||
|
||||
from ai_functions import handle_response
|
||||
from communication_subroutine import (
|
||||
IN_COMM_Q,
|
||||
OUT_COMM_Q,
|
||||
QueryControl,
|
||||
mark_delivered,
|
||||
submit_ai_query,
|
||||
)
|
||||
from constants import (
|
||||
DIR_PATH_SADOX,
|
||||
LIBRARIAN_SERVICE_ADDRESS,
|
||||
QUERY_STATUS,
|
||||
SELF_CALLBACK,
|
||||
SEND_QUERY,
|
||||
service_headers,
|
||||
)
|
||||
from librarian_watchdog import FLAG, pending_verdict
|
||||
from communication_subroutine import IN_COMM_Q, OUT_COMM_Q, QueryControl, submit_ai_query
|
||||
from constants import DIR_PATH_SADOX, LIBRARIAN_SERVICE_ADDRESS, SEND_QUERY, service_headers
|
||||
|
||||
SERVICE_HEADERS = service_headers()
|
||||
|
||||
# Per-query watchdog tuning.
|
||||
PENDING_WATCH_SECONDS = 30 # how often to ask the librarian about a uuid
|
||||
PENDING_GRACE_SECONDS = 45 # unknown-but-pending must persist this long
|
||||
PENDING_HARD_TTL = 60 * 60 * 24 * 3 # drop tracking after 3 days no matter what
|
||||
|
||||
|
||||
class DataModule(commands.Cog):
|
||||
def __init__(self, bot, logger_name):
|
||||
self.bot = bot
|
||||
self.logger = logging.getLogger(logger_name)
|
||||
# uuid -> {"ctx", "query", "created", "unknown_since"} for every search
|
||||
# dispatched but not yet answered. watch_pending polls the librarian for
|
||||
# each; check_data_q removes an entry the moment its result is rendered.
|
||||
self.pending = {}
|
||||
|
||||
def _track_pending(self, query_uuid, query, ctx):
|
||||
"""Start watching a dispatched search so a lost result can be caught."""
|
||||
self.pending[str(query_uuid)] = {
|
||||
"ctx": ctx,
|
||||
"query": query,
|
||||
"created": time.monotonic(),
|
||||
"unknown_since": None,
|
||||
}
|
||||
|
||||
@commands.hybrid_command(
|
||||
nsfw=True,
|
||||
@@ -85,13 +52,8 @@ class DataModule(commands.Cog):
|
||||
# check if current path is a file
|
||||
if os.path.isfile(os.path.join(DIR_PATH_SADOX, path)):
|
||||
res.append(path)
|
||||
if not res:
|
||||
await ctx.send("*Conjurer grzebie w pustej skrzyni* Nie ma dziś żadnych komiksów.")
|
||||
return
|
||||
# randrange(len) is 0..len-1; the old randrange(0, len-1) never picked
|
||||
# the last file and raised ValueError('empty range') on a single file.
|
||||
# trunk-ignore(bandit/B311)
|
||||
filename = res[random.randrange(len(res))]
|
||||
filename = res[random.randrange(0, len(res) - 1)]
|
||||
# select random page
|
||||
file = open(DIR_PATH_SADOX + filename, "rb")
|
||||
if True:
|
||||
@@ -133,9 +95,6 @@ class DataModule(commands.Cog):
|
||||
fresh_data = IN_COMM_Q.get(block=False)
|
||||
entries = []
|
||||
if fresh_data.stop:
|
||||
# The result arrived and is about to be rendered - stop the
|
||||
# watchdog from ever flagging this uuid as lost.
|
||||
self.pending.pop(str(fresh_data.uuid), None)
|
||||
searcher = fresh_data.author
|
||||
query = fresh_data.content
|
||||
# ai_lines is a clean, plain rendering of the SAME list in the
|
||||
@@ -147,7 +106,7 @@ class DataModule(commands.Cog):
|
||||
desc = fresh_data.entries[doi]
|
||||
title = desc["Title"][0] if desc.get("Title") else "(bez tytułu)"
|
||||
entries.append(
|
||||
f"{l_p}. {title} pod linkiem https://www.sci-hub.red/{doi} i jest to {desc['type']}\n"
|
||||
f"{l_p}. {title} pod linkiem https://www.sci-hub.se/{doi} i jest to {desc['type']}\n"
|
||||
)
|
||||
ai_lines.append(f"{l_p}. {title} (DOI: {doi}, typ: {desc['type']})")
|
||||
l_p += 1
|
||||
@@ -176,12 +135,6 @@ class DataModule(commands.Cog):
|
||||
await ctx.send(message)
|
||||
message = ""
|
||||
|
||||
# The result is now on screen: mark it delivered so the
|
||||
# librarian's resends become no-ops and it is dropped from the
|
||||
# durable inbox (never replayed again). Done after the core
|
||||
# render but before the optional AI review, which is a bonus.
|
||||
mark_delivered(str(fresh_data.uuid))
|
||||
|
||||
# Optional AI pass: re-rank the (already Crossref-relevance-
|
||||
# sorted) DOI list and review the sources. Enqueued to the AI
|
||||
# worker so it runs on whatever backend $gadaj_teraz selected;
|
||||
@@ -217,77 +170,6 @@ class DataModule(commands.Cog):
|
||||
except Empty:
|
||||
pass
|
||||
|
||||
@tasks.loop(seconds=PENDING_WATCH_SECONDS)
|
||||
async def watch_pending(self):
|
||||
"""Per-query safety net for lost results (case a).
|
||||
|
||||
For each dispatched-but-unanswered search, ask the librarian whether it
|
||||
still knows the uuid (queued or processing). While it does, the search is
|
||||
progressing - leave it alone (a busy librarian is fine). The moment a
|
||||
uuid VANISHES on the librarian while still pending here, its result was
|
||||
computed but never reached us: after a short grace window (to rule out a
|
||||
result that is merely in flight) we tell the channel - but ONLY then.
|
||||
|
||||
A normally-delivered result is popped from self.pending by check_data_q,
|
||||
so it never reaches the flag path.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
for query_uuid in list(self.pending.keys()):
|
||||
info = self.pending.get(query_uuid)
|
||||
if info is None:
|
||||
continue
|
||||
# Hard cap so a permanently-unreachable librarian can't leak entries.
|
||||
if now - info["created"] > PENDING_HARD_TTL:
|
||||
self.logger.warning("Dropping stale pending query %s (hard TTL)", query_uuid)
|
||||
self.pending.pop(query_uuid, None)
|
||||
continue
|
||||
try:
|
||||
response = await asyncio.to_thread(
|
||||
requests.post,
|
||||
f"{LIBRARIAN_SERVICE_ADDRESS}{QUERY_STATUS}",
|
||||
json={"UUID": query_uuid},
|
||||
headers=SERVICE_HEADERS,
|
||||
timeout=5,
|
||||
)
|
||||
known = (
|
||||
response.status_code == 200
|
||||
and response.json().get("data", {}).get("known", False)
|
||||
)
|
||||
except (
|
||||
requests.exceptions.RequestException,
|
||||
ValueError,
|
||||
AttributeError,
|
||||
KeyError,
|
||||
TypeError,
|
||||
) as exc:
|
||||
# Librarian unreachable / garbled or unexpected answer: we can't
|
||||
# judge, so don't cry wolf, and don't let one bad poll kill the
|
||||
# loop. Reset the clock and try again next tick.
|
||||
self.logger.info("Pending check for %s inconclusive: %s", query_uuid, exc)
|
||||
info["unknown_since"] = None
|
||||
continue
|
||||
action, info["unknown_since"] = pending_verdict(
|
||||
known, info["unknown_since"], now, PENDING_GRACE_SECONDS
|
||||
)
|
||||
# Re-check membership: the await above yields, so check_data_q may
|
||||
# have just delivered (and popped) this result.
|
||||
if action == FLAG and query_uuid in self.pending:
|
||||
await self._flag_lost(query_uuid, info)
|
||||
self.pending.pop(query_uuid, None)
|
||||
|
||||
async def _flag_lost(self, query_uuid, info):
|
||||
"""Tell the querent their finished search never made it back."""
|
||||
message = (
|
||||
"*Winda na książki z hukiem wraca z podziemi PUSTA. Z głośnika trzeszczy:* "
|
||||
f"Twoje zapytanie {query_uuid} (\"{info['query']}\") przemieliło się w "
|
||||
"bibliotece do końca, ale wynik przepadł gdzieś w drodze do baru - nic nie "
|
||||
"dotarło. Zawołaj szefa albo puść jeszcze raz."
|
||||
)
|
||||
try:
|
||||
await info["ctx"].send(message)
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
self.logger.exception("Failed to post lost-result notice for %s", query_uuid)
|
||||
|
||||
@commands.hybrid_command(
|
||||
name="wyszukaj_linki_do_dokumentow",
|
||||
description="Szuka linkow doi w bazie crossref i podaje linki do scihuba",
|
||||
@@ -314,7 +196,6 @@ class DataModule(commands.Cog):
|
||||
"query": str(query),
|
||||
"page": 1,
|
||||
"deep_search": False,
|
||||
"callback": SELF_CALLBACK,
|
||||
}
|
||||
coroutine = asyncio.to_thread(
|
||||
requests.post,
|
||||
@@ -347,7 +228,6 @@ class DataModule(commands.Cog):
|
||||
username = ctx.message.author.name
|
||||
query_object = QueryControl(username, query_uuid, query, ctx)
|
||||
OUT_COMM_Q.put(query_object)
|
||||
self._track_pending(query_uuid, query, ctx)
|
||||
await ctx.send(
|
||||
f"No dobra poszło. Zapamiętaj proszę {query_uuid} - to identyfikator twojego zapytania. Jesteś {queue_size} w kolejce."
|
||||
+ " Zapytania obsługuje algorytm zasilany czterema chomikami zapierdalającymi w kołowrotku - więc wyniki najwcześniej za kilka godzi - ale mogą być też dni."
|
||||
@@ -374,7 +254,6 @@ class DataModule(commands.Cog):
|
||||
"query": str(query),
|
||||
"page": 1,
|
||||
"deep_search": False,
|
||||
"callback": SELF_CALLBACK,
|
||||
}
|
||||
coroutine = asyncio.to_thread(
|
||||
requests.post,
|
||||
@@ -407,7 +286,6 @@ class DataModule(commands.Cog):
|
||||
username = ctx.message.author.name
|
||||
query_object = QueryControl(username, query_uuid, query, ctx, ai_review=True)
|
||||
OUT_COMM_Q.put(query_object)
|
||||
self._track_pending(query_uuid, query, ctx)
|
||||
await ctx.send(
|
||||
f"Poszło z recenzją AI. Identyfikator: {query_uuid}. Jesteś {queue_size} w kolejce."
|
||||
+ " Najpierw dojadą surowe wyniki, a zaraz po nich przesortowanie i recenzja od AI."
|
||||
@@ -472,7 +350,6 @@ class DataModule(commands.Cog):
|
||||
"query": str(query),
|
||||
"page": 1,
|
||||
"deep_search": True,
|
||||
"callback": SELF_CALLBACK,
|
||||
}
|
||||
coroutine = asyncio.to_thread(
|
||||
requests.post,
|
||||
@@ -505,7 +382,6 @@ class DataModule(commands.Cog):
|
||||
username = ctx.message.author.name
|
||||
query_object = QueryControl(username, query_uuid, query, ctx)
|
||||
OUT_COMM_Q.put(query_object)
|
||||
self._track_pending(query_uuid, query, ctx)
|
||||
await ctx.send(
|
||||
f"Wypij wypi {query_uuid} - to identyfikator twojego zapytania. Jesteś {queue_size} w kolejce niestety kolejka obowiazuje zawsze."
|
||||
+ " Zapytania obsługuje algorytm zasilany czterema chomikami zapierdalającymi w kołowrotku - więc wyniki najwcześniej za kilka godzin - ale mogą być też dni. Głębokie zabawy trwają dłużej. *Znowu mruga*"
|
||||
@@ -516,6 +392,5 @@ async def setup(bot):
|
||||
logger = logging.getLogger("discord")
|
||||
dm = DataModule(bot, "discord")
|
||||
dm.check_data_q.start()
|
||||
dm.watch_pending.start()
|
||||
await bot.add_cog(dm)
|
||||
logger.info("Loading data sharing commands module done")
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Pure decision logic for the librarian per-query watchdog.
|
||||
|
||||
Split out of ``librarian_commands`` (which pulls in discord / pdf libs, so it is
|
||||
not importable in the pytest-only unit job) so the one subtle part - the grace
|
||||
window that stops a just-delivered result from being falsely flagged as lost -
|
||||
can be unit-tested in isolation.
|
||||
|
||||
The watchdog polls the librarian's /query_status for each dispatched query:
|
||||
|
||||
* the librarian reports the uuid ``known`` while it is queued or processing,
|
||||
* once the search finishes it is dropped there, so the uuid goes ``unknown``.
|
||||
|
||||
A result that arrives normally is removed from the pending set by the result
|
||||
handler, so the watchdog never even sees it. Only a uuid that goes ``unknown``
|
||||
on the librarian *and is still pending on the bot* is a lost result - but we
|
||||
require it to stay that way for a grace window first, because there is always a
|
||||
brief moment where the librarian has finished (uuid gone) yet the result is
|
||||
still in flight / not yet rendered.
|
||||
"""
|
||||
|
||||
WAIT = "wait"
|
||||
FLAG = "flag"
|
||||
|
||||
|
||||
def pending_verdict(known, unknown_since, now, grace_seconds):
|
||||
"""Decide what to do this tick for one pending query.
|
||||
|
||||
Args:
|
||||
known: did the librarian report the uuid as still known this tick?
|
||||
unknown_since: monotonic timestamp the uuid was first seen unknown, or
|
||||
None if it was known last tick.
|
||||
now: current monotonic time.
|
||||
grace_seconds: how long a uuid must stay unknown-but-pending before it
|
||||
is declared lost.
|
||||
|
||||
Returns:
|
||||
(action, unknown_since) where action is WAIT or FLAG and the returned
|
||||
``unknown_since`` is what the caller should store for the next tick.
|
||||
"""
|
||||
if known:
|
||||
# Still queued/processing (or freshly back to known) - reset the clock.
|
||||
return WAIT, None
|
||||
if unknown_since is None:
|
||||
# First tick we see it gone: start the grace clock, don't flag yet - the
|
||||
# result may simply be in flight.
|
||||
return WAIT, now
|
||||
if now - unknown_since >= grace_seconds:
|
||||
# Gone for the whole grace window and still pending: the result was lost.
|
||||
return FLAG, unknown_since
|
||||
# Gone, but not long enough yet - keep waiting.
|
||||
return WAIT, unknown_since
|
||||
@@ -1,70 +0,0 @@
|
||||
"""Integration: the librarian's simple result cache.
|
||||
|
||||
A repeat of the same query (normalised) returns stored hits and skips the whole
|
||||
Crossref + DB scan. Nothing fancy: whitespace/case-insensitive exact match,
|
||||
disk-backed, TTL'd, size-bounded, deep/shallow kept separate.
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
from durable_queue import DiskQueue # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cache(tmp_path, monkeypatch):
|
||||
box = DiskQueue(str(tmp_path / "cache"))
|
||||
monkeypatch.setattr(lib, "_cache", box)
|
||||
monkeypatch.setattr(lib, "CACHE_TTL_SECONDS", 3600)
|
||||
monkeypatch.setattr(lib, "CACHE_MAX_ENTRIES", 500)
|
||||
return box
|
||||
|
||||
|
||||
_HITS = {"10.1000/x": {"Title": ["A Paper"], "type": "journal-article"}}
|
||||
|
||||
|
||||
def test_put_then_get_is_a_hit(cache):
|
||||
lib._cache_put("kwas foliowy", False, _HITS)
|
||||
assert lib._cache_get("kwas foliowy", False) == _HITS
|
||||
|
||||
|
||||
def test_different_query_misses(cache):
|
||||
lib._cache_put("kwas foliowy", False, _HITS)
|
||||
assert lib._cache_get("witamina c", False) is None
|
||||
|
||||
|
||||
def test_normalised_case_and_whitespace_hit_same_entry(cache):
|
||||
lib._cache_put(" Kwas Foliowy ", False, _HITS)
|
||||
assert lib._cache_get("kwas foliowy", False) == _HITS
|
||||
|
||||
|
||||
def test_deep_and_shallow_are_cached_separately(cache):
|
||||
lib._cache_put("q", False, _HITS)
|
||||
assert lib._cache_get("q", True) is None # a deep search is a different key
|
||||
assert lib._cache_get("q", False) == _HITS
|
||||
|
||||
|
||||
def test_expired_entry_is_a_miss(cache):
|
||||
cache.put(lib._cache_key("q", False), {"query": "q", "final_result": _HITS, "expires": 0})
|
||||
assert lib._cache_get("q", False) is None
|
||||
|
||||
|
||||
def test_ttl_zero_disables_the_cache(cache, monkeypatch):
|
||||
monkeypatch.setattr(lib, "CACHE_TTL_SECONDS", 0)
|
||||
lib._cache_put("q", False, _HITS) # no-op when disabled
|
||||
assert len(cache) == 0
|
||||
assert lib._cache_get("q", False) is None
|
||||
|
||||
|
||||
def test_prune_bounds_the_cache(cache, monkeypatch):
|
||||
monkeypatch.setattr(lib, "CACHE_MAX_ENTRIES", 3)
|
||||
for i in range(6):
|
||||
lib._cache_put(f"query-{i}", False, _HITS)
|
||||
assert len(cache) <= 3
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Integration: the librarian's Crossref-contact resolution.
|
||||
|
||||
CONJURER_CROSSREF_MAILTO alone is a valid, complete configuration. A missing
|
||||
netrc must NOT produce a "credentials missing" warning in that case - the old
|
||||
code warned on every single search even though the env var was set and used.
|
||||
Only a genuine absence of any contact should warn (and then raise).
|
||||
"""
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
|
||||
|
||||
class _DummyCrossref:
|
||||
"""Accepts the kwargs the real Crossref does, so Librarian() can construct."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _crossref_and_missing_netrc(monkeypatch):
|
||||
# Build with a harmless Crossref, and force the netrc read to miss so the
|
||||
# env-var path is what's exercised.
|
||||
monkeypatch.setattr(lib, "Crossref", _DummyCrossref)
|
||||
monkeypatch.setattr(lib, "NETRC_FILE", "/nonexistent/conjurer/.netrc")
|
||||
|
||||
|
||||
def test_env_mailto_alone_does_not_warn(monkeypatch, caplog):
|
||||
monkeypatch.setenv("CONJURER_CROSSREF_MAILTO", "mtuszowski@example.com")
|
||||
with caplog.at_level(logging.WARNING, logger="conjurer_librarian"):
|
||||
librarian = lib.Librarian(lib.app, "kwas foliowy", "uuid-1", False)
|
||||
assert librarian.uuid == "uuid-1" # constructed fine
|
||||
assert not any(
|
||||
"credentials missing" in r.getMessage().lower()
|
||||
or "not configured" in r.getMessage().lower()
|
||||
for r in caplog.records
|
||||
), "a missing netrc must not warn when CONJURER_CROSSREF_MAILTO is set"
|
||||
|
||||
|
||||
def test_no_contact_anywhere_raises(monkeypatch):
|
||||
monkeypatch.delenv("CONJURER_CROSSREF_MAILTO", raising=False)
|
||||
with pytest.raises(RuntimeError):
|
||||
lib.Librarian(lib.app, "kwas foliowy", "uuid-2", False)
|
||||
@@ -1,102 +0,0 @@
|
||||
"""Integration: transient Crossref failures must not destroy a search.
|
||||
|
||||
Field report: a single httpx ReadTimeout inside habanero surfaced as
|
||||
"Search <uuid> crashed", and the worker then FORGOT the search - so an
|
||||
expensive query vanished and the user got told it was eaten, all because a
|
||||
public API blinked. These pin the two defences: retry each Crossref call, and
|
||||
retry the whole search a bounded number of times before giving up.
|
||||
"""
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
from durable_queue import DiskQueue # noqa: E402
|
||||
|
||||
_LOG = logging.getLogger("test-crossref-retry")
|
||||
_LOG.addHandler(logging.NullHandler())
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_backoff(monkeypatch):
|
||||
monkeypatch.setattr(lib.time, "sleep", lambda _s: None)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def state(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(lib, "_requests", DiskQueue(str(tmp_path / "req")))
|
||||
monkeypatch.setattr(lib, "_checkpoints", DiskQueue(str(tmp_path / "cp")))
|
||||
return None
|
||||
|
||||
|
||||
def test_crossref_call_retries_then_succeeds():
|
||||
calls = {"n": 0}
|
||||
|
||||
def flaky(**_kwargs):
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 3:
|
||||
raise RuntimeError("The read operation timed out")
|
||||
return {"message": {"total-results": 1, "items": []}}
|
||||
|
||||
result = lib._crossref_call(_LOG, "works", flaky, query="q")
|
||||
assert result["message"]["total-results"] == 1
|
||||
assert calls["n"] == 3 # two failures survived
|
||||
|
||||
|
||||
def test_crossref_call_reraises_after_exhausting_attempts(monkeypatch):
|
||||
monkeypatch.setattr(lib, "CROSSREF_ATTEMPTS", 2)
|
||||
calls = {"n": 0}
|
||||
|
||||
def always_fails(**_kwargs):
|
||||
calls["n"] += 1
|
||||
raise RuntimeError("The read operation timed out")
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
lib._crossref_call(_LOG, "works", always_fails, query="q")
|
||||
assert calls["n"] == 2 # bounded, not infinite
|
||||
|
||||
|
||||
def test_crossref_call_does_not_retry_a_success():
|
||||
calls = {"n": 0}
|
||||
|
||||
def ok(**_kwargs):
|
||||
calls["n"] += 1
|
||||
return "fine"
|
||||
|
||||
assert lib._crossref_call(_LOG, "works", ok, query="q") == "fine"
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_attempt_counter_persists_and_bounds_retries(state):
|
||||
# Mirrors what the worker does on a crash: bump the persisted attempt count
|
||||
# and keep the request until SEARCH_MAX_ATTEMPTS is reached.
|
||||
uuid = "u-crash"
|
||||
lib._requests.put(uuid, {"query": "q", "deep_search": False, "callback": ""})
|
||||
|
||||
for expected in (1, 2):
|
||||
stored = lib._requests.get(uuid) or {}
|
||||
attempts = int(stored.get("attempts", 0)) + 1
|
||||
assert attempts == expected
|
||||
stored["attempts"] = attempts
|
||||
lib._requests.put(uuid, stored)
|
||||
|
||||
assert lib._requests.get(uuid)["attempts"] == 2
|
||||
# A third crash reaches the default cap (3) -> the search is forgotten.
|
||||
assert 3 >= lib.SEARCH_MAX_ATTEMPTS
|
||||
lib._forget_search(uuid)
|
||||
assert not lib._requests.contains(uuid)
|
||||
|
||||
|
||||
def test_forget_search_clears_request_and_checkpoint(state):
|
||||
lib._requests.put("u-x", {"query": "q", "deep_search": False})
|
||||
lib._checkpoints.put("u-x", {"dois": {}, "found": [], "positions": {}})
|
||||
lib._forget_search("u-x")
|
||||
assert not lib._requests.contains("u-x")
|
||||
assert not lib._checkpoints.contains("u-x")
|
||||
@@ -1,108 +0,0 @@
|
||||
"""Integration: the 'still searching' heartbeat and its cheap progress estimate.
|
||||
|
||||
The estimate must stay free: producers already record a byte offset per chunk
|
||||
file and the total is stat()'d once, so a reading is just a sum over ~40 ints.
|
||||
"""
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
import search_bot # noqa: E402
|
||||
|
||||
_LOG = logging.getLogger("test-heartbeat")
|
||||
_LOG.addHandler(logging.NullHandler())
|
||||
|
||||
|
||||
def test_progress_summary_percentages():
|
||||
progress = {"positions": {"0_chunk.txt": 250, "1_chunk.txt": 250}, "total_bytes": 1000}
|
||||
done, total, percent = lib._progress_summary(progress)
|
||||
assert (done, total) == (500, 1000)
|
||||
assert percent == pytest.approx(50.0)
|
||||
|
||||
|
||||
def test_progress_summary_unknown_total_is_zero_percent():
|
||||
done, total, percent = lib._progress_summary({"positions": {"a": 10}})
|
||||
assert (done, total, percent) == (10, 0, 0.0)
|
||||
|
||||
|
||||
def test_progress_summary_handles_empty_and_none():
|
||||
assert lib._progress_summary(None) == (0, 0, 0.0)
|
||||
assert lib._progress_summary({}) == (0, 0, 0.0)
|
||||
|
||||
|
||||
def test_progress_summary_is_clamped_to_100():
|
||||
# A partially-buffered tail can push the summed offsets past the total.
|
||||
_done, _total, percent = lib._progress_summary(
|
||||
{"positions": {"a": 1500}, "total_bytes": 1000}
|
||||
)
|
||||
assert percent == pytest.approx(100.0)
|
||||
|
||||
|
||||
def test_current_search_registration_round_trip():
|
||||
progress = {"positions": {"a": 5}, "total_bytes": 10}
|
||||
live = [{"DOI": "10.1/x"}]
|
||||
lib._set_current_search("uuid-1", "kwas foliowy", progress, live)
|
||||
with lib._current_lock:
|
||||
snapshot = dict(lib._current_search)
|
||||
assert snapshot["uuid"] == "uuid-1"
|
||||
assert snapshot["query"] == "kwas foliowy"
|
||||
assert lib._progress_summary(snapshot["progress"])[2] == pytest.approx(50.0)
|
||||
lib._clear_current_search()
|
||||
with lib._current_lock:
|
||||
assert not lib._current_search
|
||||
|
||||
|
||||
def _write_two_chunks(tmp_path):
|
||||
(tmp_path / "0_chunk.txt").write_text("10.1/a\n10.1/b\n", encoding="utf-8")
|
||||
(tmp_path / "1_chunk.txt").write_text("10.1/c\n", encoding="utf-8")
|
||||
return sum(
|
||||
(tmp_path / name).stat().st_size for name in ("0_chunk.txt", "1_chunk.txt")
|
||||
)
|
||||
|
||||
|
||||
def test_search_fills_progress_and_reaches_full_coverage(tmp_path, monkeypatch):
|
||||
# Coverage must be measured on a search that CANNOT stop early. Once every
|
||||
# queried DOI is found the consumer signals TERM and the producers stop
|
||||
# mid-file, so a search for a DOI that exists reaches an arbitrary offset -
|
||||
# asserting 100% there is a race (it failed roughly one run in two).
|
||||
# An absent DOI forces the whole database to be read.
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
expected_total = _write_two_chunks(tmp_path)
|
||||
|
||||
progress = {}
|
||||
search_bot.search_for_doi([("10.9/absent", "DATA")], [], _LOG, progress=progress)
|
||||
|
||||
assert progress["total_bytes"] == expected_total
|
||||
assert progress["chunk_files"] == 2
|
||||
done, total, percent = lib._progress_summary(progress)
|
||||
assert total == expected_total
|
||||
assert done == expected_total # nothing stopped it: whole DB scanned
|
||||
assert percent == pytest.approx(100.0)
|
||||
|
||||
|
||||
def test_progress_is_populated_for_a_search_that_finds_its_target(tmp_path, monkeypatch):
|
||||
# The early-termination case: the target is found, so coverage is whatever
|
||||
# the producers reached. Assert what IS deterministic - the total is known,
|
||||
# progress is bounded and sane, and the hit is reported.
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
expected_total = _write_two_chunks(tmp_path)
|
||||
|
||||
progress = {}
|
||||
result, _positions, _interrupted = search_bot.search_for_doi(
|
||||
[("10.1/c", "DATA")], [], _LOG, progress=progress
|
||||
)
|
||||
|
||||
assert progress["total_bytes"] == expected_total
|
||||
done, total, percent = lib._progress_summary(progress)
|
||||
assert total == expected_total
|
||||
assert 0 <= done <= total # bounded, never nonsense
|
||||
assert 0.0 <= percent <= 100.0
|
||||
assert [r for r in result if r["DOI"] == "10.1/c" and r["exists"]]
|
||||
@@ -1,127 +0,0 @@
|
||||
"""Integration: the librarian's durable result OUTBOX + retrying delivery.
|
||||
|
||||
An 8-hour search result must not be lost to a transient bot outage. The result
|
||||
is written to the OUTBOX before sending; delivery retries with backoff; the
|
||||
entry is removed only on a positive ACK; and the resender keeps flushing the
|
||||
OUTBOX (across restarts, since it is on the persistent state volume).
|
||||
"""
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
# conjurer_librarian imports `from habanero import Crossref` at import time; the
|
||||
# integration job doesn't install habanero. Stub it (we never build a real
|
||||
# Librarian here).
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
from durable_queue import DiskQueue # noqa: E402
|
||||
|
||||
_LOG = logging.getLogger("test-outbox")
|
||||
_LOG.addHandler(logging.NullHandler())
|
||||
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, status_code, text=""):
|
||||
self.status_code = status_code
|
||||
self.text = text
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_backoff(monkeypatch):
|
||||
# Never actually sleep during retry backoff in tests.
|
||||
monkeypatch.setattr(lib.time, "sleep", lambda _s: None)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def outbox(tmp_path, monkeypatch):
|
||||
box = DiskQueue(str(tmp_path / "outbox"))
|
||||
monkeypatch.setattr(lib, "_outbox", box)
|
||||
return box
|
||||
|
||||
|
||||
def test_deliver_succeeds_first_try(monkeypatch):
|
||||
urls = []
|
||||
monkeypatch.setattr(lib.requests, "post", lambda url, **k: urls.append(url) or _Resp(200))
|
||||
assert lib._deliver_result("http://bot-a:5000", "u1", {"u1": {}}, _LOG, attempts=3) is True
|
||||
assert len(urls) == 1 # no needless retries after a 200
|
||||
assert urls[0] == "http://bot-a:5000" + lib.SEND_RESULTS # to the origin bot
|
||||
|
||||
|
||||
def test_deliver_falls_back_to_main_bot_when_no_target(monkeypatch):
|
||||
urls = []
|
||||
monkeypatch.setattr(lib.requests, "post", lambda url, **k: urls.append(url) or _Resp(200))
|
||||
assert lib._deliver_result("", "u1b", {"u1b": {}}, _LOG, attempts=1) is True
|
||||
assert urls[0] == lib.MAIN_BOT_ADDRESS + lib.SEND_RESULTS # empty target -> default
|
||||
|
||||
|
||||
def test_deliver_retries_then_succeeds(monkeypatch):
|
||||
responses = iter([_Resp(503), _Resp(500), _Resp(200)])
|
||||
monkeypatch.setattr(lib.requests, "post", lambda *a, **k: next(responses))
|
||||
assert lib._deliver_result("http://bot", "u2", {"u2": {}}, _LOG, attempts=3) is True
|
||||
|
||||
|
||||
def test_deliver_returns_false_when_all_attempts_fail(monkeypatch):
|
||||
def boom(*_a, **_k):
|
||||
raise lib.requests.exceptions.RequestException("bot down")
|
||||
|
||||
monkeypatch.setattr(lib.requests, "post", boom)
|
||||
assert lib._deliver_result("http://bot", "u3", {"u3": {}}, _LOG, attempts=2) is False
|
||||
|
||||
|
||||
def test_resend_delivers_each_result_to_its_own_origin_bot(outbox, monkeypatch):
|
||||
outbox.put("ok", {"target": "http://bot-a:5000", "payload": {"ok": {}}})
|
||||
outbox.put("bad", {"target": "http://bot-b:5000", "payload": {"bad": {}}})
|
||||
seen = []
|
||||
|
||||
def fake_deliver(target, query_uuid, _payload, _logger, attempts=1):
|
||||
seen.append((target, query_uuid))
|
||||
return query_uuid == "ok"
|
||||
|
||||
monkeypatch.setattr(lib, "_deliver_result", fake_deliver)
|
||||
lib._resend_once(_LOG)
|
||||
|
||||
assert ("http://bot-a:5000", "ok") in seen # delivered to A's address
|
||||
assert ("http://bot-b:5000", "bad") in seen # attempted to B's address
|
||||
assert not outbox.contains("ok") # acked -> dropped
|
||||
assert outbox.contains("bad") # not acked -> kept for the next sweep
|
||||
|
||||
|
||||
def test_resend_handles_legacy_entry_shape(outbox, monkeypatch):
|
||||
# An OUTBOX entry from before per-origin callbacks (raw payload, no target)
|
||||
# must still be delivered - to the default bot.
|
||||
outbox.put("old", {"old": {"10.1/x": {"Title": ["P"], "type": "a"}}})
|
||||
seen = []
|
||||
|
||||
def fake_deliver(target, query_uuid, _payload, _logger, attempts=1):
|
||||
seen.append((target, query_uuid))
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(lib, "_deliver_result", fake_deliver)
|
||||
lib._resend_once(_LOG)
|
||||
assert seen == [("", "old")] # empty target -> _deliver_result uses MAIN_BOT
|
||||
assert not outbox.contains("old")
|
||||
|
||||
|
||||
def test_resend_keeps_result_until_bot_recovers(outbox, monkeypatch):
|
||||
# Simulate: bot down for the first sweep, up for the second. The result must
|
||||
# survive the outage and be delivered on recovery.
|
||||
outbox.put("u9", {"target": "http://bot", "payload": {"u9": {"10.1/x": {"Title": ["P"], "type": "article"}}}})
|
||||
state = {"up": False}
|
||||
|
||||
def flaky_post(*_a, **_k):
|
||||
return _Resp(200) if state["up"] else _Resp(502)
|
||||
|
||||
monkeypatch.setattr(lib.requests, "post", flaky_post)
|
||||
|
||||
lib._resend_once(_LOG) # bot down
|
||||
assert outbox.contains("u9") # preserved, not lost
|
||||
|
||||
state["up"] = True
|
||||
lib._resend_once(_LOG) # bot recovered
|
||||
assert not outbox.contains("u9") # now delivered and cleared
|
||||
@@ -1,125 +0,0 @@
|
||||
"""Integration: the librarian's health/liveness surface.
|
||||
|
||||
Two behaviours, both proven against the real Flask app:
|
||||
|
||||
* /ping is busy-aware - while a search is grinding it pongs back immediately
|
||||
WITHOUT queueing (busy is healthy); when idle it routes the ping through the
|
||||
internal queue for the worker to answer.
|
||||
* /query_status reports whether a uuid is still known (queued/processing), which
|
||||
is what the bot's per-query watchdog polls to catch a lost result.
|
||||
|
||||
Only the SYNC routes (/ping, /query_status) are exercised - the async /query
|
||||
route needs flask[async], which the integration job doesn't install, so query
|
||||
state is seeded directly on the module.
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
# conjurer_librarian does `from habanero import Crossref` at import time and
|
||||
# habanero isn't installed in the integration job. Stub it before importing the
|
||||
# service (we never build a real Librarian here, so Crossref is just a name).
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
|
||||
|
||||
def _client(key=None):
|
||||
lib.API_KEY = key
|
||||
return lib.app.test_client()
|
||||
|
||||
|
||||
def _reset():
|
||||
with lib._active_lock:
|
||||
lib.active_queries.clear()
|
||||
lib.worker_busy.clear()
|
||||
while not lib.librarian_queue.empty():
|
||||
lib.librarian_queue.get()
|
||||
|
||||
|
||||
def test_query_status_known_vs_unknown():
|
||||
_reset()
|
||||
client = _client()
|
||||
with lib._active_lock:
|
||||
lib.active_queries["abc"] = "queued"
|
||||
|
||||
known = client.post("/query_status", json={"UUID": "abc"}).get_json()["data"]
|
||||
assert known == {"uuid": "abc", "known": True, "state": "queued"}
|
||||
|
||||
unknown = client.post("/query_status", json={"UUID": "nope"}).get_json()["data"]
|
||||
assert unknown == {"uuid": "nope", "known": False, "state": "unknown"}
|
||||
|
||||
|
||||
def test_ping_idle_routes_through_internal_queue(monkeypatch):
|
||||
_reset()
|
||||
posted = []
|
||||
monkeypatch.setattr(lib.requests, "post", lambda *a, **k: posted.append((a, k)))
|
||||
client = _client()
|
||||
|
||||
resp = client.post("/ping", json={"UUID": "ping-idle", "callback": "http://bot-a:5000"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
# Idle => it went onto the internal queue for the worker (carrying the
|
||||
# callback so the worker pongs the right bot), NOT posted directly.
|
||||
assert posted == []
|
||||
assert lib.librarian_queue.get_nowait() == {
|
||||
"__ping__": "ping-idle",
|
||||
"callback": "http://bot-a:5000",
|
||||
}
|
||||
|
||||
|
||||
def test_ping_while_busy_pongs_directly_to_the_pinging_bot(monkeypatch):
|
||||
_reset()
|
||||
lib.worker_busy.set() # a search is grinding
|
||||
posted = []
|
||||
|
||||
class _Resp:
|
||||
status_code = 200
|
||||
|
||||
def fake_post(url, json=None, headers=None, timeout=None):
|
||||
posted.append({"url": url, "json": json})
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(lib.requests, "post", fake_post)
|
||||
client = _client()
|
||||
|
||||
resp = client.post("/ping", json={"UUID": "ping-busy", "callback": "http://bot-b:5000"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
# Busy => direct pong, NOTHING queued, and it goes to the CALLBACK bot (not
|
||||
# the static default) so a shared librarian health-checks each bot correctly.
|
||||
assert lib.librarian_queue.empty()
|
||||
assert len(posted) == 1
|
||||
assert posted[0]["json"] == {"__pong__": "ping-busy"}
|
||||
assert posted[0]["url"] == "http://bot-b:5000" + lib.SEND_RESULTS
|
||||
|
||||
|
||||
def test_ping_without_callback_pongs_to_default_bot(monkeypatch):
|
||||
_reset()
|
||||
lib.worker_busy.set()
|
||||
posted = []
|
||||
|
||||
class _Resp:
|
||||
status_code = 200
|
||||
|
||||
monkeypatch.setattr(
|
||||
lib.requests, "post",
|
||||
lambda url, **k: posted.append(url) or _Resp(),
|
||||
)
|
||||
client = _client()
|
||||
client.post("/ping", json={"UUID": "ping-nocb"}) # no callback
|
||||
|
||||
assert posted == [lib.MAIN_BOT_ADDRESS + lib.SEND_RESULTS]
|
||||
|
||||
|
||||
def test_query_status_enforces_api_key():
|
||||
_reset()
|
||||
client = _client(key="secret")
|
||||
denied = client.post("/query_status", json={"UUID": "x"})
|
||||
assert denied.status_code == 401
|
||||
ok = client.post(
|
||||
"/query_status", json={"UUID": "x"}, headers={"X-Conjurer-Api-Key": "secret"}
|
||||
)
|
||||
assert ok.status_code == 200
|
||||
@@ -1,90 +0,0 @@
|
||||
"""Integration: the librarian's persisted search state (requests + checkpoints).
|
||||
|
||||
These pin the durable-state mechanics that let a search survive a restart:
|
||||
* accepted requests are replayed (re-enqueued) after a restart,
|
||||
* a finished/abandoned search is forgotten (request + checkpoint dropped),
|
||||
* a checkpoint round-trips through disk intact.
|
||||
|
||||
The RESUME correctness itself (seek past scanned, don't miss, don't re-scan)
|
||||
lives in tests/unit/test_search_bot.py.
|
||||
"""
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "habanero" not in sys.modules:
|
||||
_habanero = types.ModuleType("habanero")
|
||||
_habanero.Crossref = object
|
||||
sys.modules["habanero"] = _habanero
|
||||
|
||||
import conjurer_librarian as lib # noqa: E402
|
||||
from durable_queue import DiskQueue # noqa: E402
|
||||
|
||||
_LOG = logging.getLogger("test-resume-state")
|
||||
_LOG.addHandler(logging.NullHandler())
|
||||
|
||||
|
||||
class _DummyCrossref:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def state(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(lib, "_requests", DiskQueue(str(tmp_path / "req")))
|
||||
monkeypatch.setattr(lib, "_checkpoints", DiskQueue(str(tmp_path / "cp")))
|
||||
monkeypatch.setattr(lib, "Crossref", _DummyCrossref)
|
||||
monkeypatch.setenv("CONJURER_CROSSREF_MAILTO", "test@example.com")
|
||||
monkeypatch.setattr(lib, "NETRC_FILE", "/nonexistent/conjurer/.netrc")
|
||||
while not lib.librarian_queue.empty():
|
||||
lib.librarian_queue.get()
|
||||
lib.librarian_list.clear()
|
||||
with lib._active_lock:
|
||||
lib.active_queries.clear()
|
||||
return None
|
||||
|
||||
|
||||
def test_accepted_request_is_replayed_after_restart(state):
|
||||
lib._requests.put("u1", {"query": "kwas foliowy", "deep_search": False})
|
||||
|
||||
lib.replay_requests(_LOG)
|
||||
|
||||
item = lib.librarian_queue.get_nowait()
|
||||
assert isinstance(item, lib.Librarian)
|
||||
assert item.uuid == "u1"
|
||||
assert item.query == "kwas foliowy"
|
||||
assert lib.active_queries["u1"] == "queued" # known again to the watchdog
|
||||
|
||||
|
||||
def test_forget_search_drops_request_and_checkpoint(state):
|
||||
lib._requests.put("u2", {"query": "x", "deep_search": False})
|
||||
lib._checkpoints.put("u2", {"dois": {}, "found": [], "positions": {}})
|
||||
|
||||
lib._forget_search("u2")
|
||||
|
||||
assert not lib._requests.contains("u2")
|
||||
assert not lib._checkpoints.contains("u2")
|
||||
|
||||
|
||||
def test_checkpoint_round_trips_through_disk(state):
|
||||
checkpoint = {
|
||||
"dois": {"10.1/x": {"DOI": "10.1/x", "title": ["T"], "type": "article"}},
|
||||
"found": ["10.1/already"],
|
||||
"positions": {"0_chunk.txt": 4096},
|
||||
}
|
||||
lib._checkpoints.put("u3", checkpoint)
|
||||
assert lib._checkpoints.get("u3") == checkpoint
|
||||
|
||||
|
||||
def test_unreadable_replayed_request_is_dropped_not_looped(state, monkeypatch):
|
||||
# A request that can't be reconstructed (e.g. missing Crossref contact) must
|
||||
# be dropped, not retried forever.
|
||||
lib._requests.put("u4", {"query": "x", "deep_search": False})
|
||||
monkeypatch.delenv("CONJURER_CROSSREF_MAILTO", raising=False)
|
||||
|
||||
lib.replay_requests(_LOG)
|
||||
|
||||
assert lib.librarian_queue.empty()
|
||||
assert not lib._requests.contains("u4") # forgotten, not left to loop
|
||||
@@ -1,130 +0,0 @@
|
||||
"""Integration: the librarian -> bot RESULT delivery contract.
|
||||
|
||||
A search that 'vanishes' (watchdog fires "zeżarło") means the result never
|
||||
reached the bot's inbound queue. These tests pin down the contract so we can
|
||||
tell a CODE break (wrong shape / uuid / auth handling) from a TRANSPORT break
|
||||
(the librarian can't reach the bot at all - wrong address/port). They prove the
|
||||
bot side is correct end to end, which isolates a systematic vanish to transport.
|
||||
|
||||
The result the librarian sends is exactly:
|
||||
{uuid: {DOI: {"Title": [<title>...], "type": <str>}}}
|
||||
(see conjurer_librarian.answer_query -> final_result, POSTed to /conjurer).
|
||||
"""
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import communication_subroutine as cs
|
||||
from durable_queue import DiskQueue
|
||||
|
||||
|
||||
def _drain(queue):
|
||||
while not queue.empty():
|
||||
queue.get()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def comm_threads(tmp_path, monkeypatch):
|
||||
# Point the durable spool at a temp dir so /conjurer's dedup/persist can't
|
||||
# leak into (or be poisoned by) the real result_inbox/ between runs.
|
||||
monkeypatch.setattr(cs, "_inbox", DiskQueue(str(tmp_path / "inbox")))
|
||||
monkeypatch.setattr(cs, "_delivered", DiskQueue(str(tmp_path / "delivered")))
|
||||
cs.awaiting_q.clear()
|
||||
_drain(cs.incoming_q)
|
||||
_drain(cs.OUT_COMM_Q)
|
||||
_drain(cs.IN_COMM_Q)
|
||||
cs.API_KEY = None
|
||||
stop = threading.Event()
|
||||
workers = [
|
||||
threading.Thread(target=cs.scan_queue, kwargs={"stop_event": stop}, daemon=True),
|
||||
threading.Thread(target=cs.scan_incoming, kwargs={"stop_event": stop}, daemon=True),
|
||||
]
|
||||
for worker in workers:
|
||||
worker.start()
|
||||
yield
|
||||
stop.set()
|
||||
for worker in workers:
|
||||
worker.join(timeout=3)
|
||||
|
||||
|
||||
def _dispatch(uuid, query="kwas foliowy"):
|
||||
"""Mimic the bot dispatching a search: a QueryControl enters the comm queue
|
||||
and scan_queue moves it into awaiting_q."""
|
||||
qc = cs.QueryControl("siara", uuid, query, None)
|
||||
cs.OUT_COMM_Q.put(qc)
|
||||
deadline = time.time() + 2
|
||||
while time.time() < deadline:
|
||||
if any(getattr(r, "uuid", None) == uuid for r in list(cs.awaiting_q)):
|
||||
return qc
|
||||
time.sleep(0.01)
|
||||
raise AssertionError("scan_queue never moved the query into awaiting_q")
|
||||
|
||||
|
||||
# The exact result the librarian emits for one found DOI.
|
||||
def _result_payload(uuid):
|
||||
return {uuid: {"10.1000/xyz": {"Title": ["A Real Paper"], "type": "journal-article"}}}
|
||||
|
||||
|
||||
def test_librarian_result_reaches_bot_when_transport_is_fine(comm_threads):
|
||||
_dispatch("uuid-ok")
|
||||
client = cs.app.test_client()
|
||||
|
||||
resp = client.post("/conjurer", json=_result_payload("uuid-ok"))
|
||||
|
||||
assert resp.status_code == 200
|
||||
got = cs.IN_COMM_Q.get(timeout=3)
|
||||
assert got.uuid == "uuid-ok"
|
||||
assert got.stop is True
|
||||
# Exactly the shape check_data_q renders: entries[DOI]["Title"][0] / ["type"].
|
||||
assert got.entries == {"10.1000/xyz": {"Title": ["A Real Paper"], "type": "journal-article"}}
|
||||
|
||||
|
||||
def test_empty_result_is_still_delivered_not_vanished(comm_threads):
|
||||
# A search that found nothing sends {uuid: {}} - it must STILL be delivered
|
||||
# (renders "niestety nie ma nic"), never look like a lost result.
|
||||
_dispatch("uuid-empty")
|
||||
client = cs.app.test_client()
|
||||
|
||||
resp = client.post("/conjurer", json={"uuid-empty": {}})
|
||||
|
||||
assert resp.status_code == 200
|
||||
got = cs.IN_COMM_Q.get(timeout=3)
|
||||
assert got.uuid == "uuid-empty"
|
||||
assert got.entries == {}
|
||||
|
||||
|
||||
def test_wrong_api_key_rejects_result_so_it_vanishes(comm_threads):
|
||||
# (b) reproduction: if the librarian's CONJURER_API_KEY differs from the
|
||||
# bot's, /conjurer returns 401 and the result is never queued - the search
|
||||
# silently vanishes exactly as reported.
|
||||
cs.API_KEY = "bot-secret"
|
||||
_dispatch("uuid-auth")
|
||||
client = cs.app.test_client()
|
||||
|
||||
resp = client.post(
|
||||
"/conjurer",
|
||||
json=_result_payload("uuid-auth"),
|
||||
headers={"X-Conjurer-Api-Key": "librarian-DIFFERENT-key"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 401
|
||||
time.sleep(0.4)
|
||||
assert cs.IN_COMM_Q.empty() # nothing delivered
|
||||
|
||||
|
||||
def test_uuid_mismatch_orphans_result_away_from_the_querent(comm_threads):
|
||||
# (b) reproduction: if the uuid the librarian echoes back doesn't byte-match
|
||||
# what the bot stored, scan_incoming can't match it -> it goes to the orphan
|
||||
# path (posted to the fallback channel, NOT the querent) and the querent's
|
||||
# pending entry is never cleared, so the watchdog still flags it lost.
|
||||
_dispatch("uuid-stored")
|
||||
client = cs.app.test_client()
|
||||
|
||||
client.post("/conjurer", json=_result_payload("uuid-DIFFERENT"))
|
||||
|
||||
got = cs.IN_COMM_Q.get(timeout=3)
|
||||
assert got.author == "Orphaned"
|
||||
assert got.uuid == "uuid-DIFFERENT"
|
||||
# The original querent's record is untouched (still awaiting) - it "vanished".
|
||||
assert any(getattr(r, "uuid", None) == "uuid-stored" for r in list(cs.awaiting_q))
|
||||
@@ -1,92 +0,0 @@
|
||||
"""Integration: the bot's DURABLE, idempotent result intake (/conjurer).
|
||||
|
||||
The expensive-result guarantees on the bot side:
|
||||
* every result is persisted to the inbox before it is acked,
|
||||
* a resend of a not-yet-delivered result is dropped (no double render),
|
||||
* once rendered (mark_delivered) further resends are dropped and it leaves the
|
||||
inbox,
|
||||
* on startup, an accepted-but-unrendered result is replayed from the inbox,
|
||||
* health-check pongs are never persisted.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import communication_subroutine as cs
|
||||
from durable_queue import DiskQueue
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def spool(tmp_path, monkeypatch):
|
||||
inbox = DiskQueue(str(tmp_path / "inbox"))
|
||||
delivered = DiskQueue(str(tmp_path / "delivered"))
|
||||
monkeypatch.setattr(cs, "_inbox", inbox)
|
||||
monkeypatch.setattr(cs, "_delivered", delivered)
|
||||
cs.API_KEY = None
|
||||
while not cs.incoming_q.empty():
|
||||
cs.incoming_q.get()
|
||||
return inbox, delivered
|
||||
|
||||
|
||||
def _drain_incoming():
|
||||
out = []
|
||||
while not cs.incoming_q.empty():
|
||||
out.append(cs.incoming_q.get())
|
||||
return out
|
||||
|
||||
|
||||
def _payload(uuid):
|
||||
return {uuid: {"10.1000/xyz": {"Title": ["A Real Paper"], "type": "journal-article"}}}
|
||||
|
||||
|
||||
def test_result_is_persisted_then_queued(spool):
|
||||
inbox, _delivered = spool
|
||||
resp = cs.app.test_client().post("/conjurer", json=_payload("u1"))
|
||||
assert resp.status_code == 200
|
||||
assert inbox.contains("u1") # durable before ack
|
||||
assert _drain_incoming() == [_payload("u1")]
|
||||
|
||||
|
||||
def test_resend_while_pending_is_not_requeued(spool):
|
||||
client = cs.app.test_client()
|
||||
client.post("/conjurer", json=_payload("u2"))
|
||||
_drain_incoming() # consume the first queueing
|
||||
# Resend before it was rendered: inbox still holds it -> dropped, not doubled.
|
||||
client.post("/conjurer", json=_payload("u2"))
|
||||
assert _drain_incoming() == []
|
||||
|
||||
|
||||
def test_resend_after_delivery_is_dropped(spool):
|
||||
inbox, delivered = spool
|
||||
client = cs.app.test_client()
|
||||
client.post("/conjurer", json=_payload("u3"))
|
||||
_drain_incoming()
|
||||
cs.mark_delivered("u3")
|
||||
assert not inbox.contains("u3")
|
||||
assert delivered.contains("u3")
|
||||
# A late resend of an already-delivered result must not re-render.
|
||||
client.post("/conjurer", json=_payload("u3"))
|
||||
assert _drain_incoming() == []
|
||||
|
||||
|
||||
def test_replay_requeues_only_undelivered(spool):
|
||||
inbox, delivered = spool
|
||||
inbox.put("u4", _payload("u4"))
|
||||
inbox.put("u5", _payload("u5"))
|
||||
delivered.put("u5", {}) # u5 already shown to the user
|
||||
cs.replay_inbox()
|
||||
keys = [list(p.keys())[0] for p in _drain_incoming()]
|
||||
assert keys == ["u4"] # only the un-rendered one replayed
|
||||
assert not inbox.contains("u5") # the delivered one is cleaned from the inbox
|
||||
|
||||
|
||||
def test_empty_result_still_persisted_and_delivered(spool):
|
||||
inbox, _delivered = spool
|
||||
cs.app.test_client().post("/conjurer", json={"u6": {}})
|
||||
assert inbox.contains("u6")
|
||||
assert _drain_incoming() == [{"u6": {}}]
|
||||
|
||||
|
||||
def test_pong_is_not_persisted(spool):
|
||||
inbox, _delivered = spool
|
||||
cs.app.test_client().post("/conjurer", json={"__pong__": "ping-1"})
|
||||
assert len(inbox) == 0
|
||||
assert _drain_incoming() == [{"__pong__": "ping-1"}]
|
||||
@@ -136,401 +136,3 @@ def test_map_openai_error_categories():
|
||||
assert ai_functions._map_openai_error(_bare(openai.RateLimitError)).category == "rate_limit"
|
||||
assert ai_functions._map_openai_error(_bare(openai.APITimeoutError)).category == "timeout"
|
||||
assert ai_functions._map_openai_error(ValueError("x")).category == "api"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- Ollama
|
||||
# The self-hosted backend is wired through Ollama's OpenAI-compatible surface,
|
||||
# so it reuses the message format and the error mapping above. What is new and
|
||||
# worth pinning: it must appear in the picker even on an upgraded settings file,
|
||||
# models come from the SERVER, and pinning one must stick.
|
||||
import asyncio # noqa: E402
|
||||
|
||||
|
||||
class _FakeModel:
|
||||
def __init__(self, ident):
|
||||
self.id = ident
|
||||
|
||||
|
||||
class _FakeModels:
|
||||
def __init__(self, ids, raises=None):
|
||||
self._ids = ids
|
||||
self._raises = raises
|
||||
|
||||
async def list(self):
|
||||
if self._raises:
|
||||
raise self._raises
|
||||
return types.SimpleNamespace(data=[_FakeModel(i) for i in self._ids])
|
||||
|
||||
|
||||
class _FakeOllamaClient:
|
||||
def __init__(self, ids=(), raises=None):
|
||||
self.models = _FakeModels(list(ids), raises)
|
||||
|
||||
|
||||
def test_ollama_config_is_offered_in_the_picker():
|
||||
assert "ollama" in ai_functions.list_ai_configs()
|
||||
cfg = ai_functions.AI_CONFIGS["ollama"]
|
||||
assert cfg["provider"] == "ollama"
|
||||
|
||||
|
||||
def test_select_model_uses_ollama_models_when_active(monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
ai_functions.AI_CONFIGS,
|
||||
"ollama",
|
||||
{"provider": "ollama", "latest_model": "llama3.1:8b", "cheap_model": "qwen2.5:3b"},
|
||||
)
|
||||
_reset_active("ollama")
|
||||
try:
|
||||
# the legacy gpt-4o default must auto-map, not leak to Ollama
|
||||
assert ai_functions.select_model("GENERAL", "gpt-4o") == "llama3.1:8b"
|
||||
assert ai_functions.select_model("MUSIC", "gpt-4o") == "qwen2.5:3b"
|
||||
# an explicit id is still honoured verbatim
|
||||
assert ai_functions.select_model("GENERAL", "mistral:7b") == "mistral:7b"
|
||||
finally:
|
||||
_reset_active("gpt")
|
||||
|
||||
|
||||
def test_list_provider_models_queries_the_ollama_server(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ai_functions, "OLLAMACLIENT", _FakeOllamaClient(["b:2", "a:1", "a:1"])
|
||||
)
|
||||
models = asyncio.run(ai_functions.list_provider_models("ollama"))
|
||||
assert models == ["a:1", "b:2"] # sorted + de-duplicated
|
||||
|
||||
|
||||
def test_list_provider_models_for_hosted_provider_reports_configured_ids():
|
||||
models = asyncio.run(ai_functions.list_provider_models("gpt"))
|
||||
assert models == [
|
||||
ai_functions.AI_CONFIGS["gpt"]["latest_model"],
|
||||
ai_functions.AI_CONFIGS["gpt"]["cheap_model"],
|
||||
]
|
||||
|
||||
|
||||
def test_list_provider_models_wraps_server_failure(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ai_functions, "OLLAMACLIENT", _FakeOllamaClient(raises=ValueError("boom"))
|
||||
)
|
||||
try:
|
||||
asyncio.run(ai_functions.list_provider_models("ollama"))
|
||||
except ai_functions.AIError as exc:
|
||||
assert exc.category == "api"
|
||||
else:
|
||||
raise AssertionError("a server failure must surface as AIError")
|
||||
|
||||
|
||||
def test_list_provider_models_without_endpoint_is_an_auth_error(monkeypatch):
|
||||
# Contrast, so the assertion cannot pass vacuously: with a client present the
|
||||
# call succeeds, and ONLY setting it to None turns it into an auth error.
|
||||
monkeypatch.setattr(ai_functions, "OLLAMACLIENT", _FakeOllamaClient(["a:1"]))
|
||||
assert asyncio.run(ai_functions.list_provider_models("ollama")) == ["a:1"]
|
||||
|
||||
monkeypatch.setattr(ai_functions, "OLLAMACLIENT", None)
|
||||
try:
|
||||
asyncio.run(ai_functions.list_provider_models("ollama"))
|
||||
except ai_functions.AIError as exc:
|
||||
assert exc.category == "auth"
|
||||
else:
|
||||
raise AssertionError("an unconfigured Ollama must surface as AIError")
|
||||
|
||||
|
||||
def test_set_active_model_pins_latest_and_keeps_cheap(monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
ai_functions.AI_CONFIGS,
|
||||
"ollama",
|
||||
{"provider": "ollama", "latest_model": "old:1", "cheap_model": "cheap:1"},
|
||||
)
|
||||
written = {}
|
||||
monkeypatch.setattr(
|
||||
ai_functions,
|
||||
"_persist_active_ai_config",
|
||||
lambda name, model_for=None: written.update(name=name, model_for=model_for),
|
||||
|
||||
)
|
||||
ai_functions.set_active_model("mistral:7b", "ollama")
|
||||
# Assert on the shared registry, not on the returned object - that object IS
|
||||
# the mutated dict, so asserting on it would pass even if nothing was stored.
|
||||
stored = ai_functions.AI_CONFIGS["ollama"]
|
||||
assert stored["latest_model"] == "mistral:7b"
|
||||
assert stored["cheap_model"] == "cheap:1" # MUSIC path untouched
|
||||
assert written["name"] # the choice was persisted...
|
||||
assert written["model_for"] == "ollama" # ...scoped to the config we changed
|
||||
|
||||
|
||||
def test_set_active_model_rejects_blank_and_unknown_config(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ai_functions, "_persist_active_ai_config", lambda _n, model_for=None: None
|
||||
)
|
||||
for bad in ("", " "):
|
||||
try:
|
||||
ai_functions.set_active_model(bad, "gpt")
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("a blank model id must be rejected")
|
||||
try:
|
||||
ai_functions.set_active_model("x", "nie-ma-takiego")
|
||||
except KeyError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("an unknown config must be rejected")
|
||||
|
||||
|
||||
def test_switching_to_ollama_without_endpoint_explains_itself(monkeypatch):
|
||||
monkeypatch.setattr(ai_functions, "OLLAMACLIENT", None)
|
||||
try:
|
||||
ai_functions.set_active_ai_config("ollama")
|
||||
except RuntimeError as exc:
|
||||
assert "CONJURER_OLLAMA_URL" in str(exc)
|
||||
else:
|
||||
raise AssertionError("switching to an unconfigured Ollama must raise")
|
||||
finally:
|
||||
_reset_active("gpt")
|
||||
|
||||
|
||||
def test_provider_generate_routes_to_ollama(monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
ai_functions.AI_CONFIGS,
|
||||
"ollama",
|
||||
{"provider": "ollama", "latest_model": "m:1", "cheap_model": "m:1"},
|
||||
)
|
||||
_reset_active("ollama")
|
||||
seen = {}
|
||||
|
||||
async def _fake_ollama_call(messages, model, cfg):
|
||||
seen.update(model=model, messages=messages)
|
||||
return "odpowiedź z domu"
|
||||
|
||||
monkeypatch.setattr(ai_functions, "_ollama_call", _fake_ollama_call)
|
||||
try:
|
||||
out = asyncio.run(
|
||||
ai_functions.provider_generate([{"role": "user", "content": "hej"}], "m:1")
|
||||
)
|
||||
finally:
|
||||
_reset_active("gpt")
|
||||
assert out == "odpowiedź z domu"
|
||||
assert seen["model"] == "m:1"
|
||||
|
||||
|
||||
# ------------------------------------------------- persistence (real disk path)
|
||||
# This path had NO coverage, which is exactly how a config-clobbering regression
|
||||
# got in: persisting the whole in-memory AI_CONFIGS (built-in defaults merged
|
||||
# under the file) overwrote operator hand-edits and resurrected deleted configs.
|
||||
import json # noqa: E402
|
||||
|
||||
|
||||
def _settings_file(tmp_path, configs, active="gpt"):
|
||||
path = tmp_path / "system_gpt_settings.json"
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{"role": "system", "content": "sys"},
|
||||
{"someuser": [1, "a", "b", "c", "asst_x"]},
|
||||
{"active": active, "configs": configs},
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def test_persist_writes_the_pin_without_clobbering_operator_edits(tmp_path, monkeypatch):
|
||||
# The file is authoritative for everything the bot does not itself change:
|
||||
# a hand-tuned cheap_model, and a config deliberately deleted from it.
|
||||
settings = _settings_file(
|
||||
tmp_path,
|
||||
{"gpt": {"provider": "openai", "latest_model": "gpt-4.1", "cheap_model": "hand-tuned"}},
|
||||
)
|
||||
monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(settings))
|
||||
monkeypatch.setitem(
|
||||
ai_functions.AI_CONFIGS,
|
||||
"ollama",
|
||||
{"provider": "ollama", "latest_model": "mistral:7b", "cheap_model": "c:1"},
|
||||
)
|
||||
|
||||
ai_functions._persist_active_ai_config("ollama", model_for="ollama")
|
||||
|
||||
data = json.loads(settings.read_text(encoding="utf-8"))
|
||||
configs = data[2]["configs"]
|
||||
assert data[2]["active"] == "ollama"
|
||||
assert configs["ollama"]["latest_model"] == "mistral:7b" # the pin landed
|
||||
assert configs["gpt"]["latest_model"] == "gpt-4.1" # edit survived
|
||||
assert configs["gpt"]["cheap_model"] == "hand-tuned" # edit survived
|
||||
assert "claude" not in configs # a deleted config is NOT resurrected
|
||||
assert data[0]["content"] == "sys" and "someuser" in data[1] # rest intact
|
||||
|
||||
|
||||
def test_plain_switch_leaves_the_configs_block_untouched(tmp_path, monkeypatch):
|
||||
original = {"gpt": {"provider": "openai", "latest_model": "gpt-4.1", "cheap_model": "hand-tuned"}}
|
||||
settings = _settings_file(tmp_path, original, active="claude")
|
||||
monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(settings))
|
||||
|
||||
# Switching backend without pinning a model must only move "active".
|
||||
ai_functions._persist_active_ai_config("gpt")
|
||||
|
||||
data = json.loads(settings.read_text(encoding="utf-8"))
|
||||
assert data[2]["active"] == "gpt"
|
||||
assert data[2]["configs"] == original
|
||||
|
||||
|
||||
def test_pinned_model_survives_a_restart(tmp_path, monkeypatch):
|
||||
# The whole point of persisting: re-reading the file must yield the pin.
|
||||
settings = _settings_file(tmp_path, {"ollama": {"provider": "ollama", "latest_model": "old:1"}})
|
||||
monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(settings))
|
||||
monkeypatch.setitem(
|
||||
ai_functions.AI_CONFIGS,
|
||||
"ollama",
|
||||
{"provider": "ollama", "latest_model": "new:2", "cheap_model": "c:1"},
|
||||
)
|
||||
|
||||
ai_functions._persist_active_ai_config("ollama", model_for="ollama")
|
||||
|
||||
reread = json.loads(settings.read_text(encoding="utf-8"))[2]
|
||||
assert reread["configs"]["ollama"]["latest_model"] == "new:2"
|
||||
|
||||
|
||||
def test_persist_survives_an_unreadable_settings_file(tmp_path, monkeypatch):
|
||||
# Best-effort by contract: a broken file must not raise into the command.
|
||||
broken = tmp_path / "broken.json"
|
||||
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) == []
|
||||
|
||||
@@ -1,76 +0,0 @@
|
||||
"""Unit tests for the disk-backed durable queue used by result delivery."""
|
||||
import json
|
||||
|
||||
from durable_queue import DiskQueue
|
||||
|
||||
|
||||
def test_put_contains_remove(tmp_path):
|
||||
q = DiskQueue(str(tmp_path / "q"))
|
||||
assert not q.contains("a")
|
||||
q.put("a", {"hello": 1})
|
||||
assert q.contains("a")
|
||||
q.remove("a")
|
||||
assert not q.contains("a")
|
||||
q.remove("a") # idempotent - no error on missing
|
||||
|
||||
|
||||
def test_put_overwrites_and_roundtrips_payload(tmp_path):
|
||||
q = DiskQueue(str(tmp_path / "q"))
|
||||
q.put("uuid-1", {"uuid-1": {"10.1/x": {"Title": ["P"], "type": "article"}}})
|
||||
q.put("uuid-1", {"uuid-1": {"changed": True}})
|
||||
items = q.items()
|
||||
assert len(items) == 1
|
||||
key, payload, _ts = items[0]
|
||||
assert key == "uuid-1"
|
||||
assert payload == {"uuid-1": {"changed": True}}
|
||||
|
||||
|
||||
def test_items_sorted_oldest_first(tmp_path, monkeypatch):
|
||||
q = DiskQueue(str(tmp_path / "q"))
|
||||
import durable_queue
|
||||
|
||||
times = iter([100.0, 200.0, 300.0])
|
||||
monkeypatch.setattr(durable_queue.time, "time", lambda: next(times))
|
||||
q.put("c", {})
|
||||
q.put("a", {})
|
||||
q.put("b", {})
|
||||
assert [k for k, _p, _ts in q.items()] == ["c", "a", "b"]
|
||||
|
||||
|
||||
def test_corrupt_file_is_skipped_not_fatal(tmp_path):
|
||||
directory = tmp_path / "q"
|
||||
q = DiskQueue(str(directory))
|
||||
q.put("good", {"ok": 1})
|
||||
(directory / "broken.json").write_text("{ this is not json", encoding="utf-8")
|
||||
keys = q.keys()
|
||||
assert keys == ["good"] # broken file skipped, good one survives
|
||||
|
||||
|
||||
def test_prune_keeps_newest(tmp_path, monkeypatch):
|
||||
q = DiskQueue(str(tmp_path / "q"))
|
||||
import durable_queue
|
||||
|
||||
times = iter([1.0, 2.0, 3.0, 4.0, 5.0])
|
||||
monkeypatch.setattr(durable_queue.time, "time", lambda: next(times))
|
||||
for key in ("k1", "k2", "k3", "k4", "k5"):
|
||||
q.put(key, {})
|
||||
dropped = q.prune(2)
|
||||
assert dropped == 3
|
||||
assert set(q.keys()) == {"k4", "k5"}
|
||||
|
||||
|
||||
def test_atomic_write_leaves_no_tmp_files(tmp_path):
|
||||
directory = tmp_path / "q"
|
||||
q = DiskQueue(str(directory))
|
||||
q.put("a", {"x": 1})
|
||||
leftover = [p.name for p in directory.iterdir() if p.suffix == ".tmp"]
|
||||
assert leftover == []
|
||||
|
||||
|
||||
def test_key_with_slashes_is_sanitised(tmp_path):
|
||||
q = DiskQueue(str(tmp_path / "q"))
|
||||
q.put("../../etc/passwd", {"evil": 1})
|
||||
# Stays inside the directory (no traversal), and round-trips by key.
|
||||
files = list((tmp_path / "q").iterdir())
|
||||
assert all(f.parent == tmp_path / "q" for f in files)
|
||||
assert q.items()[0][1] == {"evil": 1}
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Unit tests for the per-query watchdog verdict logic.
|
||||
|
||||
The grace window is the whole point: a search that has just finished is briefly
|
||||
'unknown' on the librarian while its result is still in flight, and we must NOT
|
||||
flag that as lost. Only a uuid that stays unknown-but-pending past the grace
|
||||
window is a genuinely lost result.
|
||||
"""
|
||||
from librarian_watchdog import FLAG, WAIT, pending_verdict
|
||||
|
||||
GRACE = 45
|
||||
|
||||
|
||||
def test_known_resets_clock_and_waits():
|
||||
action, unknown_since = pending_verdict(
|
||||
known=True, unknown_since=100.0, now=200.0, grace_seconds=GRACE
|
||||
)
|
||||
assert action == WAIT
|
||||
assert unknown_since is None # clock reset while it's still known
|
||||
|
||||
|
||||
def test_first_unknown_starts_grace_but_does_not_flag():
|
||||
action, unknown_since = pending_verdict(
|
||||
known=False, unknown_since=None, now=1000.0, grace_seconds=GRACE
|
||||
)
|
||||
assert action == WAIT
|
||||
assert unknown_since == 1000.0 # clock started now
|
||||
|
||||
|
||||
def test_unknown_within_grace_keeps_waiting():
|
||||
action, unknown_since = pending_verdict(
|
||||
known=False, unknown_since=1000.0, now=1000.0 + GRACE - 1, grace_seconds=GRACE
|
||||
)
|
||||
assert action == WAIT
|
||||
assert unknown_since == 1000.0 # unchanged, still counting
|
||||
|
||||
|
||||
def test_unknown_past_grace_flags_lost():
|
||||
action, unknown_since = pending_verdict(
|
||||
known=False, unknown_since=1000.0, now=1000.0 + GRACE, grace_seconds=GRACE
|
||||
)
|
||||
assert action == FLAG
|
||||
assert unknown_since == 1000.0
|
||||
|
||||
|
||||
def test_recovered_to_known_after_being_unknown_resets():
|
||||
# It reappeared (e.g. requeued / status flapped): do not flag, reset.
|
||||
action, unknown_since = pending_verdict(
|
||||
known=True, unknown_since=1000.0, now=1000.0 + GRACE + 10, grace_seconds=GRACE
|
||||
)
|
||||
assert action == WAIT
|
||||
assert unknown_since is None
|
||||
@@ -27,32 +27,17 @@ def _write_chunks(directory, count, target=None, target_index=None):
|
||||
(directory / f"{n}_chunk.txt").write_text("".join(lines), encoding="utf-8")
|
||||
|
||||
|
||||
def _run_full(dois, timeout=20, stop_event=None, resume=None):
|
||||
"""Run search_for_doi in a thread; return the whole result box.
|
||||
|
||||
search_for_doi now returns (result_list, positions, interrupted); the box
|
||||
exposes all three (plus 'finished' and 'live') for the resume tests.
|
||||
"""
|
||||
def _run_bounded(dois, timeout=20):
|
||||
"""Run search_for_doi in a thread; return (finished_in_time, result)."""
|
||||
box = {}
|
||||
live = []
|
||||
|
||||
def _run():
|
||||
result_list, positions, interrupted = search_bot.search_for_doi(
|
||||
dois, live, _LOG, stop_event=stop_event, resume=resume
|
||||
worker = threading.Thread(
|
||||
target=lambda: box.update(result=search_bot.search_for_doi(dois, live, _LOG)),
|
||||
daemon=True,
|
||||
)
|
||||
box.update(result=result_list, positions=positions, interrupted=interrupted, live=live)
|
||||
|
||||
worker = threading.Thread(target=_run, daemon=True)
|
||||
worker.start()
|
||||
worker.join(timeout)
|
||||
box["finished"] = not worker.is_alive()
|
||||
return box
|
||||
|
||||
|
||||
def _run_bounded(dois, timeout=20, stop_event=None, resume=None):
|
||||
"""Back-compat wrapper: return (finished_in_time, result_list)."""
|
||||
box = _run_full(dois, timeout, stop_event, resume)
|
||||
return box.get("finished"), box.get("result")
|
||||
return (not worker.is_alive()), box.get("result")
|
||||
|
||||
|
||||
def test_finds_doi_in_trailing_chunk(tmp_path, monkeypatch):
|
||||
@@ -109,33 +94,6 @@ def test_survives_invalid_utf8_byte_and_still_finds_later_doi(tmp_path, monkeypa
|
||||
assert hit, "DOI after the bad byte was not found - the file was aborted mid-read"
|
||||
|
||||
|
||||
def test_doi_match_is_exact_not_substring(tmp_path, monkeypatch):
|
||||
# A DB line "10.1/12" must NOT satisfy a search for "10.1/1" (the old
|
||||
# `doi in line` substring test did). The exact DOI must still be found.
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
(tmp_path / "0_chunk.txt").write_text(
|
||||
"10.1/12\n10.1/1\n10.2/999\n", encoding="utf-8"
|
||||
)
|
||||
|
||||
finished, result = _run_bounded([("10.1/1", "DATA"), ("10.9/absent", "DATA")])
|
||||
|
||||
assert finished
|
||||
by_doi = {r["DOI"]: r["exists"] for r in result}
|
||||
assert by_doi["10.1/1"] is True # exact line present -> found
|
||||
assert by_doi["10.9/absent"] is False
|
||||
|
||||
|
||||
def test_doi_match_handles_line_with_trailing_metadata(tmp_path, monkeypatch):
|
||||
# Lines of the form "<DOI>\t<metadata>" still match on the first token.
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
(tmp_path / "0_chunk.txt").write_text("10.5/abc\tsome title here\n", encoding="utf-8")
|
||||
|
||||
finished, result = _run_bounded([("10.5/abc", "DATA")])
|
||||
|
||||
assert finished
|
||||
assert result[0]["exists"] is True
|
||||
|
||||
|
||||
def test_discover_chunk_files_sorted_numerically(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
for n in (0, 2, 10, 1):
|
||||
@@ -146,71 +104,3 @@ def test_discover_chunk_files_sorted_numerically(tmp_path, monkeypatch):
|
||||
|
||||
# Numeric order (10 after 2, not lexicographic), and non-chunk files ignored.
|
||||
assert found == ["0_chunk.txt", "1_chunk.txt", "2_chunk.txt", "10_chunk.txt"]
|
||||
|
||||
|
||||
def _offset_after(path, marker):
|
||||
"""Byte-cookie (tell) just past the line equal to `marker` in `path`."""
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
while True:
|
||||
line = handle.readline()
|
||||
if not line:
|
||||
raise AssertionError(f"marker {marker!r} not found")
|
||||
if line.strip() == marker:
|
||||
return handle.tell()
|
||||
|
||||
|
||||
def test_resume_seeks_past_scanned_part_and_continues(tmp_path, monkeypatch):
|
||||
# Chunk: early | first-half decoy | MIDDLE | late. Resume from just past
|
||||
# MIDDLE with 'early' pre-found. The scan must: keep 'early' (pre-marked),
|
||||
# find 'late' (after the resume point), and NOT find the first-half decoy
|
||||
# (proving it seeked past it instead of re-reading from the top).
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
path = tmp_path / "0_chunk.txt"
|
||||
path.write_text(
|
||||
"10.1/early\n10.1/only-first-half\nMIDDLE\n10.1/late\n", encoding="utf-8"
|
||||
)
|
||||
offset = _offset_after(str(path), "MIDDLE")
|
||||
|
||||
resume = {"found": ["10.1/early"], "positions": {"0_chunk.txt": offset}}
|
||||
box = _run_full(
|
||||
[("10.1/early", "D"), ("10.1/late", "D"), ("10.1/only-first-half", "D")],
|
||||
resume=resume,
|
||||
)
|
||||
|
||||
assert box["finished"]
|
||||
by_doi = {r["DOI"]: r["exists"] for r in box["result"]}
|
||||
assert by_doi["10.1/early"] is True # carried over from the checkpoint
|
||||
assert by_doi["10.1/late"] is True # found after the resume offset
|
||||
assert by_doi["10.1/only-first-half"] is False # skipped - not re-scanned
|
||||
|
||||
|
||||
def test_stop_event_interrupts_and_reports_positions(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
_write_chunks(tmp_path, count=2)
|
||||
stop = __import__("threading").Event()
|
||||
stop.set() # already asked to stop before it starts
|
||||
|
||||
box = _run_full([("10.0000/decoy-0-a", "D")], stop_event=stop)
|
||||
|
||||
assert box["finished"], "an already-set stop must not hang the search"
|
||||
assert box["interrupted"] is True
|
||||
assert isinstance(box["positions"], dict)
|
||||
|
||||
|
||||
def test_bounded_queue_does_not_deadlock_on_early_termination(tmp_path, monkeypatch):
|
||||
# The OOM fix bounds the work queue. That means a producer can block on a
|
||||
# FULL queue - and if the consumers have already finished (all DOIs found)
|
||||
# it must notice the TERM sentinel instead of hanging forever. Tiny queue +
|
||||
# target on the first line + thousands of trailing decoys the producer still
|
||||
# holds is exactly that situation.
|
||||
monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
|
||||
monkeypatch.setattr(search_bot, "WORK_Q_SIZE", 3) # force the producer to block
|
||||
target = "10.1234/found.on.line.one"
|
||||
lines = [target + "\n"] + [f"10.0000/decoy-{i}\n" for i in range(5000)]
|
||||
(tmp_path / "0_chunk.txt").write_text("".join(lines), encoding="utf-8")
|
||||
|
||||
finished, result = _run_bounded([(target, "DATA")], timeout=20)
|
||||
|
||||
assert finished, "a full bounded queue deadlocked the producer on early termination"
|
||||
hit = [r for r in result if r["DOI"] == target and r["exists"]]
|
||||
assert hit, "the target on the first line should have been found"
|
||||
|
||||
Reference in New Issue
Block a user