diff --git a/.github/workflows/build-windows.yml b/.github/workflows/build-windows.yml index a8cd58f1..770287e3 100644 --- a/.github/workflows/build-windows.yml +++ b/.github/workflows/build-windows.yml @@ -141,14 +141,18 @@ jobs: - name: Create zip for release run: | python -c " - import zipfile, pathlib - src = pathlib.Path('artifacts/AudioMuse-AI') + import zipfile, pathlib, sys + # upload-artifact strips the directory pointed at by 'path:', so the + # bundle's files land directly under artifacts/ (not artifacts/AudioMuse-AI/). + src = pathlib.Path('artifacts') out = 'AudioMuse-AI-amd64-windows.zip' + files = [f for f in sorted(src.rglob('*')) if f.is_file()] + if not files: + sys.exit('No files found under artifacts/ -- the build artifact is missing or empty; refusing to ship an empty zip') with zipfile.ZipFile(out, 'w', zipfile.ZIP_DEFLATED, allowZip64=True) as zf: - for f in sorted(src.rglob('*')): - if f.is_file(): - zf.write(f, f.relative_to(src.parent).as_posix()) - print(f'Created {out}') + for f in files: + zf.write(f, 'AudioMuse-AI/' + f.relative_to(src).as_posix()) + print(f'Created {out} with {len(files)} files') " - name: Attach packages to the release diff --git a/.github/workflows/lint-flake8.yml b/.github/workflows/lint-flake8.yml index 412e4a82..1b216e58 100644 --- a/.github/workflows/lint-flake8.yml +++ b/.github/workflows/lint-flake8.yml @@ -28,4 +28,4 @@ jobs: - name: Run flake8 static analysis run: | - flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics + flake8 . --count --select=E9,F63,F7,F82,F401,F811 --show-source --statistics diff --git a/app.py b/app.py index d316579d..612f5041 100644 --- a/app.py +++ b/app.py @@ -32,10 +32,13 @@ # Import helper functions from app_helper import ( - init_db, get_db, close_db, + get_db, close_db, redis_conn, get_task_info_from_db, cancel_job_and_children_recursive, +) +from database import init_db +from config import ( TASK_STATUS_PENDING, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED ) @@ -46,8 +49,6 @@ resolve_jwt_secret, ) -from app_provider_migration import migration_bp - from error import error_manager from error.error_manager import AudioMuseError from error.error_dictionary import UNKNOWN_ERROR_CODE @@ -172,7 +173,7 @@ def teardown_db(e=None): else: app.logger.info("RQ worker mode: skipping startup database schema bootstrap.") -import app_setup +import app_setup # noqa: F401 # --- API Endpoints --- @@ -705,7 +706,7 @@ def listen_for_index_reloads(): load_voyager_index_for_querying(force_reload=True) from tasks.artist_gmm_manager import load_artist_index_for_querying load_artist_index_for_querying(force_reload=True) - from app_helper import load_map_projection, load_artist_projection + from database import load_map_projection, load_artist_projection load_map_projection('main_map', force_reload=True) load_artist_projection('artist_map', force_reload=True) # Rebuild the map JSON cache used by the /api/map endpoint @@ -757,48 +758,56 @@ def listen_for_index_reloads(): -# --- Import and Register Blueprints --- -# This is the original, working structure. -from app_chat import chat_bp -from app_clustering import clustering_bp -from app_analysis import analysis_bp -from app_cron import cron_bp, run_due_cron_jobs -from app_voyager import voyager_bp -from app_sonic_fingerprint import sonic_fingerprint_bp -from app_path import path_bp -from app_external import external_bp # --- NEW: Import the external blueprint --- -from app_alchemy import alchemy_bp -from app_map import map_bp -from app_waveform import waveform_bp -from app_artist_similarity import artist_similarity_bp -from app_clap_search import clap_search_bp -from app_lyrics import lyrics_search_bp -from app_sem_grove import sem_grove_bp -from app_backup import backup_bp -from app_dashboard import dashboard_bp -from app_users import users_bp -from app_sync import sync_bp - -app.register_blueprint(chat_bp, url_prefix='/chat') -app.register_blueprint(clustering_bp) -app.register_blueprint(analysis_bp) -app.register_blueprint(cron_bp) -app.register_blueprint(voyager_bp) -app.register_blueprint(sonic_fingerprint_bp) -app.register_blueprint(path_bp) -app.register_blueprint(external_bp, url_prefix='/external') # --- NEW: Register the external blueprint --- -app.register_blueprint(alchemy_bp) -app.register_blueprint(map_bp) -app.register_blueprint(waveform_bp) -app.register_blueprint(artist_similarity_bp) -app.register_blueprint(clap_search_bp) -app.register_blueprint(lyrics_search_bp) -app.register_blueprint(sem_grove_bp) -app.register_blueprint(backup_bp) -app.register_blueprint(migration_bp) -app.register_blueprint(dashboard_bp) -app.register_blueprint(users_bp) -app.register_blueprint(sync_bp) +# --- Blueprint Registration --- +# Standard Flask factory pattern: blueprint imports are inside +# this function so the eager import graph stays flat. + + +def _register_blueprints(flask_app): + from app_chat import chat_bp + from app_clustering import clustering_bp + from app_analysis import analysis_bp + from app_cron import cron_bp + from app_voyager import voyager_bp + from app_sonic_fingerprint import sonic_fingerprint_bp + from app_path import path_bp + from app_external import external_bp + from app_alchemy import alchemy_bp + from app_map import map_bp + from app_waveform import waveform_bp + from app_artist_similarity import artist_similarity_bp + from app_clap_search import clap_search_bp + from app_lyrics import lyrics_search_bp + from app_sem_grove import sem_grove_bp + from app_backup import backup_bp + from app_provider_migration import migration_bp + from app_dashboard import dashboard_bp + from app_users import users_bp + from app_sync import sync_bp + + flask_app.register_blueprint(chat_bp, url_prefix='/chat') + flask_app.register_blueprint(clustering_bp) + flask_app.register_blueprint(analysis_bp) + flask_app.register_blueprint(cron_bp) + flask_app.register_blueprint(voyager_bp) + flask_app.register_blueprint(sonic_fingerprint_bp) + flask_app.register_blueprint(path_bp) + flask_app.register_blueprint(external_bp, url_prefix='/external') + flask_app.register_blueprint(alchemy_bp) + flask_app.register_blueprint(map_bp) + flask_app.register_blueprint(waveform_bp) + flask_app.register_blueprint(artist_similarity_bp) + flask_app.register_blueprint(clap_search_bp) + flask_app.register_blueprint(lyrics_search_bp) + flask_app.register_blueprint(sem_grove_bp) + flask_app.register_blueprint(backup_bp) + flask_app.register_blueprint(migration_bp) + flask_app.register_blueprint(dashboard_bp) + flask_app.register_blueprint(users_bp) + flask_app.register_blueprint(sync_bp) + + +_register_blueprints(app) # --- Startup: Load indexes and caches (Flask server only, NOT RQ workers) --- # RQ workers import app.py but should NOT load indexes or start background threads. @@ -828,7 +837,7 @@ def listen_for_index_reloads(): logger.debug(f"No precomputed map projection to load at startup or load failed: {e}") # Also try to load artist component projection into memory try: - from app_helper import load_artist_projection + from database import load_artist_projection load_artist_projection('artist_map') logger.info("In-memory artist component projection loaded at startup.") except Exception as e: @@ -895,6 +904,7 @@ def _start_map_init_background(): def _cron_manager_loop(): try: from time import sleep + from app_cron import run_due_cron_jobs while True: try: with app.app_context(): diff --git a/app_alchemy.py b/app_alchemy.py index 3c143d60..e849c368 100644 --- a/app_alchemy.py +++ b/app_alchemy.py @@ -185,7 +185,7 @@ def list_anchors(): 500: description: Database error. """ - from app_helper import get_alchemy_anchors + from database import get_alchemy_anchors try: anchors = get_alchemy_anchors() # no centroid returned here (name-only list) @@ -227,7 +227,7 @@ def create_anchor(): 500: description: Database failure. """ - from app_helper import save_alchemy_anchor + from database import save_alchemy_anchor payload = request.get_json() or {} name = (payload.get('name') or '').strip() centroid = payload.get('centroid') @@ -260,7 +260,7 @@ def remove_anchor(anchor_id): 404: description: Anchor not found. """ - from app_helper import delete_alchemy_anchor + from database import delete_alchemy_anchor ok = delete_alchemy_anchor(anchor_id) if not ok: return jsonify({'error': 'Anchor not found'}), 404 @@ -298,7 +298,7 @@ def rename_anchor(anchor_id): 404: description: Anchor not found. """ - from app_helper import update_alchemy_anchor_name + from database import update_alchemy_anchor_name payload = request.get_json() or {} name = (payload.get('name') or '').strip() if not name: @@ -371,7 +371,7 @@ def list_radios(): 500: description: Database error. """ - from app_helper import get_alchemy_radios + from database import get_alchemy_radios try: radios = get_alchemy_radios() return jsonify({'radios': [{ @@ -418,7 +418,7 @@ def create_radio(): 500: description: Database failure. """ - from app_helper import create_alchemy_radio + from database import create_alchemy_radio payload = request.get_json() or {} anchor_id = payload.get('anchor_id') try: @@ -471,7 +471,7 @@ def update_radio(radio_id): 404: description: Radio not found. """ - from app_helper import update_alchemy_radio + from database import update_alchemy_radio payload = request.get_json() or {} temperature, n_results, error = _parse_radio_settings(payload) if error: @@ -502,7 +502,7 @@ def remove_radio(radio_id): 404: description: Radio not found. """ - from app_helper import delete_alchemy_radio + from database import delete_alchemy_radio ok = delete_alchemy_radio(radio_id) if not ok: return jsonify({'error': 'Radio not found'}), 404 @@ -588,7 +588,7 @@ def artist_projections_api(): 500: description: Failure to read cache. """ - from app_helper import ARTIST_PROJECTION_CACHE + from database import ARTIST_PROJECTION_CACHE try: if not ARTIST_PROJECTION_CACHE: diff --git a/app_analysis.py b/app_analysis.py index f88c4480..0176f95a 100644 --- a/app_analysis.py +++ b/app_analysis.py @@ -1,14 +1,18 @@ # app_analysis.py -from flask import Blueprint, jsonify, request +from flask import Blueprint, jsonify, request, render_template import uuid import logging # Import configuration from the main config.py -from config import NUM_RECENT_ALBUMS, TOP_N_MOODS +from config import NUM_RECENT_ALBUMS, TOP_N_MOODS, TASK_STATUS_PENDING # RQ import from rq import Retry +# App helper functions +from app_helper import rq_queue_high, save_task_status +from database import clean_up_previous_main_tasks, get_active_main_task + logger = logging.getLogger(__name__) # Create a Blueprint for analysis-related routes @@ -29,7 +33,6 @@ def cleaning_page(): schema: type: string """ - from flask import render_template return render_template('cleaning.html', title = 'AudioMuse-AI - Database Cleaning', active='cleaning') @analysis_bp.route('/api/analysis/start', methods=['POST']) @@ -80,9 +83,6 @@ def start_analysis_endpoint(): 500: description: Server error during task enqueue. """ - # Local imports to prevent circular dependency at startup - from app_helper import rq_queue_high, clean_up_previous_main_tasks, save_task_status, TASK_STATUS_PENDING, get_active_main_task - # Check for any existing active main task to prevent parallel batch runs. active_task = get_active_main_task() if active_task: @@ -146,9 +146,6 @@ def start_cleaning_endpoint(): 500: description: Server error during task enqueue. """ - # Local imports to prevent circular dependency at startup - from app_helper import rq_queue_high, clean_up_previous_main_tasks, save_task_status, TASK_STATUS_PENDING, get_active_main_task - active_task = get_active_main_task() if active_task: return jsonify({ diff --git a/app_clustering.py b/app_clustering.py index 8e4eb3ed..a56383e2 100644 --- a/app_clustering.py +++ b/app_clustering.py @@ -5,7 +5,21 @@ import traceback # Import all necessary configuration variables -from config import MAX_SONGS_PER_CLUSTER, SCORE_WEIGHT_DIVERSITY, SCORE_WEIGHT_SILHOUETTE, SCORE_WEIGHT_DAVIES_BOULDIN, SCORE_WEIGHT_CALINSKI_HARABASZ, SCORE_WEIGHT_PURITY, SCORE_WEIGHT_OTHER_FEATURE_DIVERSITY, SCORE_WEIGHT_OTHER_FEATURE_PURITY, MIN_SONGS_PER_GENRE_FOR_STRATIFICATION, STRATIFIED_SAMPLING_TARGET_PERCENTILE, CLUSTER_ALGORITHM, NUM_CLUSTERS_MIN, NUM_CLUSTERS_MAX, DBSCAN_EPS_MIN, DBSCAN_EPS_MAX, DBSCAN_MIN_SAMPLES_MIN, DBSCAN_MIN_SAMPLES_MAX, GMM_N_COMPONENTS_MIN, GMM_N_COMPONENTS_MAX, SPECTRAL_N_CLUSTERS_MIN, SPECTRAL_N_CLUSTERS_MAX, ENABLE_CLUSTERING_EMBEDDINGS, PCA_COMPONENTS_MIN, PCA_COMPONENTS_MAX, CLUSTERING_RUNS, TOP_N_MOODS, AI_MODEL_PROVIDER, OLLAMA_SERVER_URL, OLLAMA_MODEL_NAME, OPENAI_SERVER_URL, OPENAI_MODEL_NAME, OPENAI_API_KEY, GEMINI_API_KEY, GEMINI_MODEL_NAME, TOP_N_PLAYLISTS, MISTRAL_API_KEY, MISTRAL_MODEL_NAME +from config import ( + MAX_SONGS_PER_CLUSTER, SCORE_WEIGHT_DIVERSITY, SCORE_WEIGHT_SILHOUETTE, + SCORE_WEIGHT_DAVIES_BOULDIN, SCORE_WEIGHT_CALINSKI_HARABASZ, + SCORE_WEIGHT_PURITY, SCORE_WEIGHT_OTHER_FEATURE_DIVERSITY, + SCORE_WEIGHT_OTHER_FEATURE_PURITY, MIN_SONGS_PER_GENRE_FOR_STRATIFICATION, + STRATIFIED_SAMPLING_TARGET_PERCENTILE, CLUSTER_ALGORITHM, NUM_CLUSTERS_MIN, + NUM_CLUSTERS_MAX, DBSCAN_EPS_MIN, DBSCAN_EPS_MAX, DBSCAN_MIN_SAMPLES_MIN, + DBSCAN_MIN_SAMPLES_MAX, GMM_N_COMPONENTS_MIN, GMM_N_COMPONENTS_MAX, + SPECTRAL_N_CLUSTERS_MIN, SPECTRAL_N_CLUSTERS_MAX, ENABLE_CLUSTERING_EMBEDDINGS, + PCA_COMPONENTS_MIN, PCA_COMPONENTS_MAX, CLUSTERING_RUNS, TOP_N_MOODS, + AI_MODEL_PROVIDER, OLLAMA_SERVER_URL, OLLAMA_MODEL_NAME, OPENAI_SERVER_URL, + OPENAI_MODEL_NAME, OPENAI_API_KEY, GEMINI_API_KEY, GEMINI_MODEL_NAME, + TOP_N_PLAYLISTS, MISTRAL_API_KEY, MISTRAL_MODEL_NAME, + TASK_STATUS_PENDING, TASK_STATUS_FAILURE, +) # RQ import from rq import Retry @@ -13,6 +27,10 @@ from error import error_manager from error.error_dictionary import ERR_CLUSTERING_FAILED +# App helper functions +from app_helper import rq_queue_high, save_task_status +from database import clean_up_previous_main_tasks, get_active_main_task + logger = logging.getLogger(__name__) @@ -22,7 +40,6 @@ def clustering_task_failure_handler(job, connection, type, value, tb): """A failure handler for the main clustering task, executed by the worker.""" from flask_app import app - from app_helper import save_task_status, TASK_STATUS_FAILURE with app.app_context(): task_id = getattr(job, 'id', None) or getattr(job, 'get_id', lambda: None)() @@ -245,10 +262,6 @@ def start_clustering_endpoint(): status: type: string """ - # Local imports to prevent circular dependency at startup - from app_helper import rq_queue_high, get_active_main_task - from app_helper import clean_up_previous_main_tasks, save_task_status, TASK_STATUS_PENDING - # Check for any existing active main task to prevent parallel batch runs active_task = get_active_main_task() if active_task: diff --git a/app_cron.py b/app_cron.py index b4b41c93..471226e4 100644 --- a/app_cron.py +++ b/app_cron.py @@ -1,8 +1,8 @@ from flask import Blueprint, render_template, jsonify, request from psycopg2.extras import DictCursor -from database import get_db +from database import get_db, save_task_status from taskqueue import rq_queue_high -from app_helper import save_task_status, TASK_STATUS_PENDING +from config import TASK_STATUS_PENDING import uuid, time, logging from config import ( TOP_N_MOODS, @@ -240,8 +240,6 @@ def run_due_cron_jobs(): "pca_components_max": int(PCA_COMPONENTS_MAX), "num_clustering_runs": int(CLUSTERING_RUNS), "max_songs_per_cluster_val": int(MAX_SONGS_PER_CLUSTER), - "gmm_n_components_min": int(GMM_N_COMPONENTS_MIN), - "gmm_n_components_max": int(GMM_N_COMPONENTS_MAX), "top_n_playlists_param": int(TOP_N_PLAYLISTS), "min_songs_per_genre_for_stratification_param": int(MIN_SONGS_PER_GENRE_FOR_STRATIFICATION), "stratified_sampling_target_percentile_param": int(STRATIFIED_SAMPLING_TARGET_PERCENTILE), diff --git a/app_helper.py b/app_helper.py index 3651850f..fc8d4b30 100644 --- a/app_helper.py +++ b/app_helper.py @@ -1,17 +1,33 @@ -# app_helper.py -import ipaddress +"""App-layer helpers that compose the data (``database``) and queue +(``taskqueue``) layers for the web and task tiers. + +This is NOT the database layer -- all SQL lives in ``database.py``. What remains +here is orchestration and presentation glue: + +- ``cancel_job_and_children_recursive`` -- recursively cancel an RQ job tree. +- ``build_and_store_map_projection`` / ``build_and_store_artist_projection`` -- + compute a 2D projection and persist it via ``database``. +- ``attach_song_features`` / ``top_stratified_genre`` -- enrich API result rows. + +It also re-exports the most commonly used ``database`` / ``taskqueue`` handles so +the many modules doing ``from app_helper import get_db, redis_conn, ...`` are +untouched. +""" import json import logging -import socket -import sys import time -from urllib.parse import urlparse -import psycopg2 from psycopg2.extras import DictCursor import numpy as np -from database import get_db, close_db +import database +from database import ( # noqa: F401 + get_db, close_db, save_task_status, record_task_history, _build_task_note, + get_score_data_by_ids, load_map_projection, get_task_info_from_db, get_tracks_by_ids, + save_track_analysis_and_embedding, + # Used internally by the build_and_store_* projection orchestration below. + save_map_projection, save_artist_projection, +) from taskqueue import ( redis_conn, rq_queue_high, @@ -21,1070 +37,19 @@ send_stop_job_command, ) -from config import STRATIFIED_GENRES -from tz_helper import UTC_NOW_SQL +from config import ( # noqa: F401 + STRATIFIED_GENRES, + TASK_STATUS_PENDING, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, + TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED, +) logger = logging.getLogger(__name__) -# Import app object after it's defined to break circular dependency -# Avoid importing the Flask `app` object here to prevent circular imports. -# Use the module-level `logger` defined above for logging instead of `app.logger`. - -# In-memory cache for the precomputed 2D map projection (optional) -MAP_PROJECTION_CACHE = None - - -def validate_outbound_url(url): - """SSRF guard for user-supplied outbound HTTP(S) URLs. - - Returns ``(True, None)`` when the URL is safe to fetch, else - ``(False, reason)``. - - Self-hosted media servers and APIs (e.g. a private Lyrics API) legitimately - live on the LAN (RFC1918) or the same host (loopback), so those are allowed. - Only what is never a real user service and is a classic SSRF target is - rejected: non-HTTP(S) schemes and link-local / multicast / reserved / - unspecified addresses (notably 169.254.169.254 cloud metadata). - """ - if not url: - return False, 'URL is required' - - try: - parsed = urlparse(str(url)) - except Exception: - return False, 'Invalid URL' - - if parsed.scheme not in ('http', 'https'): - return False, 'Only http and https URLs are supported' - - host = parsed.hostname - if not host: - return False, 'URL host is required' - - try: - addrinfo = socket.getaddrinfo( - host, parsed.port or (443 if parsed.scheme == 'https' else 80), - type=socket.SOCK_STREAM, - ) - except Exception: - return False, 'Could not resolve host' - - for entry in addrinfo: - try: - ip_obj = ipaddress.ip_address(entry[4][0]) - except ValueError: - return False, 'Resolved host to invalid IP' - if ( - ip_obj.is_link_local - or ip_obj.is_multicast - or ip_obj.is_reserved - or ip_obj.is_unspecified - ): - return False, 'Target host resolves to a disallowed IP address' - - return True, None - -# In-memory cache for the precomputed 2D artist component projections -ARTIST_PROJECTION_CACHE = None - -# --- Constants --- -MAX_LOG_ENTRIES_STORED = 10 # Max number of recent log entries to store in the database per task - -def init_db(): - db = get_db() - with db.cursor() as cur: - # Serialize concurrent init_db() runs across gunicorn workers/containers. - # Multiple workers racing on CREATE EXTENSION / CREATE OR REPLACE FUNCTION - # causes Postgres "tuple concurrently updated" errors on pg_proc/pg_extension. - # A session-level advisory lock forces other workers to wait here. - # The key is an arbitrary stable bigint specific to this app's init. - # Safety: session-level advisory locks are auto-released by Postgres - # when the connection ends (normal close, crash, kill, or network drop), - # so this lock can NEVER leak permanently even if init_db() raises. - cur.execute("SELECT pg_advisory_lock(726354821)") - try: - # Enable extensions to fix and assist in searches - if sys.platform == 'win32': - for ext in ('unaccent', 'pg_trgm'): - cur.execute("SAVEPOINT ext_create") - try: - cur.execute(f'CREATE EXTENSION IF NOT EXISTS {ext}') - cur.execute("RELEASE SAVEPOINT ext_create") - except Exception: - logger.warning("Extension %s not available -- skipping", ext) - cur.execute("ROLLBACK TO SAVEPOINT ext_create") - else: - cur.execute('CREATE EXTENSION IF NOT EXISTS unaccent') - cur.execute('CREATE EXTENSION IF NOT EXISTS pg_trgm') - # Create 'score' table - cur.execute("CREATE TABLE IF NOT EXISTS score (item_id TEXT PRIMARY KEY, title TEXT, author TEXT, album TEXT, album_artist TEXT, tempo REAL, key TEXT, scale TEXT, mood_vector TEXT)") - # Add 'energy' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'energy')") - if not cur.fetchone()[0]: - logger.info("Adding 'energy' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN energy REAL") - # Add 'other_features' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'other_features')") - if not cur.fetchone()[0]: - logger.info("Adding 'other_features' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN other_features TEXT") - # Add 'album' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'album')") - if not cur.fetchone()[0]: - logger.info("Adding 'album' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN album TEXT") - # Add 'album_artist' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'album_artist')") - if not cur.fetchone()[0]: - logger.info("Adding 'album_artist' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN album_artist TEXT") - # Add 'year' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'year')") - if not cur.fetchone()[0]: - logger.info("Adding 'year' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN year INTEGER") - # Add 'rating' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'rating')") - if not cur.fetchone()[0]: - logger.info("Adding 'rating' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN rating INTEGER") - # Add 'file_path' column if not exists - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'file_path')") - if not cur.fetchone()[0]: - logger.info("Adding 'file_path' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN file_path TEXT") - - # Ensure we have a searchable, accent-stripped `search_u` column. - # Postgres does not allow generated columns to call `unaccent()` (it's not marked immutable), - # so we store the value in a normal column and keep it in sync via trigger. - cur.execute("SELECT is_generated FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'search_u'") - row = cur.fetchone() - search_u_generated = (row and row[0] == 'ALWAYS') - - if search_u_generated: - logger.info("Dropping legacy generated 'search_u' column to replace it with a trigger-updated column.") - cur.execute("ALTER TABLE score DROP COLUMN IF EXISTS search_u") - row = None - - # Create plain `search_u` column if missing - if not row: - logger.info("Adding 'search_u' column to 'score' table.") - cur.execute("ALTER TABLE score ADD COLUMN search_u TEXT") - - # Create helper function for accent stripping (safe to run multiple times) - if sys.platform == 'win32': - cur.execute("SAVEPOINT search_setup") - try: - cur.execute("CREATE OR REPLACE FUNCTION immutable_unaccent(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$ SELECT public.unaccent($1) $$;") - cur.execute(""" - CREATE OR REPLACE FUNCTION score_search_u_sync() RETURNS trigger LANGUAGE plpgsql AS $$ - BEGIN - NEW.search_u := lower(immutable_unaccent(concat_ws(' ', NEW.title, NEW.author, NEW.album))); - RETURN NEW; - END; - $$; - """) - cur.execute("DROP TRIGGER IF EXISTS score_search_u_sync_trigger ON score") - cur.execute(""" - CREATE TRIGGER score_search_u_sync_trigger - BEFORE INSERT OR UPDATE ON score - FOR EACH ROW - EXECUTE FUNCTION score_search_u_sync(); - """) - cur.execute("UPDATE score SET search_u = lower(immutable_unaccent(concat_ws(' ', title, author, album))) WHERE search_u IS NULL") - cur.execute("CREATE INDEX IF NOT EXISTS score_search_u_trgm ON score USING gin (search_u gin_trgm_ops)") - cur.execute("RELEASE SAVEPOINT search_setup") - except Exception: - logger.warning("unaccent/pg_trgm extensions not available -- accent-insensitive search disabled") - cur.execute("ROLLBACK TO SAVEPOINT search_setup") - else: - cur.execute("CREATE OR REPLACE FUNCTION immutable_unaccent(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$ SELECT public.unaccent($1) $$;") - cur.execute(""" - CREATE OR REPLACE FUNCTION score_search_u_sync() RETURNS trigger LANGUAGE plpgsql AS $$ - BEGIN - NEW.search_u := lower(immutable_unaccent(concat_ws(' ', NEW.title, NEW.author, NEW.album))); - RETURN NEW; - END; - $$; - """) - cur.execute("DROP TRIGGER IF EXISTS score_search_u_sync_trigger ON score") - cur.execute(""" - CREATE TRIGGER score_search_u_sync_trigger - BEFORE INSERT OR UPDATE ON score - FOR EACH ROW - EXECUTE FUNCTION score_search_u_sync(); - """) - cur.execute("UPDATE score SET search_u = lower(immutable_unaccent(concat_ws(' ', title, author, album))) WHERE search_u IS NULL") - cur.execute("CREATE INDEX IF NOT EXISTS score_search_u_trgm ON score USING gin (search_u gin_trgm_ops)") - - # Create 'playlist' table - cur.execute("CREATE TABLE IF NOT EXISTS playlist (id SERIAL PRIMARY KEY, playlist_name TEXT, item_id TEXT, title TEXT, author TEXT, UNIQUE (playlist_name, item_id))") - # Create 'task_status' table - cur.execute("CREATE TABLE IF NOT EXISTS task_status (id SERIAL PRIMARY KEY, task_id TEXT UNIQUE NOT NULL, parent_task_id TEXT, task_type TEXT NOT NULL, sub_type_identifier TEXT, status TEXT, progress INTEGER DEFAULT 0, details TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Migrate 'start_time' and 'end_time' columns - for col_name in ['start_time', 'end_time']: - cur.execute("SELECT data_type FROM information_schema.columns WHERE table_name = 'task_status' AND column_name = %s", (col_name,)) - if not cur.fetchone(): cur.execute(f"ALTER TABLE task_status ADD COLUMN {col_name} DOUBLE PRECISION") - # Create 'task_history' table — a small, persistent log of the last - # completed/cancelled MAIN tasks. Survives the global Cancel button - # which wipes `task_status`. Capped to the most recent 10 rows. - cur.execute(""" - CREATE TABLE IF NOT EXISTS task_history ( - id SERIAL PRIMARY KEY, - recorded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - task_id TEXT, - task_type TEXT, - status TEXT, - duration_seconds DOUBLE PRECISION, - note TEXT - ) - """) - # Create 'embedding' table - cur.execute("CREATE TABLE IF NOT EXISTS embedding (item_id TEXT PRIMARY KEY, FOREIGN KEY (item_id) REFERENCES score (item_id) ON DELETE CASCADE)") - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'embedding' AND column_name = 'embedding')") - if not cur.fetchone()[0]: cur.execute("ALTER TABLE embedding ADD COLUMN embedding BYTEA") - # Create 'lyrics_embedding' table for lyrics similarity and axis scores - cur.execute("CREATE TABLE IF NOT EXISTS lyrics_embedding (item_id TEXT PRIMARY KEY, FOREIGN KEY (item_id) REFERENCES score (item_id) ON DELETE CASCADE)") - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'lyrics_embedding' AND column_name = 'embedding')") - if not cur.fetchone()[0]: cur.execute("ALTER TABLE lyrics_embedding ADD COLUMN embedding BYTEA") - # axis_vector: float32 BYTEA, fixed-order flattened over MUSIC_ANALYSIS_AXES. - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'lyrics_embedding' AND column_name = 'axis_vector')") - if not cur.fetchone()[0]: cur.execute("ALTER TABLE lyrics_embedding ADD COLUMN axis_vector BYTEA") - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'lyrics_embedding' AND column_name = 'updated_at')") - if not cur.fetchone()[0]: cur.execute("ALTER TABLE lyrics_embedding ADD COLUMN updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") - # Create 'clap_embedding' table for CLAP text search embeddings - cur.execute("CREATE TABLE IF NOT EXISTS clap_embedding (item_id TEXT PRIMARY KEY, FOREIGN KEY (item_id) REFERENCES score (item_id) ON DELETE CASCADE)") - cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'clap_embedding' AND column_name = 'embedding')") - if not cur.fetchone()[0]: cur.execute("ALTER TABLE clap_embedding ADD COLUMN embedding BYTEA") - # Create 'voyager_index_data' table - cur.execute("CREATE TABLE IF NOT EXISTS voyager_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'clap_index_data' table for stored CLAP text search indexes - cur.execute("CREATE TABLE IF NOT EXISTS clap_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'lyrics_index_data' table for stored Lyrics voyager indexes (mirrors clap_index_data; supports chunked storage). - cur.execute("CREATE TABLE IF NOT EXISTS lyrics_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'lyrics_axes_index_data' table for the axis-vector voyager index (one binary-friendly vector per song over MUSIC_ANALYSIS_AXES labels). - cur.execute("CREATE TABLE IF NOT EXISTS lyrics_axes_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'artist_index_data' table for artist GMM-based HNSW index - cur.execute("CREATE TABLE IF NOT EXISTS artist_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, artist_map_json TEXT NOT NULL, gmm_params_json TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'artist_metadata_data' table for the per-artist auxiliary - # metadata blob (artist_map + GMM params). Decoupled from the Voyager - # index binary and segmented independently so a single column value - # never crosses PG's 1 GB MaxAllocSize cap, regardless of library size. - cur.execute("CREATE TABLE IF NOT EXISTS artist_metadata_data (name VARCHAR(255) PRIMARY KEY, blob_data BYTEA NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'map_projection_data' table for precomputed 2D map projections - cur.execute("CREATE TABLE IF NOT EXISTS map_projection_data (index_name VARCHAR(255) PRIMARY KEY, projection_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'artist_component_projection' table for precomputed 2D artist component projections - cur.execute("CREATE TABLE IF NOT EXISTS artist_component_projection (index_name VARCHAR(255) PRIMARY KEY, projection_data BYTEA NOT NULL, artist_component_map_json TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'cron' table to hold scheduled jobs (very small and simple) - cur.execute("CREATE TABLE IF NOT EXISTS cron (id SERIAL PRIMARY KEY, name TEXT, task_type TEXT NOT NULL, cron_expr TEXT NOT NULL, enabled BOOLEAN DEFAULT FALSE, last_run DOUBLE PRECISION, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Create 'audiomuse_users' table. Every account (including the - # install-time admin) lives here. 'role' is 'admin' or 'user'. - cur.execute("CREATE TABLE IF NOT EXISTS audiomuse_users (id SERIAL PRIMARY KEY, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, role TEXT NOT NULL DEFAULT 'user', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Lightweight migration for installs that already have the table without a role column. - cur.execute("ALTER TABLE audiomuse_users ADD COLUMN IF NOT EXISTS role TEXT NOT NULL DEFAULT 'user'") - # Create 'dashboard_stats' singleton table (id fixed to 1) that holds - # precomputed content/library aggregates and index counts. Refreshed - # at app startup and hourly by a background job so the dashboard - # does not have to scan the whole `score` table on every poll. - cur.execute( - "CREATE TABLE IF NOT EXISTS dashboard_stats (" - "id INTEGER PRIMARY KEY, " - "updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " - "content JSONB NOT NULL DEFAULT '{}'::jsonb, " - "indexes JSONB NOT NULL DEFAULT '[]'::jsonb, " - "CONSTRAINT dashboard_stats_singleton CHECK (id = 1))" - ) - # Ensure older restored DBs still have the primary key constraint. - cur.execute( - "SELECT COUNT(*) FROM information_schema.table_constraints " - "WHERE table_name = 'dashboard_stats' AND constraint_type = 'PRIMARY KEY'" - ) - row = cur.fetchone() - if row and row[0] == 0: - logger.info("Cleaning dashboard_stats and adding missing primary key constraint to dashboard_stats.id") - cur.execute("DELETE FROM dashboard_stats") - cur.execute("ALTER TABLE dashboard_stats ADD CONSTRAINT dashboard_stats_pkey PRIMARY KEY (id)") - # Create 'artist_mapping' table to map artist names to media server artist IDs - cur.execute("CREATE TABLE IF NOT EXISTS artist_mapping (artist_name TEXT PRIMARY KEY, artist_id TEXT)") - # Create application configuration table to persist setup values. - cur.execute( - "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'app_config')" - ) - if not cur.fetchone()[0]: - cur.execute( - "CREATE TABLE app_config (" - "key TEXT PRIMARY KEY, value TEXT NOT NULL, " - "updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)" - ) - # Create 'alchemy_anchors' table to persist named user anchors for reuse - cur.execute("CREATE TABLE IF NOT EXISTS alchemy_anchors (id SERIAL PRIMARY KEY, name TEXT UNIQUE NOT NULL, centroid JSONB NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - cur.execute("CREATE TABLE IF NOT EXISTS alchemy_radios (id SERIAL PRIMARY KEY, anchor_id INTEGER UNIQUE NOT NULL REFERENCES alchemy_anchors(id) ON DELETE CASCADE, temperature DOUBLE PRECISION NOT NULL, n_results INTEGER NOT NULL, enabled BOOLEAN NOT NULL DEFAULT TRUE, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") - # Provider migration tool: wizard session state (one row per migration attempt) - cur.execute(""" - CREATE TABLE IF NOT EXISTS migration_session ( - id SERIAL PRIMARY KEY, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - completed_at TIMESTAMP, - status TEXT NOT NULL DEFAULT 'in_progress', - source_type TEXT NOT NULL, - target_type TEXT NOT NULL, - target_creds TEXT NOT NULL, - state JSONB NOT NULL DEFAULT '{}' - ) - """) - # Create 'text_search_queries' table for precomputed CLAP text search queries - cur.execute(""" - CREATE TABLE IF NOT EXISTS text_search_queries ( - id SERIAL PRIMARY KEY, - query_text TEXT NOT NULL, - score REAL NOT NULL, - rank INTEGER NOT NULL, - created_at TIMESTAMP DEFAULT NOW(), - UNIQUE(rank) - ) - """) - cur.execute("CREATE INDEX IF NOT EXISTS idx_text_search_queries_rank ON text_search_queries(rank)") - - # Insert default queries if table is empty - cur.execute("SELECT COUNT(*) FROM text_search_queries") - count = cur.fetchone()[0] - - if count == 0: - default_queries = [ - "female vocal romantic trap", - "synth indie pop raspy", - "sad hard rock male vocal", - "funk falsetto energetic", - "groovy sax blues", - "classical relaxed piano", - "belting jazz happy", - "tabla afrobeat fast-paced", - "harmonized vocals slow-paced electronica", - "autotuned gospel excited", - "breathy aggressive house", - "smooth folk mid-tempo", - "deep voice r&b dark", - "punk guitar angry", - "metal choir dreamy", - "chant reggae trumpet", - "high-pitched brass hip-hop", - "disco whispered drum machine", - "happy whispered indie pop", - "synth energetic raspy", - "rock slow-paced cello", - "falsetto jazz excited", - "r&b male vocal romantic", - "harmonized vocals dark trap", - "smooth blues sax", - "high-pitched fast-paced soul", - "female vocal sad hip-hop", - "congas aggressive soul", - "mid-tempo afrobeat autotuned", - "belting funk groovy", - "angry alternative breathy", - "gospel choir steelpan", - "viola relaxed folk", - "dreamy rhodes metal", - "acoustic guitar country chant", - "deep voice orchestra reggae", - "fast-paced synth progressive rock", - "hard rock raspy romantic", - "fast-paced electric guitar progressive rock", - "hard rock aggressive breathy", - "rock high-pitched energetic", - "autotuned energetic hip-hop", - "raspy fast-paced blues", - "belting electronica energetic", - "whispered indie pop aggressive", - "harmonized vocals aggressive synth", - "orchestra whispered romantic", - "belting mid-tempo progressive rock", - "autotuned pop mid-tempo", - "pop energetic synthesizer" - ] - - for rank, query in enumerate(default_queries, start=1): - cur.execute(""" - INSERT INTO text_search_queries (query_text, score, rank, created_at) - VALUES (%s, %s, %s, NOW()) - """, (query, 1.0, rank)) - - logger.info(f"Inserted {len(default_queries)} default DCLAP search queries") - - db.commit() - # Release the advisory lock acquired at the top of init_db(). - finally: - cur.execute("SELECT pg_advisory_unlock(726354821)") - -# --- Status Constants --- -TASK_STATUS_PENDING = "PENDING" -TASK_STATUS_STARTED = "STARTED" -TASK_STATUS_PROGRESS = "PROGRESS" -TASK_STATUS_SUCCESS = "SUCCESS" -TASK_STATUS_FAILURE = "FAILURE" -TASK_STATUS_REVOKED = "REVOKED" - -# --- DB Cleanup Utility --- -def clean_up_previous_main_tasks(): - """ - Cleans up all previous main tasks before a new one starts. - - Archives tasks in SUCCESS state. - - Archives stale tasks stuck in PENDING, STARTED, or PROGRESS states. - - DELETES all child tasks associated with archived parent tasks to prevent DB bloat. - A main task is identified by having a NULL parent_task_id. - """ - db = get_db() # This now calls the function within this file - cur = db.cursor(cursor_factory=DictCursor) - logger.info("Starting cleanup of all previous main tasks.") - - non_terminal_statuses = (TASK_STATUS_PENDING, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS) - - try: - cur.execute("SELECT task_id, status, details, task_type, start_time, end_time FROM task_status WHERE status IN %s AND parent_task_id IS NULL", (non_terminal_statuses,)) - tasks_to_archive = cur.fetchall() - - archived_count = 0 - deleted_children_count = 0 - - for task_row in tasks_to_archive: - task_id = task_row['task_id'] - original_status = task_row['status'] - - original_details_json = task_row['details'] - original_status_message = f"Task was in '{original_status}' state." - - original_details_dict = None - if original_details_json: - try: - original_details_dict = json.loads(original_details_json) - original_status_message = original_details_dict.get("status_message", original_status_message) - except (json.JSONDecodeError, TypeError): - logger.warning(f"Could not parse original details for task {task_id} during archival.") - - # Record into persistent history BEFORE deleting children — the - # note builder needs to query subtasks (e.g. tracks_analyzed). - try: - duration_s = None - if task_row['start_time'] is not None: - end = task_row['end_time'] if task_row['end_time'] is not None else time.time() - duration_s = max(0.0, float(end) - float(task_row['start_time'])) - final_status = TASK_STATUS_SUCCESS if original_status == TASK_STATUS_SUCCESS else TASK_STATUS_REVOKED - record_task_history( - task_id, task_row['task_type'], final_status, - duration_s, details=original_details_dict, - ) - except Exception as e_hist: - logger.debug(f"history record skipped during archive of {task_id}: {e_hist}") - - if original_status == TASK_STATUS_SUCCESS: - archival_reason = "New main task started, old successful task archived." - else: - archival_reason = f"New main task started, stale task (status: {original_status}) has been archived." - - archived_details = { - "log": [f"[Archived] {archival_reason}. Original summary: {original_status_message}"], - "original_status_before_archival": original_status, - "archival_reason": archival_reason - } - archived_details_json = json.dumps(archived_details) - - with db.cursor() as update_cur: - # First, delete all child tasks to prevent DB bloat and avoid counting old tasks - update_cur.execute( - "DELETE FROM task_status WHERE parent_task_id = %s", - (task_id,) - ) - children_deleted = update_cur.rowcount - deleted_children_count += children_deleted - - if children_deleted > 0: - logger.info(f"Deleted {children_deleted} child tasks for parent task {task_id}") - - # Then archive the parent task - update_cur.execute( - "UPDATE task_status SET status = %s, details = %s, progress = 100, timestamp = NOW() WHERE task_id = %s AND status = %s", - (TASK_STATUS_REVOKED, archived_details_json, task_id, original_status) - ) - archived_count += 1 - - if archived_count > 0: - db.commit() - logger.info(f"Archived {archived_count} previous main tasks and deleted {deleted_children_count} child tasks.") - else: - logger.info("No previous main tasks found to clean up.") - except Exception as e_main_clean: - db.rollback() - logger.error(f"Error during the main task cleanup process: {e_main_clean}") - finally: - cur.close() - - -# --------------------------------------------------------------------------- -# Task history (separate from task_status — survives the global Cancel button) -# --------------------------------------------------------------------------- - -TASK_HISTORY_MAX_ROWS = 10 - - -def _build_task_note(task_type, details_obj, db): - """Build a short, human-readable note for a finished task. - - Looks at the ``details`` JSON we stored on the main task and, when needed, - queries subtasks to compute a meaningful number (e.g. total songs analyzed - across all album_analysis subtasks).""" - if not isinstance(details_obj, dict): - details_obj = {} - t = (task_type or '').lower() - - try: - if 'analysis' in t: - # Prefer summing tracks_analyzed from album_analysis subtasks. - try: - with db.cursor() as cur: - cur.execute( - "SELECT details FROM task_status WHERE parent_task_id = %s AND status = 'SUCCESS'", - (details_obj.get('_task_id') or '',), - ) - rows = cur.fetchall() - except Exception: - rows = [] - songs = 0 - for (d,) in rows or []: - if not d: - continue - try: - obj = json.loads(d) - if isinstance(obj, dict): - v = obj.get('tracks_analyzed') - if isinstance(v, (int, float)): - songs += int(v) - except Exception: - continue - if songs > 0: - return f"Songs analyzed: {songs}" - # Fallback to album-level info from the main task details. - albums = details_obj.get('albums_completed') or details_obj.get('total_albums_processed') - if albums: - return f"Albums analyzed: {albums}" - return '' - - if 'clean' in t: - for k in ('tracks_deleted', 'orphans_removed', 'songs_cleaned', - 'tracks_removed', 'deleted_count', 'cleaned_tracks'): - v = details_obj.get(k) - if isinstance(v, (int, float)): - return f"Songs cleaned: {int(v)}" - return '' - - if 'cluster' in t: - sampled = (details_obj.get('best_params') or {}).get('initial_subset_size') \ - if isinstance(details_obj.get('best_params'), dict) else None - if sampled is None: - sampled = details_obj.get('sampled_songs') or details_obj.get('num_sampled_songs') - n_clusters = details_obj.get('num_playlists_created') or details_obj.get('num_clusters') - parts = [] - if sampled: - parts.append(f"sampled: {int(sampled)}") - if n_clusters: - parts.append(f"clusters: {int(n_clusters)}") - return ' • '.join(parts) - except Exception as e: - logger.debug(f"task note builder failed for type={task_type}: {e}") - return '' -def record_task_history(task_id, task_type, status, duration_seconds=None, note=None, details=None): - """Insert a row into ``task_history`` and trim the table to the most - recent ``TASK_HISTORY_MAX_ROWS`` entries. - - Safe to call from anywhere; never raises. ``details`` (dict or None) is - used to build a default ``note`` when one is not provided explicitly. - If a short note cannot be inferred, fall back to the task's final - status_message or message text when available. - - The history table is treated as immutable per task_id: once a task has - been recorded, we do not insert a second history row for the same task. - """ - if not task_id: - return - try: - db = get_db() - # If no note was supplied, try to infer one from details. - if note is None: - details_obj = details if isinstance(details, dict) else {} - # Pass task_id through so the analysis branch can query subtasks. - details_obj = dict(details_obj) - details_obj['_task_id'] = task_id - note = _build_task_note(task_type, details_obj, db) or '' - if not note: - note = details_obj.get('status_message') or details_obj.get('message') or '' - - with db.cursor() as cur: - cur.execute( - "SELECT 1 FROM task_history WHERE task_id = %s LIMIT 1", - (task_id,) - ) - if cur.fetchone(): - return - cur.execute( - f""" - INSERT INTO task_history (task_id, task_type, status, duration_seconds, note, recorded_at) - VALUES (%s, %s, %s, %s, %s, {UTC_NOW_SQL}) - """, - (task_id, task_type, status, duration_seconds, note), - ) - # Trim — keep only the most recent rows. - cur.execute( - """ - DELETE FROM task_history - WHERE id NOT IN ( - SELECT id FROM task_history ORDER BY recorded_at DESC, id DESC LIMIT %s - ) - """, - (TASK_HISTORY_MAX_ROWS,), - ) - db.commit() - except Exception as e: - logger.warning(f"record_task_history failed for {task_id}: {e}") - try: - db.rollback() - except Exception: - pass - - -def get_active_main_task(task_type=None): - """Return the currently active main task. - - If task_type is provided, only return an active task of that type. - If task_type is None, return any active main task. - """ - db = get_db() - cur = db.cursor(cursor_factory=DictCursor) - non_terminal_statuses = (TASK_STATUS_PENDING, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS) - - if task_type: - cur.execute(""" - SELECT task_id, task_type, status, details - FROM task_status - WHERE task_type = %s AND status IN %s AND parent_task_id IS NULL - ORDER BY timestamp DESC - LIMIT 1 - """, (task_type, non_terminal_statuses)) - else: - cur.execute(""" - SELECT task_id, task_type, status, details - FROM task_status - WHERE status IN %s AND parent_task_id IS NULL - ORDER BY timestamp DESC - LIMIT 1 - """, (non_terminal_statuses,)) - - active_task = cur.fetchone() - cur.close() - return dict(active_task) if active_task else None - - -# --- DB Utility Functions (used by tasks.py and API) --- -def save_task_status(task_id, task_type, status=TASK_STATUS_PENDING, parent_task_id=None, sub_type_identifier=None, progress=0, details=None): - """ - Saves or updates a task's status in the database, using Unix timestamps for start and end times. - """ - db = get_db() # This now calls the function within this file - cur = db.cursor() - current_unix_time = time.time() - - if details is not None and isinstance(details, dict): - # Log truncation logic remains the same - if status != TASK_STATUS_SUCCESS and 'log' in details and isinstance(details['log'], list): - log_list = details['log'] - if len(log_list) > MAX_LOG_ENTRIES_STORED: - original_log_length = len(log_list) - details['log'] = log_list[-MAX_LOG_ENTRIES_STORED:] - details['log_storage_info'] = f"Log in DB truncated to last {MAX_LOG_ENTRIES_STORED} entries. Original length: {original_log_length}." - else: - details.pop('log_storage_info', None) - elif status == TASK_STATUS_SUCCESS: - details.pop('log_storage_info', None) - if 'log' not in details or not isinstance(details.get('log'), list) or not details.get('log'): - details['log'] = ["Task completed successfully."] - - details_json = json.dumps(details) if details is not None else None - - try: - # This query now handles start_time and end_time using Unix timestamps - cur.execute(""" - INSERT INTO task_status (task_id, parent_task_id, task_type, sub_type_identifier, status, progress, details, timestamp, start_time, end_time) - VALUES (%s, %s, %s, %s, %s, %s, %s, NOW(), %s, CASE WHEN %s IN ('SUCCESS', 'FAILURE', 'REVOKED') THEN %s ELSE NULL END) - ON CONFLICT (task_id) DO UPDATE SET - status = EXCLUDED.status, - parent_task_id = EXCLUDED.parent_task_id, - sub_type_identifier = EXCLUDED.sub_type_identifier, - progress = EXCLUDED.progress, - details = EXCLUDED.details, - timestamp = NOW(), - start_time = COALESCE(task_status.start_time, %s), - end_time = CASE - WHEN EXCLUDED.status IN ('SUCCESS', 'FAILURE', 'REVOKED') AND task_status.end_time IS NULL - THEN %s - ELSE task_status.end_time - END - """, (task_id, parent_task_id, task_type, sub_type_identifier, status, progress, details_json, current_unix_time, status, current_unix_time, current_unix_time, current_unix_time)) - db.commit() - except psycopg2.Error as e: - logger.error(f"DB Error saving task status for {task_id}: {e}") - try: - db.rollback() - logger.info(f"DB transaction rolled back for task status update of {task_id}.") - except psycopg2.Error as rb_e: - logger.error(f"DB Error during rollback for task status {task_id}: {rb_e}") - finally: - cur.close() - - # Record persistent history for MAIN tasks that just reached a terminal state. - # Skip the synthetic 'unknown' placeholder inserted by the global cancel - # path (app_helper.cancel_all_jobs) — it has no real type and would show - # up as an 'unknown' row in the dashboard's recent activity table. - try: - if ( - parent_task_id is None - and status in (TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED) - and task_type and task_type != 'unknown' - ): - duration_s = None - try: - hist_cur = db.cursor() - hist_cur.execute( - "SELECT start_time, end_time FROM task_status WHERE task_id = %s", - (task_id,), - ) - row = hist_cur.fetchone() - hist_cur.close() - if row and row[0] is not None: - end = row[1] if row[1] is not None else current_unix_time - duration_s = max(0.0, float(end) - float(row[0])) - except Exception: - pass - record_task_history(task_id, task_type, status, duration_s, details=details) - except Exception as e_hist: - logger.debug(f"history record skipped for {task_id}: {e_hist}") - - -def get_task_info_from_db(task_id): - """Fetches task info from DB and calculates running time in Python.""" - db = get_db() # This now calls the function within this file - cur = db.cursor(cursor_factory=DictCursor) - # Fetch raw columns including the Unix timestamps - cur.execute(""" - SELECT - task_id, parent_task_id, task_type, sub_type_identifier, status, progress, details, timestamp, start_time, end_time - FROM task_status - WHERE task_id = %s - """, (task_id,)) - row = cur.fetchone() - cur.close() - if not row: - return None - - row_dict = dict(row) - current_unix_time = time.time() - - start_time = row_dict.get('start_time') - end_time = row_dict.get('end_time') - - # If start_time is null (old record or pre-start), duration is 0. - if start_time is None: - row_dict['running_time_seconds'] = 0.0 - else: - # If end_time is null, task is running. Use current time. - effective_end_time = end_time if end_time is not None else current_unix_time - row_dict['running_time_seconds'] = max(0, effective_end_time - start_time) - - return row_dict - -def get_child_tasks_from_db(parent_task_id): - """Fetches all child tasks for a given parent_task_id from the database.""" - conn = get_db() # This now calls the function within this file - cur = conn.cursor(cursor_factory=DictCursor) - # MODIFIED: Select the 'details' column as well for the final check. - cur.execute("SELECT task_id, status, sub_type_identifier, details FROM task_status WHERE parent_task_id = %s", (parent_task_id,)) - tasks = cur.fetchall() - cur.close() - # DictCursor returns a list of dictionary-like objects, convert to plain dicts - return [dict(row) for row in tasks] - -def save_track_analysis_and_embedding(item_id, title, author, tempo, key, scale, moods, embedding_vector, energy=None, other_features=None, album=None, album_artist=None, year=None, rating=None, file_path=None): - """Saves track analysis and embedding in a single transaction.""" - - def _sanitize_string(s, max_length=1000, field_name="field"): - """Sanitize string for PostgreSQL insertion.""" - if s is None: - return None - - # Ensure it's a string - if not isinstance(s, str): - try: - s = str(s) - except Exception: - logger.warning(f"Could not convert {field_name} to string, using empty string") - return "" - - # Remove problematic characters - # NUL byte (0x00) - PostgreSQL cannot store - s = s.replace('\x00', '') - - # Remove other control characters that could cause issues - # Keep only printable ASCII, space, tab, newline, and common Unicode - s = ''.join(char for char in s if char.isprintable() or char in '\n\t ') - - # Truncate to max length to prevent overly long strings - if len(s) > max_length: - logger.warning(f"{field_name} truncated from {len(s)} to {max_length} characters") - s = s[:max_length] - - # Strip leading/trailing whitespace - s = s.strip() - - return s - - # Sanitize all string inputs with field-specific limits - title = _sanitize_string(title, max_length=500, field_name="title") - author = _sanitize_string(author, max_length=200, field_name="author") - album = _sanitize_string(album, max_length=200, field_name="album") - album_artist = _sanitize_string(album_artist, max_length=200, field_name="album_artist") - key = _sanitize_string(key, max_length=10, field_name="key") - scale = _sanitize_string(scale, max_length=10, field_name="scale") - other_features = _sanitize_string(other_features, max_length=2000, field_name="other_features") - - # year: parse from various date formats and validate - def _parse_year_from_date(year_value): - """ - Parse year from various date formats. - Supports: YYYY, YYYY-MM-DD, MM-DD-YYYY, DD-MM-YYYY (with - or / separators) - """ - if year_value is None: - return None - - year_str = str(year_value).strip() - if not year_str: - return None - - # Try parsing as pure integer first (YYYY) - try: - year = int(year_str) - if 1000 <= year <= 2100: - return year - except (ValueError, TypeError): - pass - - # Normalize separators - normalized = year_str.replace('/', '-') - parts = normalized.split('-') - - if len(parts) == 3: - try: - # YYYY-MM-DD format - if len(parts[0]) == 4: - year = int(parts[0]) - if 1000 <= year <= 2100: - return year - - # MM-DD-YYYY or DD-MM-YYYY format - if len(parts[2]) == 4: - year = int(parts[2]) - if 1000 <= year <= 2100: - return year - - # 2-digit year (MM-DD-YY) - if len(parts[2]) == 2: - year = int(parts[2]) - year += 2000 if year < 30 else 1900 - if 1000 <= year <= 2100: - return year - except (ValueError, TypeError, IndexError): - pass - - return None - - year = _parse_year_from_date(year) - - # rating: validate as integer 0-5 (5-star rating system) - if rating is not None: - try: - rating = int(rating) - if rating < 0 or rating > 5: - rating = None - except (ValueError, TypeError): - rating = None - - file_path = _sanitize_string(file_path, max_length=1000, field_name="file_path") - - mood_str = ','.join(f"{k}:{v:.3f}" for k, v in moods.items()) - - conn = get_db() # This now calls the function within this file - cur = conn.cursor() - try: - # Save analysis to score table - cur.execute(""" - INSERT INTO score (item_id, title, author, tempo, key, scale, mood_vector, energy, other_features, album, album_artist, year, rating, file_path) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (item_id) DO UPDATE SET - title = EXCLUDED.title, - author = EXCLUDED.author, - tempo = EXCLUDED.tempo, - key = EXCLUDED.key, - scale = EXCLUDED.scale, - mood_vector = EXCLUDED.mood_vector, - energy = EXCLUDED.energy, - other_features = EXCLUDED.other_features, - album = EXCLUDED.album, - album_artist = EXCLUDED.album_artist, - year = EXCLUDED.year, - rating = EXCLUDED.rating, - file_path = EXCLUDED.file_path - """, (item_id, title, author, tempo, key, scale, mood_str, energy, other_features, album, album_artist, year, rating, file_path)) - - # Save embedding - if isinstance(embedding_vector, np.ndarray) and embedding_vector.size > 0: - embedding_blob = embedding_vector.astype(np.float32).tobytes() - cur.execute(""" - INSERT INTO embedding (item_id, embedding) VALUES (%s, %s) - ON CONFLICT (item_id) DO UPDATE SET embedding = EXCLUDED.embedding - """, (item_id, psycopg2.Binary(embedding_blob))) - - conn.commit() - except Exception as e: - conn.rollback() - logger.error("Error saving track analysis and embedding for %s: %s", item_id, e) - raise - finally: - cur.close() - -def save_clap_embedding(item_id, clap_embedding_vector): - """Saves CLAP embedding for a track.""" - if clap_embedding_vector is None or (isinstance(clap_embedding_vector, np.ndarray) and clap_embedding_vector.size == 0): - return - - conn = get_db() - cur = conn.cursor() - try: - embedding_blob = clap_embedding_vector.astype(np.float32).tobytes() - cur.execute(""" - INSERT INTO clap_embedding (item_id, embedding) VALUES (%s, %s) - ON CONFLICT (item_id) DO UPDATE SET embedding = EXCLUDED.embedding - """, (item_id, psycopg2.Binary(embedding_blob))) - conn.commit() - except Exception as e: - conn.rollback() - logger.error(f"Error saving CLAP embedding for {item_id}: {e}") - raise - finally: - cur.close() - - -def get_clap_embedding(item_id): - """Load CLAP embedding for a track from the database. - - Returns: - numpy array (512-dim float32) or None if not found - """ - conn = get_db() - cur = conn.cursor() - try: - cur.execute("SELECT embedding FROM clap_embedding WHERE item_id = %s", (item_id,)) - row = cur.fetchone() - if row and row[0]: - return np.frombuffer(row[0], dtype=np.float32) - return None - except Exception as e: - logger.error(f"Error loading CLAP embedding for {item_id}: {e}") - return None - finally: - cur.close() - - -def save_lyrics_embedding(item_id, lyrics_embedding_vector, axis_vector=None): - """Saves the lyrics embedding (gte-multilingual-base) and the fixed-order axis vector. - - ``axis_vector`` must be a numpy array (float32) already in canonical - MUSIC_ANALYSIS_AXES order (use ``_score_axes`` to produce it). May be None. - """ - if lyrics_embedding_vector is None or (isinstance(lyrics_embedding_vector, np.ndarray) and lyrics_embedding_vector.size == 0): - return - - conn = get_db() - cur = conn.cursor() - try: - embedding_blob = lyrics_embedding_vector.astype(np.float32).tobytes() if isinstance(lyrics_embedding_vector, np.ndarray) else np.asarray(lyrics_embedding_vector, dtype=np.float32).tobytes() - axis_blob = None - if axis_vector is not None: - arr = axis_vector if isinstance(axis_vector, np.ndarray) else np.asarray(axis_vector, dtype=np.float32) - if arr.size > 0: - axis_blob = arr.astype(np.float32, copy=False).tobytes() - cur.execute(""" - INSERT INTO lyrics_embedding (item_id, embedding, axis_vector) VALUES (%s, %s, %s) - ON CONFLICT (item_id) DO UPDATE SET embedding = EXCLUDED.embedding, axis_vector = EXCLUDED.axis_vector, updated_at = CURRENT_TIMESTAMP - """, (item_id, psycopg2.Binary(embedding_blob), - psycopg2.Binary(axis_blob) if axis_blob is not None else None)) - conn.commit() - except Exception as e: - conn.rollback() - logger.error(f"Error saving lyrics embedding for {item_id}: {e}") - raise - finally: - cur.close() - - -def get_tracks_by_ids(item_ids_list): - """Fetches full track data (including embeddings) for a specific list of item_ids.""" - if not item_ids_list: - return [] - conn = get_db() # This now calls the function within this file - cur = conn.cursor(cursor_factory=DictCursor) - - # Convert item_ids to strings to match the text type in database - item_ids_str = [str(item_id) for item_id in item_ids_list] - - query = """ - SELECT s.item_id, s.title, s.author, s.album, s.album_artist, s.tempo, s.key, s.scale, s.mood_vector, s.energy, s.other_features, s.year, s.rating, s.file_path, e.embedding - FROM score s - LEFT JOIN embedding e ON s.item_id = e.item_id - WHERE s.item_id IN %s - """ - cur.execute(query, (tuple(item_ids_str),)) - rows = cur.fetchall() - cur.close() - - # Convert DictRow objects to regular dicts to allow adding new keys. - processed_rows = [] - for row in rows: - row_dict = dict(row) - if row_dict.get('embedding'): - row_dict['embedding_vector'] = np.frombuffer(row_dict['embedding'], dtype=np.float32) - else: - row_dict['embedding_vector'] = np.array([]) - processed_rows.append(row_dict) - - return processed_rows - -def get_score_data_by_ids(item_ids_list): - """Fetches only score-related data (excluding embeddings) for a specific list of item_ids.""" - if not item_ids_list: - return [] - conn = get_db() # This now calls the function within this file - cur = conn.cursor(cursor_factory=DictCursor) - query = """ - SELECT s.item_id, s.title, s.author, s.album, s.album_artist, s.tempo, s.key, s.scale, s.mood_vector, s.energy, s.other_features, s.year, s.rating, s.file_path - FROM score s - WHERE s.item_id IN %s - """ - try: - cur.execute(query, (tuple(item_ids_list),)) - rows = cur.fetchall() - except Exception as e: - logger.error(f"Error fetching score data by IDs: {e}") - rows = [] # Return empty list on error - finally: - cur.close() - return [dict(row) for row in rows] +# The Flask `app` object is intentionally NOT imported here (circular import); +# use the module-level `logger` above. The 2D map/artist projection caches live +# in database.MAP_PROJECTION_CACHE / database.ARTIST_PROJECTION_CACHE, written by +# the build_and_store_* helpers below and read by database.load_*_projection. def top_stratified_genre(mood_vector): @@ -1136,280 +101,24 @@ def attach_song_features(rows, id_key='item_id'): return rows -def save_alchemy_anchor(name, centroid): - """Save a named anchor centroid into DB.""" - if not name or not centroid or not isinstance(centroid, list): - raise ValueError('Anchor name and centroid list are required.') - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - centroid_json = json.dumps(centroid) - cur.execute( - "INSERT INTO alchemy_anchors (name, centroid) VALUES (%s, %s) " - "ON CONFLICT (name) DO UPDATE SET centroid = EXCLUDED.centroid, created_at = NOW() " - "RETURNING id, name, created_at", - (name, centroid_json) - ) - row = cur.fetchone() - conn.commit() - return dict(row) if row else None - except Exception as e: - conn.rollback() - logger.error(f"Failed to save alchemy anchor '{name}': {e}") - return None - finally: - cur.close() -def get_alchemy_anchors(): - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - cur.execute("SELECT id, name, created_at FROM alchemy_anchors ORDER BY created_at DESC") - rows = cur.fetchall() - return [dict(row) for row in rows] - except Exception as e: - logger.error(f"Failed to load alchemy anchors: {e}") - return [] - finally: - cur.close() -def delete_alchemy_anchor(anchor_id): - conn = get_db() - cur = conn.cursor() - try: - cur.execute("DELETE FROM alchemy_anchors WHERE id = %s", (anchor_id,)) - conn.commit() - return cur.rowcount > 0 - except Exception as e: - conn.rollback() - logger.error(f"Failed to delete alchemy anchor id={anchor_id}: {e}") - return False - finally: - cur.close() -def get_alchemy_anchor_by_id(anchor_id): - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - cur.execute("SELECT id, name, centroid, created_at FROM alchemy_anchors WHERE id = %s", (anchor_id,)) - row = cur.fetchone() - if not row: - return None - anchor = dict(row) - if isinstance(anchor.get('centroid'), str): - try: - anchor['centroid'] = json.loads(anchor['centroid']) - except Exception: - anchor['centroid'] = None - return anchor - except Exception as e: - logger.error(f"Failed to fetch alchemy anchor id={anchor_id}: {e}") - return None - finally: - cur.close() -def update_alchemy_anchor_name(anchor_id, name): - if not name or not isinstance(name, str): - return None - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - cur.execute( - "UPDATE alchemy_anchors SET name = %s WHERE id = %s RETURNING id, name", - (name.strip(), anchor_id) - ) - row = cur.fetchone() - conn.commit() - if not row: - return None - return dict(row) - except Exception as e: - conn.rollback() - logger.error(f"Failed to rename alchemy anchor id={anchor_id}: {e}") - return None - finally: - cur.close() -def get_alchemy_radios(): - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - cur.execute( - "SELECT r.id, r.anchor_id, a.name, r.temperature, r.n_results, r.enabled " - "FROM alchemy_radios r JOIN alchemy_anchors a ON a.id = r.anchor_id " - "ORDER BY a.name" - ) - rows = cur.fetchall() - return [dict(row) for row in rows] - except Exception: - logger.exception("Failed to load alchemy radios") - return [] - finally: - cur.close() -def create_alchemy_radio(anchor_id, temperature, n_results, enabled=True): - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - cur.execute( - "INSERT INTO alchemy_radios (anchor_id, temperature, n_results, enabled) " - "VALUES (%s, %s, %s, %s) RETURNING id, anchor_id, temperature, n_results, enabled", - (anchor_id, temperature, n_results, bool(enabled)) - ) - row = cur.fetchone() - conn.commit() - return dict(row) if row else None - except Exception: - conn.rollback() - logger.exception(f"Failed to create alchemy radio for anchor_id={anchor_id}") - return None - finally: - cur.close() -def update_alchemy_radio(radio_id, temperature, n_results, enabled): - conn = get_db() - cur = conn.cursor(cursor_factory=DictCursor) - try: - cur.execute( - "UPDATE alchemy_radios SET temperature = %s, n_results = %s, enabled = %s " - "WHERE id = %s RETURNING id, anchor_id, temperature, n_results, enabled", - (temperature, n_results, bool(enabled), radio_id) - ) - row = cur.fetchone() - conn.commit() - return dict(row) if row else None - except Exception: - conn.rollback() - logger.exception(f"Failed to update alchemy radio id={radio_id}") - return None - finally: - cur.close() -def delete_alchemy_radio(radio_id): - conn = get_db() - cur = conn.cursor() - try: - cur.execute("DELETE FROM alchemy_radios WHERE id = %s", (radio_id,)) - conn.commit() - return cur.rowcount > 0 - except Exception: - conn.rollback() - logger.exception(f"Failed to delete alchemy radio id={radio_id}") - return False - finally: - cur.close() -def save_map_projection(index_name, id_map, projection_array): - """ - Save a precomputed 2D projection into the map_projection_data table. - projection_array: numpy array of shape (N,2), dtype=float32 - id_map: JSON-serializable list/dict mapping rows to item_ids - """ - conn = get_db() - try: - blob = projection_array.astype(np.float32).tobytes() - if not blob: - logger.info(f"Map projection '{index_name}' has no data; clearing existing store.") - with conn.cursor() as cur: - cur.execute( - "DELETE FROM map_projection_data WHERE index_name = %s OR index_name LIKE %s ESCAPE '\\'", - (index_name, index_name.replace('_', r'\_') + r"\_%\_%"), - ) - conn.commit() - return - embedding_dim = projection_array.shape[1] if projection_array.ndim == 2 else 0 - from tasks.index_build_helpers import store_voyager_index_segmented - store_voyager_index_segmented( - conn, - target_table="map_projection_data", - index_name=index_name, - index_bytes=blob, - id_map=id_map, - embedding_dimension=embedding_dim, - binary_column="projection_data", - ) - conn.commit() - try: - id_count = len(id_map) if hasattr(id_map, '__len__') else None - logger.info(f"Saved map projection '{index_name}' to DB: {len(blob)} bytes, ids={id_count}") - except Exception: - logger.debug("Saved map projection but failed to compute size/id_count for log.") - except Exception as e: - conn.rollback() - logger.error(f"Failed to save map projection: {e}") - raise - - -def load_map_projection(index_name, force_reload=False): - """Load precomputed projection from DB. Returns (id_map, numpy_array) or (None, None)""" - global MAP_PROJECTION_CACHE - # Try cache first (unless force_reload is True) - if not force_reload and MAP_PROJECTION_CACHE and MAP_PROJECTION_CACHE.get('index_name') == index_name: - logger.info(f"Map projection '{index_name}' already loaded in cache. Skipping reload.") - return MAP_PROJECTION_CACHE.get('id_map'), MAP_PROJECTION_CACHE.get('projection') - - logger.info(f"Attempting to load map projection '{index_name}' from database into memory...") - conn = get_db() - cur = conn.cursor() - try: - cur.execute("SELECT projection_data, id_map_json FROM map_projection_data WHERE index_name = %s", (index_name,)) - row = cur.fetchone() - if row and row[0] is not None: - proj_blob, id_map_json = row[0], row[1] - else: - import re - from tasks.index_build_helpers import reassemble_segmented_id_map - cur.execute( - "SELECT index_name, projection_data, id_map_json FROM map_projection_data WHERE index_name LIKE %s ESCAPE '\\'", - (index_name.replace('_', r'\_') + r"\_%\_%",), - ) - candidates = cur.fetchall() - if not candidates: - logger.warning(f"Map projection '{index_name}' not found in the database. Cache will be empty.") - return None, None - seg_pattern = re.compile(rf"^{re.escape(index_name)}_(\d+)_(\d+)$") - parts = [] - total_expected = None - for name, part_blob, part_id_map in candidates: - m = seg_pattern.match(name) - if not m: - continue - part_no = int(m.group(1)) - total = int(m.group(2)) - if total_expected is None: - total_expected = total - elif total_expected != total: - logger.error(f"Map projection segment total mismatch for '{index_name}' ({total_expected} vs {total}). Aborting load.") - return None, None - parts.append((part_no, part_blob, part_id_map)) - if total_expected is None or len(parts) != total_expected: - logger.error(f"Incomplete map projection segments for '{index_name}': expected {total_expected}, found {len(parts)}. Aborting load.") - return None, None - parts.sort(key=lambda p: p[0]) - proj_blob = b"".join(bytes(p[1]) for p in parts if p[1]) - id_map_json = reassemble_segmented_id_map((p[0], p[2]) for p in parts) - proj = np.frombuffer(proj_blob, dtype=np.float32) - # infer shape as (-1,2) if length divisible by 2 - if proj.size % 2 == 0: - proj = proj.reshape((-1, 2)) - id_map = json.loads(id_map_json) - MAP_PROJECTION_CACHE = {'index_name': index_name, 'id_map': id_map, 'projection': proj} - logger.info(f"Map projection '{index_name}' with {len(id_map)} items loaded successfully into memory.") - return id_map, proj - except Exception as e: - logger.error(f"Failed to load map projection: {e}", exc_info=True) - return None, None - finally: - cur.close() def build_and_store_map_projection(index_name='main_map'): @@ -1467,9 +176,8 @@ def build_and_store_map_projection(index_name='main_map'): # Save to DB try: save_map_projection(index_name, ids, projections) - # update in-memory cache - global MAP_PROJECTION_CACHE - MAP_PROJECTION_CACHE = {'index_name': index_name, 'id_map': ids, 'projection': projections} + # Update the canonical in-memory cache (read by database.load_map_projection). + database.MAP_PROJECTION_CACHE = {'index_name': index_name, 'id_map': ids, 'projection': projections} # Note: Caller (analysis task) is responsible for publishing reload message after all builds complete return True except Exception as e: @@ -1477,60 +185,8 @@ def build_and_store_map_projection(index_name='main_map'): return False -def load_artist_projection(index_name='artist_map', force_reload=False): - """Load precomputed artist component projection from DB. - Returns (artist_component_map, numpy_array) or (None, None). - artist_component_map format: [{'artist_id': '...', 'component_idx': 0, 'weight': 0.3}, ...] - """ - global ARTIST_PROJECTION_CACHE - # Try cache first (unless force_reload is True) - if not force_reload and ARTIST_PROJECTION_CACHE and ARTIST_PROJECTION_CACHE.get('index_name') == index_name: - logger.info(f"Artist projection '{index_name}' already loaded in cache. Skipping reload.") - return ARTIST_PROJECTION_CACHE.get('component_map'), ARTIST_PROJECTION_CACHE.get('projection') - - logger.info(f"Attempting to load artist projection '{index_name}' from database into memory...") - conn = get_db() - cur = conn.cursor() - try: - cur.execute("SELECT projection_data, artist_component_map_json FROM artist_component_projection WHERE index_name = %s", (index_name,)) - row = cur.fetchone() - if not row: - logger.warning(f"Artist projection '{index_name}' not found in the database. Cache will be empty.") - return None, None - proj_blob, component_map_json = row[0], row[1] - proj = np.frombuffer(proj_blob, dtype=np.float32) - # infer shape as (-1,2) if length divisible by 2 - if proj.size % 2 == 0: - proj = proj.reshape((-1, 2)) - component_map = json.loads(component_map_json) - ARTIST_PROJECTION_CACHE = {'index_name': index_name, 'component_map': component_map, 'projection': proj} - logger.info(f"Artist projection '{index_name}' with {len(component_map)} components loaded successfully into memory.") - return component_map, proj - except Exception as e: - logger.error(f"Failed to load artist projection: {e}", exc_info=True) - return None, None - finally: - cur.close() -def save_artist_projection(index_name, component_map, projections): - """Save artist component projection to database. - component_map: [{'artist_id': '...', 'component_idx': 0, 'weight': 0.3}, ...] - projections: numpy array of shape (N, 2) - """ - conn = get_db() - cur = conn.cursor() - try: - component_map_json = json.dumps(component_map) - proj_blob = projections.astype(np.float32).tobytes() - cur.execute("INSERT INTO artist_component_projection (index_name, projection_data, artist_component_map_json) VALUES (%s, %s, %s) ON CONFLICT (index_name) DO UPDATE SET projection_data = EXCLUDED.projection_data, artist_component_map_json = EXCLUDED.artist_component_map_json, created_at = CURRENT_TIMESTAMP", (index_name, proj_blob, component_map_json)) - conn.commit() - logger.info(f"Saved artist projection '{index_name}' with {len(component_map)} components to database.") - except Exception as e: - conn.rollback() - logger.error(f"Failed to save artist projection: {e}", exc_info=True) - finally: - cur.close() def build_and_store_artist_projection(index_name='artist_map'): @@ -1619,9 +275,8 @@ def build_and_store_artist_projection(index_name='artist_map'): try: save_artist_projection(index_name, component_map, projections) - # Update in-memory cache - global ARTIST_PROJECTION_CACHE - ARTIST_PROJECTION_CACHE = {'index_name': index_name, 'component_map': component_map, 'projection': projections} + # Update the canonical in-memory cache (read by database.load_artist_projection). + database.ARTIST_PROJECTION_CACHE = {'index_name': index_name, 'component_map': component_map, 'projection': projections} # Note: Caller (analysis task) is responsible for publishing reload message after all builds complete return True except Exception as e: @@ -1629,21 +284,6 @@ def build_and_store_artist_projection(index_name='artist_map'): return False -def update_playlist_table(playlists): # Removed db_path - conn = get_db() # This now calls the function within this file - cur = conn.cursor() - try: - # Clear all previous conceptual playlists to reflect only the current run. - cur.execute("DELETE FROM playlist") - for name, cluster in playlists.items(): - for item_id, title, author in cluster: - cur.execute("INSERT INTO playlist (playlist_name, item_id, title, author) VALUES (%s, %s, %s, %s) ON CONFLICT (playlist_name, item_id) DO NOTHING", (name, item_id, title, author)) - conn.commit() - except Exception as e: - conn.rollback() - logger.error("Error updating playlist table: %s", e) - finally: - cur.close() def cancel_job_and_children_recursive(job_id, task_type_from_db=None, reason="Task cancellation processed by API."): """Helper to cancel a job and its children based on DB records. diff --git a/app_helper_artist.py b/app_helper_artist.py index f285c944..970d0e1d 100644 --- a/app_helper_artist.py +++ b/app_helper_artist.py @@ -6,7 +6,7 @@ import logging from database import get_db -from tasks.memory_utils import sanitize_string_for_db +from sanitization import sanitize_string_for_db logger = logging.getLogger(__name__) diff --git a/app_logging.py b/app_logging.py index 314731cd..e637311b 100644 --- a/app_logging.py +++ b/app_logging.py @@ -11,70 +11,94 @@ ``logger.info(...)`` from task modules fell through to Python's ``lastResort`` handler — silently dropping INFO-level output during long-running jobs. -Emoji safety: the ``EmojiStrippingFilter`` removes emoji and other non-Latin-1 -symbols from every log record before it reaches a handler. This prevents -``UnicodeEncodeError`` / ``UnicodeDecodeError`` crashes on Windows when stdout -is a pipe (PyInstaller native build) or when the console code-page cannot -represent the character. HTML templates and web-UI progress messages are -unaffected — only the Python ``logging`` pipeline is sanitised. +Record sanitization: the ``LogSanitizingFilter`` cleans every log record before +it reaches a handler. It removes emoji / non-Latin-1 symbols (which raise +``UnicodeEncodeError`` / ``UnicodeDecodeError`` on Windows when stdout is a pipe +or the console code-page cannot represent the character) and neutralizes CR/LF +and other control characters so an attacker-controlled value embedded in a +message cannot forge or split log lines (CWE-117, log injection). This is the +single, centralized place where log-message sanitization happens — call sites +log the raw value and the filter cleans it. HTML templates and web-UI progress +messages are unaffected — only the Python ``logging`` pipeline is sanitised, and +the traceback ``logger.exception`` appends is left intact (the formatter renders +it after filtering, so the full error is always visible in the log). """ import logging import re +from typing import Any LOG_FORMAT = "[%(levelname)s]-[%(asctime)s]-%(message)s" LOG_DATEFMT = "%d-%m-%Y %H-%M-%S" # --------------------------------------------------------------------------- -# Emoji / symbol stripping for console-safe logging +# Console-safe + injection-safe log record sanitization # --------------------------------------------------------------------------- # Ranges cover all common emoji blocks plus Dingbats, Misc Symbols, # Geometric Shapes, Supplemental Symbols, and variation selectors. -# Characters within Latin-1 (U+0000–U+00FF) are *not* stripped, so -# European accented letters (e.g. é, ñ, ü) pass through unchanged. +# Characters within Latin-1 (U+0000-U+00FF) are *not* stripped, so +# European accented letters pass through unchanged. _EMOJI_RE = re.compile( "[\U0001F300-\U0001F9FF" # Misc Symbols, Emoticons, Transport, Supplemental "\U0001FA00-\U0001FAFF" # Chess Symbols, Symbols Extended-A - "\U00002190-\U000027BF" # Arrows (→ ← ↑ ↓ ↔), Misc Technical, Dingbats (✓ ✗ ✕ ★ ☆ ♯ ♭ etc.) - "\U000025A0-\U000025FF" # Geometric Shapes (● ○ ■ □ ◆ ◇ ▲ ▼ etc.) + "\U00002190-\U000027BF" # Arrows, Misc Technical, Dingbats + "\U000025A0-\U000025FF" # Geometric Shapes "\U00002B00-\U00002BFF" # Misc Symbols & Arrows "\U0001F000-\U0001F02F" # Mahjong Tiles "\U0001F0A0-\U0001F0FF" # Playing Cards - "\uFE0F\u200D" # Variation Selector-16, Zero-Width Joiner + "\\uFE0F\\u200D" # Variation Selector-16, Zero-Width Joiner "]+" ) +# Control codes plus DEL and the Unicode line/paragraph separators, EXCEPT tab +# (0x09). Includes LF (0x0A), CR (0x0D), NEL (U+0085), LINE SEPARATOR (U+2028) +# and PARAGRAPH SEPARATOR (U+2029): replacing these with a space prevents a +# value with an embedded line break from forging additional log lines, including +# for consumers (and Windows code-pages) that treat the Unicode separators as +# line boundaries. +_CONTROL_RE = re.compile(r"[\x00-\x08\x0A-\x1F\x7F\x85" + chr(0x2028) + chr(0x2029) + "]") -def _strip_emoji(text: str) -> str: - """Remove emoji and symbol characters from *text*, returning a plain string.""" + +def _sanitize_log_text(text: Any) -> Any: + """Make *text* safe for a single console log line. + + Strips emoji/symbol characters (Windows code-page safety) and replaces CR/LF + and other C0 control codes with a space so an attacker-controlled value + cannot forge or split log lines (CWE-117). Tabs are preserved. Non-string + values (e.g. numeric ``logging`` args) are returned unchanged. + """ if not isinstance(text, str): return text cleaned = _EMOJI_RE.sub("", text) - # Collapse multiple spaces that may result from removing a symbol + cleaned = _CONTROL_RE.sub(" ", cleaned) + # Collapse runs of spaces left behind by removed symbols / control codes. return re.sub(r" {2,}", " ", cleaned).strip() -class EmojiStrippingFilter(logging.Filter): - """Logging filter that strips emoji/symbols from ``record.msg`` and ``record.args``. +class LogSanitizingFilter(logging.Filter): + """Logging filter that sanitizes ``record.msg`` and ``record.args``. - Attach this to the root logger's handlers so every log record is sanitised - before it reaches a ``StreamHandler`` (console / pipe). File-based handlers + Attach this to the root logger's handlers so every log record is cleaned + before it reaches a ``StreamHandler`` (console / pipe): emoji/symbols are + removed and CR/LF + control characters are neutralised. File-based handlers with ``propagate=False`` (e.g. the Windows supervisor's own log) are not - affected. + affected. The exception traceback added by ``logger.exception`` lives in + ``record.exc_info`` and is rendered by the formatter after this filter runs, + so it is never altered here — the full error always reaches the log. """ def filter(self, record): if isinstance(record.msg, str): - record.msg = _strip_emoji(record.msg) + record.msg = _sanitize_log_text(record.msg) if record.args: if isinstance(record.args, dict): record.args = { - k: _strip_emoji(v) if isinstance(v, str) else v + k: _sanitize_log_text(v) if isinstance(v, str) else v for k, v in record.args.items() } elif isinstance(record.args, (list, tuple)): record.args = tuple( - _strip_emoji(a) if isinstance(a, str) else a + _sanitize_log_text(a) if isinstance(a, str) else a for a in record.args ) return True @@ -83,10 +107,11 @@ def filter(self, record): def configure_logging(level: int = logging.INFO) -> None: """Install the project-wide root logger format. Idempotent. - An ``EmojiStrippingFilter`` is attached to every handler on the root logger, - making all console / pipe output safe on Windows regardless of code-page. + A ``LogSanitizingFilter`` is attached to every handler on the root logger, + making all console / pipe output safe on Windows regardless of code-page and + neutralising log-injection attempts from untrusted message data. """ logging.basicConfig(level=level, format=LOG_FORMAT, datefmt=LOG_DATEFMT) for handler in logging.root.handlers: - if not any(isinstance(f, EmojiStrippingFilter) for f in handler.filters): - handler.addFilter(EmojiStrippingFilter()) + if not any(isinstance(f, LogSanitizingFilter) for f in handler.filters): + handler.addFilter(LogSanitizingFilter()) diff --git a/app_path.py b/app_path.py index 11a59fef..7a104b16 100644 --- a/app_path.py +++ b/app_path.py @@ -73,7 +73,7 @@ def _resolve_mood_to_song_id(mood, other_song_id, pct=100): def _resolve_anchor_to_song_id(anchor_id, other_song_id=None, pct=100): - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id try: anchor = get_alchemy_anchor_by_id(int(anchor_id)) except Exception: diff --git a/app_provider_migration.py b/app_provider_migration.py index b4942d56..40305701 100644 --- a/app_provider_migration.py +++ b/app_provider_migration.py @@ -23,7 +23,7 @@ # anything in. from database import get_db from taskqueue import redis_conn, rq_queue_high -from app_helper import validate_outbound_url +from ssrf_guard import validate_outbound_url from tasks.mediaserver.helper import detect_path_format as _detect_path_format logger = logging.getLogger(__name__) @@ -66,7 +66,7 @@ def __getattr__(self, name): # --------------------------------------------------------------------------- # SSRF guard for the user-supplied media-server URL. Delegates to the shared -# ``app_helper.validate_outbound_url`` (allows LAN/loopback, blocks non-HTTP(S) +# ``ssrf_guard.validate_outbound_url`` (allows LAN/loopback, blocks non-HTTP(S) # schemes and link-local/cloud-metadata). A missing url is allowed and left to # the downstream probe. # --------------------------------------------------------------------------- @@ -1672,10 +1672,8 @@ def _load_state(session_id): def _sanitize_json_value(value): - """Wrapper kept for backward compatibility — delegates to the shared - sanitizer in :mod:`tasks.memory_utils`. - """ - from tasks.memory_utils import sanitize_json_for_db + """Local alias for the shared JSON sanitizer in :mod:`sanitization`.""" + from sanitization import sanitize_json_for_db return sanitize_json_for_db(value) diff --git a/app_setup.py b/app_setup.py index 347ab4cc..cfcb8012 100644 --- a/app_setup.py +++ b/app_setup.py @@ -5,7 +5,7 @@ from flask_app import app from tasks.setup_manager import setup_manager from app_auth import check_setup_needed -from app_helper import validate_outbound_url +from ssrf_guard import validate_outbound_url import restart_manager import tasks.mediaserver as mediaserver from error import error_manager diff --git a/app_voyager.py b/app_voyager.py index 04253483..0796e9fc 100644 --- a/app_voyager.py +++ b/app_voyager.py @@ -367,7 +367,7 @@ def get_similar_tracks_endpoint(): # --- Anchor mode: use anchor's centroid vector --- if anchor_id_param is not None: - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id anchor = get_alchemy_anchor_by_id(anchor_id_param) if not anchor or not anchor.get('centroid'): return jsonify({"error": f"Anchor with id {anchor_id_param} not found or has no centroid."}), 404 diff --git a/config.py b/config.py index e3350f4c..60a5f93b 100644 --- a/config.py +++ b/config.py @@ -1,6 +1,16 @@ #AudioMuse-AI/config.py import os +# --- Task Status Constants --- +# These are used across the application for task tracking. Placed here so they're +# available everywhere without creating import chains. +TASK_STATUS_PENDING = 'PENDING' +TASK_STATUS_STARTED = 'STARTED' +TASK_STATUS_PROGRESS = 'PROGRESS' +TASK_STATUS_SUCCESS = 'SUCCESS' +TASK_STATUS_FAILURE = 'FAILURE' +TASK_STATUS_REVOKED = 'REVOKED' + # --- Media Server Type --- MEDIASERVER_TYPE = os.environ.get("MEDIASERVER_TYPE", "jellyfin").lower() # Possible values: jellyfin, navidrome, lyrion, emby @@ -573,7 +583,16 @@ def _compute_headers(): TEMPO_MAX_BPM = float(os.getenv("TEMPO_MAX_BPM", "200.0")) OTHER_FEATURE_LABELS = ['danceable', 'aggressive', 'happy', 'party', 'relaxed', 'sad'] -# Redis cache key for CLAP text embeddings of OTHER_FEATURE_LABELS +# Voice vocabulary used in MCP system prompts +VOICE_VOCAB = ["female vocalists", "female vocalist", "male vocalists"] + +# Fallback genre list used when library context has no top genres +AI_FALLBACK_GENRES = ( + "rock, pop, metal, jazz, electronic, dance, alternative, indie, punk, blues, " + "hard rock, heavy metal, hip-hop, funk, country, soul" +) + +# Redis cache key for CLAP text embeddings CLAP_OTHER_FEATURES_REDIS_KEY = os.environ.get("CLAP_OTHER_FEATURES_REDIS_KEY", "audiomuse:clap_other_feature_text_embeddings") # --- Sonic Fingerprint Constants --- @@ -645,6 +664,17 @@ def _compute_headers(): MAX_SONGS_PER_ARTIST_PLAYLIST = int(os.environ.get("MAX_SONGS_PER_ARTIST_PLAYLIST", "5")) # Enable energy-arc shaping for playlist ordering (gentle start -> peak -> cool down) PLAYLIST_ENERGY_ARC = os.environ.get("PLAYLIST_ENERGY_ARC", "False").lower() == "true" + +# --- Instant Playlist AI Brainstorm --- +AI_BRAINSTORM_SOUND_DESCRIPTIONS_MAX = int(os.environ.get("AI_BRAINSTORM_SOUND_DESCRIPTIONS_MAX", "3")) +AI_BRAINSTORM_SEED_ARTISTS_MAX = int(os.environ.get("AI_BRAINSTORM_SEED_ARTISTS_MAX", "4")) +AI_BRAINSTORM_USE_ARTIST_SEEDS = os.environ.get("AI_BRAINSTORM_USE_ARTIST_SEEDS", "true").lower() == "true" +AI_BRAINSTORM_SIMILAR_ARTISTS_PER_SEED = int(os.environ.get("AI_BRAINSTORM_SIMILAR_ARTISTS_PER_SEED", "8")) +AI_BRAINSTORM_LYRIC_THEMES_MAX = int(os.environ.get("AI_BRAINSTORM_LYRIC_THEMES_MAX", "2")) +AI_BRAINSTORM_GENRE_SCORE_THRESHOLD = float(os.environ.get("AI_BRAINSTORM_GENRE_SCORE_THRESHOLD", "0.3")) +AI_BRAINSTORM_POOL_FLOOR = int(os.environ.get("AI_BRAINSTORM_POOL_FLOOR", "40")) +AI_BRAINSTORM_RELAX_YEAR_PAD = int(os.environ.get("AI_BRAINSTORM_RELAX_YEAR_PAD", "5")) + # --- Authentication --- # Set all three to enable authentication. Leave any blank to disable (legacy mode). AUDIOMUSE_USER = os.environ.get("AUDIOMUSE_USER", "") diff --git a/database.py b/database.py index cfb96182..ab09e608 100644 --- a/database.py +++ b/database.py @@ -11,15 +11,44 @@ first (pgserver), exports the resulting DSN as ``DATABASE_URL``, then boots the app. """ +import json import logging +import sys +import time +import numpy as np import psycopg2 from flask import g +from psycopg2.extras import DictCursor import config logger = logging.getLogger(__name__) +# UTC "now" SQL fragment is owned by tz_helper, a leaf module -- importing it +# adds no depth to the eager import graph, so there is no need to duplicate it. +from tz_helper import UTC_NOW_SQL + +# Shared input sanitizer (leaf module) for cleaning string columns before writes. +from sanitization import sanitize_db_field + +# Task status constants (imported from config, re-exported for backward compatibility) +from config import ( + TASK_STATUS_PENDING, + TASK_STATUS_STARTED, + TASK_STATUS_PROGRESS, + TASK_STATUS_SUCCESS, + TASK_STATUS_FAILURE, + TASK_STATUS_REVOKED, +) + +# Task history constants +TASK_HISTORY_MAX_ROWS = 10 +MAX_LOG_ENTRIES_STORED = 10 # Max number of recent log entries to store in the database per task + +# In-memory cache for the precomputed 2D map projection (optional) +MAP_PROJECTION_CACHE = None + _embedded_server = None @@ -98,3 +127,1335 @@ def stop_embedded(): if _embedded_server is not None: _embedded_server.cleanup() _embedded_server = None + + +# --------------------------------------------------------------------------- +# Task status and history operations +# --------------------------------------------------------------------------- + +def _build_task_note(task_type, details_obj, db): + """Build a short, human-readable note for a finished task. + + Looks at the ``details`` JSON we stored on the main task and, when needed, + queries subtasks to compute a meaningful number (e.g. total songs analyzed + across all album_analysis subtasks).""" + if not isinstance(details_obj, dict): + details_obj = {} + t = (task_type or '').lower() + + try: + if 'analysis' in t: + # Prefer summing tracks_analyzed from album_analysis subtasks. + try: + with db.cursor() as cur: + cur.execute( + "SELECT details FROM task_status WHERE parent_task_id = %s AND status = 'SUCCESS'", + (details_obj.get('_task_id') or '',), + ) + rows = cur.fetchall() + except Exception: + rows = [] + songs = 0 + for (d,) in rows or []: + if not d: + continue + try: + obj = json.loads(d) + if isinstance(obj, dict): + v = obj.get('tracks_analyzed') + if isinstance(v, (int, float)): + songs += int(v) + except Exception: + continue + if songs > 0: + return f"Songs analyzed: {songs}" + # Fallback to album-level info from the main task details. + albums = details_obj.get('albums_completed') or details_obj.get('total_albums_processed') + if albums: + return f"Albums analyzed: {albums}" + return '' + + if 'clean' in t: + for k in ('tracks_deleted', 'orphans_removed', 'songs_cleaned', + 'tracks_removed', 'deleted_count', 'cleaned_tracks'): + v = details_obj.get(k) + if isinstance(v, (int, float)): + return f"Songs cleaned: {int(v)}" + return '' + + if 'cluster' in t: + sampled = (details_obj.get('best_params') or {}).get('initial_subset_size') \ + if isinstance(details_obj.get('best_params'), dict) else None + if sampled is None: + sampled = details_obj.get('sampled_songs') or details_obj.get('num_sampled_songs') + n_clusters = details_obj.get('num_playlists_created') or details_obj.get('num_clusters') + parts = [] + if sampled: + parts.append(f"sampled: {int(sampled)}") + if n_clusters: + parts.append(f"clusters: {int(n_clusters)}") + return ' | '.join(parts) + except Exception as e: + logger.debug(f"task note builder failed for type={task_type}: {e}") + return '' + + +def record_task_history(task_id, task_type, status, duration_seconds=None, note=None, details=None): + """Insert a row into ``task_history`` and trim the table to the most + recent ``TASK_HISTORY_MAX_ROWS`` entries. + + Safe to call from anywhere; never raises. ``details`` (dict or None) is + used to build a default ``note`` when one is not provided explicitly. + If a short note cannot be inferred, fall back to the task's final + status_message or message text when available. + + The history table is treated as immutable per task_id: once a task has + been recorded, we do not insert a second history row for the same task. + """ + if not task_id: + return + try: + db = get_db() + # If no note was supplied, try to infer one from details. + if note is None: + details_obj = details if isinstance(details, dict) else {} + # Pass task_id through so the analysis branch can query subtasks. + details_obj = dict(details_obj) + details_obj['_task_id'] = task_id + note = _build_task_note(task_type, details_obj, db) or '' + if not note: + note = details_obj.get('status_message') or details_obj.get('message') or '' + + with db.cursor() as cur: + cur.execute( + "SELECT 1 FROM task_history WHERE task_id = %s LIMIT 1", + (task_id,) + ) + if cur.fetchone(): + return + cur.execute( + f""" + INSERT INTO task_history (task_id, task_type, status, duration_seconds, note, recorded_at) + VALUES (%s, %s, %s, %s, %s, {UTC_NOW_SQL}) + """, + (task_id, task_type, status, duration_seconds, note), + ) + # Trim — keep only the most recent rows. + cur.execute( + """ + DELETE FROM task_history + WHERE id NOT IN ( + SELECT id FROM task_history ORDER BY recorded_at DESC, id DESC LIMIT %s + ) + """, + (TASK_HISTORY_MAX_ROWS,), + ) + db.commit() + except Exception as e: + logger.warning(f"record_task_history failed for {task_id}: {e}") + try: + db.rollback() + except Exception: + pass + + +def save_task_status(task_id, task_type, status=TASK_STATUS_PENDING, parent_task_id=None, sub_type_identifier=None, progress=0, details=None): + """ + Saves or updates a task's status in the database, using Unix timestamps for start and end times. + """ + db = get_db() + cur = db.cursor() + current_unix_time = time.time() + + if details is not None and isinstance(details, dict): + # Log truncation logic remains the same + if status != TASK_STATUS_SUCCESS and 'log' in details and isinstance(details['log'], list): + log_list = details['log'] + if len(log_list) > MAX_LOG_ENTRIES_STORED: + original_log_length = len(log_list) + details['log'] = log_list[-MAX_LOG_ENTRIES_STORED:] + details['log_storage_info'] = f"Log in DB truncated to last {MAX_LOG_ENTRIES_STORED} entries. Original length: {original_log_length}." + else: + details.pop('log_storage_info', None) + elif status == TASK_STATUS_SUCCESS: + details.pop('log_storage_info', None) + if 'log' not in details or not isinstance(details.get('log'), list) or not details.get('log'): + details['log'] = ["Task completed successfully."] + + details_json = json.dumps(details) if details is not None else None + + try: + # This query now handles start_time and end_time using Unix timestamps + cur.execute(""" + INSERT INTO task_status (task_id, parent_task_id, task_type, sub_type_identifier, status, progress, details, timestamp, start_time, end_time) + VALUES (%s, %s, %s, %s, %s, %s, %s, NOW(), %s, CASE WHEN %s IN ('SUCCESS', 'FAILURE', 'REVOKED') THEN %s ELSE NULL END) + ON CONFLICT (task_id) DO UPDATE SET + status = EXCLUDED.status, + parent_task_id = EXCLUDED.parent_task_id, + sub_type_identifier = EXCLUDED.sub_type_identifier, + progress = EXCLUDED.progress, + details = EXCLUDED.details, + timestamp = NOW(), + start_time = COALESCE(task_status.start_time, %s), + end_time = CASE + WHEN EXCLUDED.status IN ('SUCCESS', 'FAILURE', 'REVOKED') AND task_status.end_time IS NULL + THEN %s + ELSE task_status.end_time + END + """, (task_id, parent_task_id, task_type, sub_type_identifier, status, progress, details_json, current_unix_time, status, current_unix_time, current_unix_time, current_unix_time)) + db.commit() + except psycopg2.Error: + logger.exception(f"DB Error saving task status for {task_id}") + try: + db.rollback() + logger.info(f"DB transaction rolled back for task status update of {task_id}.") + except psycopg2.Error: + logger.exception(f"DB Error during rollback for task status {task_id}") + finally: + cur.close() + + # Record persistent history for MAIN tasks that just reached a terminal state. + # Skip the synthetic 'unknown' placeholder inserted by the global cancel + # path (app_helper.cancel_all_jobs) — it has no real type and would show + # up as an 'unknown' row in the dashboard's recent activity table. + try: + if ( + parent_task_id is None + and status in (TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED) + and task_type and task_type != 'unknown' + ): + duration_s = None + try: + hist_cur = db.cursor() + hist_cur.execute( + "SELECT start_time, end_time FROM task_status WHERE task_id = %s", + (task_id,), + ) + row = hist_cur.fetchone() + hist_cur.close() + if row and row[0] is not None: + end = row[1] if row[1] is not None else current_unix_time + duration_s = max(0.0, float(end) - float(row[0])) + except Exception: + pass + record_task_history(task_id, task_type, status, duration_s, details=details) + except Exception as e_hist: + logger.debug(f"history record skipped for {task_id}: {e_hist}") + + +def get_task_info_from_db(task_id): + """Fetches task info from DB and calculates running time in Python.""" + db = get_db() + cur = db.cursor(cursor_factory=DictCursor) + # Fetch raw columns including the Unix timestamps + cur.execute(""" + SELECT + task_id, parent_task_id, task_type, sub_type_identifier, status, progress, details, timestamp, start_time, end_time + FROM task_status + WHERE task_id = %s + """, (task_id,)) + row = cur.fetchone() + cur.close() + if not row: + return None + + row_dict = dict(row) + current_unix_time = time.time() + + start_time = row_dict.get('start_time') + end_time = row_dict.get('end_time') + + # If start_time is null (old record or pre-start), duration is 0. + if start_time is None: + row_dict['running_time_seconds'] = 0.0 + else: + # If end_time is null, task is running. Use current time. + effective_end_time = end_time if end_time is not None else current_unix_time + row_dict['running_time_seconds'] = max(0, effective_end_time - start_time) + + return row_dict + + +# --------------------------------------------------------------------------- +# Score and map projection utilities +# --------------------------------------------------------------------------- + +def get_score_data_by_ids(item_ids_list): + """Fetches only score-related data (excluding embeddings) for a specific list of item_ids.""" + if not item_ids_list: + return [] + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + query = """ + SELECT s.item_id, s.title, s.author, s.album, s.album_artist, s.tempo, s.key, s.scale, s.mood_vector, s.energy, s.other_features, s.year, s.rating, s.file_path + FROM score s + WHERE s.item_id IN %s + """ + try: + cur.execute(query, (tuple(item_ids_list),)) + rows = cur.fetchall() + except Exception: + logger.exception("Error fetching score data by IDs") + rows = [] + finally: + cur.close() + return [dict(row) for row in rows] + + +def get_tracks_by_ids(item_ids_list): + """Fetches full track data (including embeddings) for a specific list of item_ids.""" + if not item_ids_list: + return [] + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + + # Convert item_ids to strings to match the text type in database + item_ids_str = [str(item_id) for item_id in item_ids_list] + + query = """ + SELECT s.item_id, s.title, s.author, s.album, s.album_artist, s.tempo, s.key, s.scale, s.mood_vector, s.energy, s.other_features, s.year, s.rating, s.file_path, e.embedding + FROM score s + LEFT JOIN embedding e ON s.item_id = e.item_id + WHERE s.item_id IN %s + """ + cur.execute(query, (tuple(item_ids_str),)) + rows = cur.fetchall() + cur.close() + + # Convert DictRow objects to regular dicts to allow adding new keys. + processed_rows = [] + for row in rows: + row_dict = dict(row) + if row_dict.get('embedding'): + row_dict['embedding_vector'] = np.frombuffer(row_dict['embedding'], dtype=np.float32) + else: + row_dict['embedding_vector'] = np.array([]) + processed_rows.append(row_dict) + + return processed_rows + + +def load_map_projection(index_name, force_reload=False): + """Load precomputed projection from DB. Returns (id_map, numpy_array) or (None, None)""" + global MAP_PROJECTION_CACHE + # Try cache first (unless force_reload is True) + if not force_reload and MAP_PROJECTION_CACHE and MAP_PROJECTION_CACHE.get('index_name') == index_name: + logger.info(f"Map projection '{index_name}' already loaded in cache. Skipping reload.") + return MAP_PROJECTION_CACHE.get('id_map'), MAP_PROJECTION_CACHE.get('projection') + + logger.info(f"Attempting to load map projection '{index_name}' from database into memory...") + conn = get_db() + cur = conn.cursor() + try: + cur.execute("SELECT projection_data, id_map_json FROM map_projection_data WHERE index_name = %s", (index_name,)) + row = cur.fetchone() + if row and row[0] is not None: + proj_blob, id_map_json = row[0], row[1] + else: + import re + from tasks.index_build_helpers import reassemble_segmented_id_map + cur.execute( + "SELECT index_name, projection_data, id_map_json FROM map_projection_data WHERE index_name LIKE %s ESCAPE '\\'", + (index_name.replace('_', r'\_') + r"\_%\_%",), + ) + candidates = cur.fetchall() + if not candidates: + logger.warning(f"Map projection '{index_name}' not found in the database. Cache will be empty.") + return None, None + seg_pattern = re.compile(rf"^{re.escape(index_name)}_(\d+)_(\d+)$") + parts = [] + total_expected = None + for name, part_blob, part_id_map in candidates: + m = seg_pattern.match(name) + if not m: + continue + part_no = int(m.group(1)) + total = int(m.group(2)) + if total_expected is None: + total_expected = total + elif total_expected != total: + logger.error(f"Map projection segment total mismatch for '{index_name}' ({total_expected} vs {total}). Aborting load.") + return None, None + parts.append((part_no, part_blob, part_id_map)) + if total_expected is None or len(parts) != total_expected: + logger.error(f"Incomplete map projection segments for '{index_name}': expected {total_expected}, found {len(parts)}. Aborting load.") + return None, None + parts.sort(key=lambda p: p[0]) + proj_blob = b"".join(bytes(p[1]) for p in parts if p[1]) + id_map_json = reassemble_segmented_id_map((p[0], p[2]) for p in parts) + proj = np.frombuffer(proj_blob, dtype=np.float32) + # infer shape as (-1,2) if length divisible by 2 + if proj.size % 2 == 0: + proj = proj.reshape((-1, 2)) + id_map = json.loads(id_map_json) + MAP_PROJECTION_CACHE = {'index_name': index_name, 'id_map': id_map, 'projection': proj} + logger.info(f"Map projection '{index_name}' with {len(id_map)} items loaded successfully into memory.") + return id_map, proj + except Exception: + logger.exception("Failed to load map projection") + return None, None + finally: + cur.close() + + +# --------------------------------------------------------------------------- +# Analysis and embedding utilities +# --------------------------------------------------------------------------- + +def save_track_analysis_and_embedding(item_id, title, author, tempo, key, scale, moods, embedding_vector, energy=None, other_features=None, album=None, album_artist=None, year=None, rating=None, file_path=None): + """Saves track analysis and embedding in a single transaction.""" + + # Sanitize all string inputs with field-specific limits + title = sanitize_db_field(title, max_length=500, field_name="title") + author = sanitize_db_field(author, max_length=200, field_name="author") + album = sanitize_db_field(album, max_length=200, field_name="album") + album_artist = sanitize_db_field(album_artist, max_length=200, field_name="album_artist") + key = sanitize_db_field(key, max_length=10, field_name="key") + scale = sanitize_db_field(scale, max_length=10, field_name="scale") + other_features = sanitize_db_field(other_features, max_length=2000, field_name="other_features") + + # year: parse from various date formats and validate + def _parse_year_from_date(year_value): + """ + Parse year from various date formats. + Supports: YYYY, YYYY-MM-DD, MM-DD-YYYY, DD-MM-YYYY (with - or / separators) + """ + if year_value is None: + return None + + year_str = str(year_value).strip() + if not year_str: + return None + + # Try parsing as pure integer first (YYYY) + try: + year = int(year_str) + if 1000 <= year <= 2100: + return year + except (ValueError, TypeError): + pass + + # Normalize separators + normalized = year_str.replace('/', '-') + parts = normalized.split('-') + + if len(parts) == 3: + try: + # YYYY-MM-DD format + if len(parts[0]) == 4: + year = int(parts[0]) + if 1000 <= year <= 2100: + return year + + # MM-DD-YYYY or DD-MM-YYYY format + if len(parts[2]) == 4: + year = int(parts[2]) + if 1000 <= year <= 2100: + return year + + # 2-digit year (MM-DD-YY) + if len(parts[2]) == 2: + year = int(parts[2]) + year += 2000 if year < 30 else 1900 + if 1000 <= year <= 2100: + return year + except (ValueError, TypeError, IndexError): + pass + + return None + + year = _parse_year_from_date(year) + + # rating: validate as integer 0-5 (5-star rating system) + if rating is not None: + try: + rating = int(rating) + if rating < 0 or rating > 5: + rating = None + except (ValueError, TypeError): + rating = None + + file_path = sanitize_db_field(file_path, max_length=1000, field_name="file_path") + + mood_str = ','.join(f"{k}:{v:.3f}" for k, v in moods.items()) + + conn = get_db() + cur = conn.cursor() + try: + # Save analysis to score table + cur.execute(""" + INSERT INTO score (item_id, title, author, tempo, key, scale, mood_vector, energy, other_features, album, album_artist, year, rating, file_path) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ON CONFLICT (item_id) DO UPDATE SET + title = EXCLUDED.title, + author = EXCLUDED.author, + tempo = EXCLUDED.tempo, + key = EXCLUDED.key, + scale = EXCLUDED.scale, + mood_vector = EXCLUDED.mood_vector, + energy = EXCLUDED.energy, + other_features = EXCLUDED.other_features, + album = EXCLUDED.album, + album_artist = EXCLUDED.album_artist, + year = EXCLUDED.year, + rating = EXCLUDED.rating, + file_path = EXCLUDED.file_path + """, (item_id, title, author, tempo, key, scale, mood_str, energy, other_features, album, album_artist, year, rating, file_path)) + + # Save embedding + if isinstance(embedding_vector, np.ndarray) and embedding_vector.size > 0: + embedding_blob = embedding_vector.astype(np.float32).tobytes() + cur.execute(""" + INSERT INTO embedding (item_id, embedding) VALUES (%s, %s) + ON CONFLICT (item_id) DO UPDATE SET embedding = EXCLUDED.embedding + """, (item_id, psycopg2.Binary(embedding_blob))) + + conn.commit() + except Exception: + conn.rollback() + logger.exception("Error saving track analysis and embedding for %s", item_id) + raise + finally: + cur.close() + + +def save_clap_embedding(item_id, clap_embedding_vector): + """Saves CLAP embedding for a track.""" + if clap_embedding_vector is None or (isinstance(clap_embedding_vector, np.ndarray) and clap_embedding_vector.size == 0): + return + + conn = get_db() + cur = conn.cursor() + try: + embedding_blob = clap_embedding_vector.astype(np.float32).tobytes() + cur.execute(""" + INSERT INTO clap_embedding (item_id, embedding) VALUES (%s, %s) + ON CONFLICT (item_id) DO UPDATE SET embedding = EXCLUDED.embedding + """, (item_id, psycopg2.Binary(embedding_blob))) + conn.commit() + except Exception: + conn.rollback() + logger.exception(f"Error saving CLAP embedding for {item_id}") + raise + finally: + cur.close() + + +def get_clap_embedding(item_id): + """Load CLAP embedding for a track from the database. + + Returns: + numpy array (512-dim float32) or None if not found + """ + conn = get_db() + cur = conn.cursor() + try: + cur.execute("SELECT embedding FROM clap_embedding WHERE item_id = %s", (item_id,)) + row = cur.fetchone() + if row and row[0]: + return np.frombuffer(row[0], dtype=np.float32) + return None + except Exception: + logger.exception(f"Error loading CLAP embedding for {item_id}") + return None + finally: + cur.close() + + +def save_lyrics_embedding(item_id, lyrics_embedding_vector, axis_vector=None): + """Saves the lyrics embedding (gte-multilingual-base) and the fixed-order axis vector. + + ``axis_vector`` must be a numpy array (float32) already in canonical + MUSIC_ANALYSIS_AXES order (use ``_score_axes`` to produce it). May be None. + """ + if lyrics_embedding_vector is None or (isinstance(lyrics_embedding_vector, np.ndarray) and lyrics_embedding_vector.size == 0): + return + + conn = get_db() + cur = conn.cursor() + try: + embedding_blob = lyrics_embedding_vector.astype(np.float32).tobytes() if isinstance(lyrics_embedding_vector, np.ndarray) else np.asarray(lyrics_embedding_vector, dtype=np.float32).tobytes() + axis_blob = None + if axis_vector is not None: + arr = axis_vector if isinstance(axis_vector, np.ndarray) else np.asarray(axis_vector, dtype=np.float32) + if arr.size > 0: + axis_blob = arr.astype(np.float32, copy=False).tobytes() + cur.execute(""" + INSERT INTO lyrics_embedding (item_id, embedding, axis_vector) VALUES (%s, %s, %s) + ON CONFLICT (item_id) DO UPDATE SET embedding = EXCLUDED.embedding, axis_vector = EXCLUDED.axis_vector, updated_at = CURRENT_TIMESTAMP + """, (item_id, psycopg2.Binary(embedding_blob), + psycopg2.Binary(axis_blob) if axis_blob is not None else None)) + conn.commit() + except Exception: + conn.rollback() + logger.exception(f"Error saving lyrics embedding for {item_id}") + raise + finally: + cur.close() + + +# In-memory cache for the precomputed 2D artist component projections (optional). +ARTIST_PROJECTION_CACHE = None + + +# --------------------------------------------------------------------------- +# Database schema +# --------------------------------------------------------------------------- + +def init_db(): + db = get_db() + with db.cursor() as cur: + # Serialize concurrent init_db() runs across gunicorn workers/containers. + # Multiple workers racing on CREATE EXTENSION / CREATE OR REPLACE FUNCTION + # causes Postgres "tuple concurrently updated" errors on pg_proc/pg_extension. + # A session-level advisory lock forces other workers to wait here. + # The key is an arbitrary stable bigint specific to this app's init. + # Safety: session-level advisory locks are auto-released by Postgres + # when the connection ends (normal close, crash, kill, or network drop), + # so this lock can NEVER leak permanently even if init_db() raises. + cur.execute("SELECT pg_advisory_lock(726354821)") + try: + # Enable extensions to fix and assist in searches + if sys.platform == 'win32': + for ext in ('unaccent', 'pg_trgm'): + cur.execute("SAVEPOINT ext_create") + try: + cur.execute(f'CREATE EXTENSION IF NOT EXISTS {ext}') + cur.execute("RELEASE SAVEPOINT ext_create") + except Exception: + logger.warning("Extension %s not available -- skipping", ext) + cur.execute("ROLLBACK TO SAVEPOINT ext_create") + else: + cur.execute('CREATE EXTENSION IF NOT EXISTS unaccent') + cur.execute('CREATE EXTENSION IF NOT EXISTS pg_trgm') + # Create 'score' table + cur.execute("CREATE TABLE IF NOT EXISTS score (item_id TEXT PRIMARY KEY, title TEXT, author TEXT, album TEXT, album_artist TEXT, tempo REAL, key TEXT, scale TEXT, mood_vector TEXT)") + # Add 'energy' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'energy')") + if not cur.fetchone()[0]: + logger.info("Adding 'energy' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN energy REAL") + # Add 'other_features' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'other_features')") + if not cur.fetchone()[0]: + logger.info("Adding 'other_features' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN other_features TEXT") + # Add 'album' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'album')") + if not cur.fetchone()[0]: + logger.info("Adding 'album' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN album TEXT") + # Add 'album_artist' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'album_artist')") + if not cur.fetchone()[0]: + logger.info("Adding 'album_artist' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN album_artist TEXT") + # Add 'year' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'year')") + if not cur.fetchone()[0]: + logger.info("Adding 'year' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN year INTEGER") + # Add 'rating' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'rating')") + if not cur.fetchone()[0]: + logger.info("Adding 'rating' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN rating INTEGER") + # Add 'file_path' column if not exists + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'file_path')") + if not cur.fetchone()[0]: + logger.info("Adding 'file_path' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN file_path TEXT") + + # Ensure we have a searchable, accent-stripped `search_u` column. + # Postgres does not allow generated columns to call `unaccent()` (it's not marked immutable), + # so we store the value in a normal column and keep it in sync via trigger. + cur.execute("SELECT is_generated FROM information_schema.columns WHERE table_name = 'score' AND column_name = 'search_u'") + row = cur.fetchone() + search_u_generated = (row and row[0] == 'ALWAYS') + + if search_u_generated: + logger.info("Dropping legacy generated 'search_u' column to replace it with a trigger-updated column.") + cur.execute("ALTER TABLE score DROP COLUMN IF EXISTS search_u") + row = None + + # Create plain `search_u` column if missing + if not row: + logger.info("Adding 'search_u' column to 'score' table.") + cur.execute("ALTER TABLE score ADD COLUMN search_u TEXT") + + # Create helper function for accent stripping (safe to run multiple times) + if sys.platform == 'win32': + cur.execute("SAVEPOINT search_setup") + try: + cur.execute("CREATE OR REPLACE FUNCTION immutable_unaccent(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$ SELECT public.unaccent($1) $$;") + cur.execute(""" + CREATE OR REPLACE FUNCTION score_search_u_sync() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + NEW.search_u := lower(immutable_unaccent(concat_ws(' ', NEW.title, NEW.author, NEW.album))); + RETURN NEW; + END; + $$; + """) + cur.execute("DROP TRIGGER IF EXISTS score_search_u_sync_trigger ON score") + cur.execute(""" + CREATE TRIGGER score_search_u_sync_trigger + BEFORE INSERT OR UPDATE ON score + FOR EACH ROW + EXECUTE FUNCTION score_search_u_sync(); + """) + cur.execute("UPDATE score SET search_u = lower(immutable_unaccent(concat_ws(' ', title, author, album))) WHERE search_u IS NULL") + cur.execute("CREATE INDEX IF NOT EXISTS score_search_u_trgm ON score USING gin (search_u gin_trgm_ops)") + cur.execute("RELEASE SAVEPOINT search_setup") + except Exception: + logger.warning("unaccent/pg_trgm extensions not available -- accent-insensitive search disabled") + cur.execute("ROLLBACK TO SAVEPOINT search_setup") + else: + cur.execute("CREATE OR REPLACE FUNCTION immutable_unaccent(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$ SELECT public.unaccent($1) $$;") + cur.execute(""" + CREATE OR REPLACE FUNCTION score_search_u_sync() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + NEW.search_u := lower(immutable_unaccent(concat_ws(' ', NEW.title, NEW.author, NEW.album))); + RETURN NEW; + END; + $$; + """) + cur.execute("DROP TRIGGER IF EXISTS score_search_u_sync_trigger ON score") + cur.execute(""" + CREATE TRIGGER score_search_u_sync_trigger + BEFORE INSERT OR UPDATE ON score + FOR EACH ROW + EXECUTE FUNCTION score_search_u_sync(); + """) + cur.execute("UPDATE score SET search_u = lower(immutable_unaccent(concat_ws(' ', title, author, album))) WHERE search_u IS NULL") + cur.execute("CREATE INDEX IF NOT EXISTS score_search_u_trgm ON score USING gin (search_u gin_trgm_ops)") + + # Create 'playlist' table + cur.execute("CREATE TABLE IF NOT EXISTS playlist (id SERIAL PRIMARY KEY, playlist_name TEXT, item_id TEXT, title TEXT, author TEXT, UNIQUE (playlist_name, item_id))") + # Create 'task_status' table + cur.execute("CREATE TABLE IF NOT EXISTS task_status (id SERIAL PRIMARY KEY, task_id TEXT UNIQUE NOT NULL, parent_task_id TEXT, task_type TEXT NOT NULL, sub_type_identifier TEXT, status TEXT, progress INTEGER DEFAULT 0, details TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Migrate 'start_time' and 'end_time' columns + for col_name in ['start_time', 'end_time']: + cur.execute("SELECT data_type FROM information_schema.columns WHERE table_name = 'task_status' AND column_name = %s", (col_name,)) + if not cur.fetchone(): cur.execute(f"ALTER TABLE task_status ADD COLUMN {col_name} DOUBLE PRECISION") + # Create 'task_history' table — a small, persistent log of the last + # completed/cancelled MAIN tasks. Survives the global Cancel button + # which wipes `task_status`. Capped to the most recent 10 rows. + cur.execute(""" + CREATE TABLE IF NOT EXISTS task_history ( + id SERIAL PRIMARY KEY, + recorded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + task_id TEXT, + task_type TEXT, + status TEXT, + duration_seconds DOUBLE PRECISION, + note TEXT + ) + """) + # Create 'embedding' table + cur.execute("CREATE TABLE IF NOT EXISTS embedding (item_id TEXT PRIMARY KEY, FOREIGN KEY (item_id) REFERENCES score (item_id) ON DELETE CASCADE)") + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'embedding' AND column_name = 'embedding')") + if not cur.fetchone()[0]: cur.execute("ALTER TABLE embedding ADD COLUMN embedding BYTEA") + # Create 'lyrics_embedding' table for lyrics similarity and axis scores + cur.execute("CREATE TABLE IF NOT EXISTS lyrics_embedding (item_id TEXT PRIMARY KEY, FOREIGN KEY (item_id) REFERENCES score (item_id) ON DELETE CASCADE)") + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'lyrics_embedding' AND column_name = 'embedding')") + if not cur.fetchone()[0]: cur.execute("ALTER TABLE lyrics_embedding ADD COLUMN embedding BYTEA") + # axis_vector: float32 BYTEA, fixed-order flattened over MUSIC_ANALYSIS_AXES. + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'lyrics_embedding' AND column_name = 'axis_vector')") + if not cur.fetchone()[0]: cur.execute("ALTER TABLE lyrics_embedding ADD COLUMN axis_vector BYTEA") + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'lyrics_embedding' AND column_name = 'updated_at')") + if not cur.fetchone()[0]: cur.execute("ALTER TABLE lyrics_embedding ADD COLUMN updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") + # Create 'clap_embedding' table for CLAP text search embeddings + cur.execute("CREATE TABLE IF NOT EXISTS clap_embedding (item_id TEXT PRIMARY KEY, FOREIGN KEY (item_id) REFERENCES score (item_id) ON DELETE CASCADE)") + cur.execute("SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'clap_embedding' AND column_name = 'embedding')") + if not cur.fetchone()[0]: cur.execute("ALTER TABLE clap_embedding ADD COLUMN embedding BYTEA") + # Create 'voyager_index_data' table + cur.execute("CREATE TABLE IF NOT EXISTS voyager_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'clap_index_data' table for stored CLAP text search indexes + cur.execute("CREATE TABLE IF NOT EXISTS clap_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'lyrics_index_data' table for stored Lyrics voyager indexes (mirrors clap_index_data; supports chunked storage). + cur.execute("CREATE TABLE IF NOT EXISTS lyrics_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'lyrics_axes_index_data' table for the axis-vector voyager index (one binary-friendly vector per song over MUSIC_ANALYSIS_AXES labels). + cur.execute("CREATE TABLE IF NOT EXISTS lyrics_axes_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'artist_index_data' table for artist GMM-based HNSW index + cur.execute("CREATE TABLE IF NOT EXISTS artist_index_data (index_name VARCHAR(255) PRIMARY KEY, index_data BYTEA NOT NULL, artist_map_json TEXT NOT NULL, gmm_params_json TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'artist_metadata_data' table for the per-artist auxiliary + # metadata blob (artist_map + GMM params). Decoupled from the Voyager + # index binary and segmented independently so a single column value + # never crosses PG's 1 GB MaxAllocSize cap, regardless of library size. + cur.execute("CREATE TABLE IF NOT EXISTS artist_metadata_data (name VARCHAR(255) PRIMARY KEY, blob_data BYTEA NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'map_projection_data' table for precomputed 2D map projections + cur.execute("CREATE TABLE IF NOT EXISTS map_projection_data (index_name VARCHAR(255) PRIMARY KEY, projection_data BYTEA NOT NULL, id_map_json TEXT NOT NULL, embedding_dimension INTEGER NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'artist_component_projection' table for precomputed 2D artist component projections + cur.execute("CREATE TABLE IF NOT EXISTS artist_component_projection (index_name VARCHAR(255) PRIMARY KEY, projection_data BYTEA NOT NULL, artist_component_map_json TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'cron' table to hold scheduled jobs (very small and simple) + cur.execute("CREATE TABLE IF NOT EXISTS cron (id SERIAL PRIMARY KEY, name TEXT, task_type TEXT NOT NULL, cron_expr TEXT NOT NULL, enabled BOOLEAN DEFAULT FALSE, last_run DOUBLE PRECISION, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Create 'audiomuse_users' table. Every account (including the + # install-time admin) lives here. 'role' is 'admin' or 'user'. + cur.execute("CREATE TABLE IF NOT EXISTS audiomuse_users (id SERIAL PRIMARY KEY, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, role TEXT NOT NULL DEFAULT 'user', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Lightweight migration for installs that already have the table without a role column. + cur.execute("ALTER TABLE audiomuse_users ADD COLUMN IF NOT EXISTS role TEXT NOT NULL DEFAULT 'user'") + # Create 'dashboard_stats' singleton table (id fixed to 1) that holds + # precomputed content/library aggregates and index counts. Refreshed + # at app startup and hourly by a background job so the dashboard + # does not have to scan the whole `score` table on every poll. + cur.execute( + "CREATE TABLE IF NOT EXISTS dashboard_stats (" + "id INTEGER PRIMARY KEY, " + "updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " + "content JSONB NOT NULL DEFAULT '{}'::jsonb, " + "indexes JSONB NOT NULL DEFAULT '[]'::jsonb, " + "CONSTRAINT dashboard_stats_singleton CHECK (id = 1))" + ) + # Ensure older restored DBs still have the primary key constraint. + cur.execute( + "SELECT COUNT(*) FROM information_schema.table_constraints " + "WHERE table_name = 'dashboard_stats' AND constraint_type = 'PRIMARY KEY'" + ) + row = cur.fetchone() + if row and row[0] == 0: + logger.info("Cleaning dashboard_stats and adding missing primary key constraint to dashboard_stats.id") + cur.execute("DELETE FROM dashboard_stats") + cur.execute("ALTER TABLE dashboard_stats ADD CONSTRAINT dashboard_stats_pkey PRIMARY KEY (id)") + # Create 'artist_mapping' table to map artist names to media server artist IDs + cur.execute("CREATE TABLE IF NOT EXISTS artist_mapping (artist_name TEXT PRIMARY KEY, artist_id TEXT)") + # Create application configuration table to persist setup values. + cur.execute( + "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'app_config')" + ) + if not cur.fetchone()[0]: + cur.execute( + "CREATE TABLE app_config (" + "key TEXT PRIMARY KEY, value TEXT NOT NULL, " + "updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)" + ) + # Create 'alchemy_anchors' table to persist named user anchors for reuse + cur.execute("CREATE TABLE IF NOT EXISTS alchemy_anchors (id SERIAL PRIMARY KEY, name TEXT UNIQUE NOT NULL, centroid JSONB NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + cur.execute("CREATE TABLE IF NOT EXISTS alchemy_radios (id SERIAL PRIMARY KEY, anchor_id INTEGER UNIQUE NOT NULL REFERENCES alchemy_anchors(id) ON DELETE CASCADE, temperature DOUBLE PRECISION NOT NULL, n_results INTEGER NOT NULL, enabled BOOLEAN NOT NULL DEFAULT TRUE, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") + # Provider migration tool: wizard session state (one row per migration attempt) + cur.execute(""" + CREATE TABLE IF NOT EXISTS migration_session ( + id SERIAL PRIMARY KEY, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + completed_at TIMESTAMP, + status TEXT NOT NULL DEFAULT 'in_progress', + source_type TEXT NOT NULL, + target_type TEXT NOT NULL, + target_creds TEXT NOT NULL, + state JSONB NOT NULL DEFAULT '{}' + ) + """) + # Create 'text_search_queries' table for precomputed CLAP text search queries + cur.execute(""" + CREATE TABLE IF NOT EXISTS text_search_queries ( + id SERIAL PRIMARY KEY, + query_text TEXT NOT NULL, + score REAL NOT NULL, + rank INTEGER NOT NULL, + created_at TIMESTAMP DEFAULT NOW(), + UNIQUE(rank) + ) + """) + cur.execute("CREATE INDEX IF NOT EXISTS idx_text_search_queries_rank ON text_search_queries(rank)") + + # Insert default queries if table is empty + cur.execute("SELECT COUNT(*) FROM text_search_queries") + count = cur.fetchone()[0] + + if count == 0: + default_queries = [ + "female vocal romantic trap", + "synth indie pop raspy", + "sad hard rock male vocal", + "funk falsetto energetic", + "groovy sax blues", + "classical relaxed piano", + "belting jazz happy", + "tabla afrobeat fast-paced", + "harmonized vocals slow-paced electronica", + "autotuned gospel excited", + "breathy aggressive house", + "smooth folk mid-tempo", + "deep voice r&b dark", + "punk guitar angry", + "metal choir dreamy", + "chant reggae trumpet", + "high-pitched brass hip-hop", + "disco whispered drum machine", + "happy whispered indie pop", + "synth energetic raspy", + "rock slow-paced cello", + "falsetto jazz excited", + "r&b male vocal romantic", + "harmonized vocals dark trap", + "smooth blues sax", + "high-pitched fast-paced soul", + "female vocal sad hip-hop", + "congas aggressive soul", + "mid-tempo afrobeat autotuned", + "belting funk groovy", + "angry alternative breathy", + "gospel choir steelpan", + "viola relaxed folk", + "dreamy rhodes metal", + "acoustic guitar country chant", + "deep voice orchestra reggae", + "fast-paced synth progressive rock", + "hard rock raspy romantic", + "fast-paced electric guitar progressive rock", + "hard rock aggressive breathy", + "rock high-pitched energetic", + "autotuned energetic hip-hop", + "raspy fast-paced blues", + "belting electronica energetic", + "whispered indie pop aggressive", + "harmonized vocals aggressive synth", + "orchestra whispered romantic", + "belting mid-tempo progressive rock", + "autotuned pop mid-tempo", + "pop energetic synthesizer" + ] + + for rank, query in enumerate(default_queries, start=1): + cur.execute(""" + INSERT INTO text_search_queries (query_text, score, rank, created_at) + VALUES (%s, %s, %s, NOW()) + """, (query, 1.0, rank)) + + logger.info(f"Inserted {len(default_queries)} default DCLAP search queries") + + db.commit() + # Release the advisory lock acquired at the top of init_db(). + finally: + cur.execute("SELECT pg_advisory_unlock(726354821)") + + + +# --------------------------------------------------------------------------- +# Task lifecycle: status archival, active-task lookup, child tasks +# --------------------------------------------------------------------------- + +def clean_up_previous_main_tasks(): + """ + Cleans up all previous main tasks before a new one starts. + - Archives tasks in SUCCESS state. + - Archives stale tasks stuck in PENDING, STARTED, or PROGRESS states. + - DELETES all child tasks associated with archived parent tasks to prevent DB bloat. + A main task is identified by having a NULL parent_task_id. + """ + db = get_db() # This now calls the function within this file + cur = db.cursor(cursor_factory=DictCursor) + logger.info("Starting cleanup of all previous main tasks.") + + non_terminal_statuses = (TASK_STATUS_PENDING, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS) + + try: + cur.execute("SELECT task_id, status, details, task_type, start_time, end_time FROM task_status WHERE status IN %s AND parent_task_id IS NULL", (non_terminal_statuses,)) + tasks_to_archive = cur.fetchall() + + archived_count = 0 + deleted_children_count = 0 + + for task_row in tasks_to_archive: + task_id = task_row['task_id'] + original_status = task_row['status'] + + original_details_json = task_row['details'] + original_status_message = f"Task was in '{original_status}' state." + + original_details_dict = None + if original_details_json: + try: + original_details_dict = json.loads(original_details_json) + original_status_message = original_details_dict.get("status_message", original_status_message) + except (json.JSONDecodeError, TypeError): + logger.warning(f"Could not parse original details for task {task_id} during archival.") + + # Record into persistent history BEFORE deleting children — the + # note builder needs to query subtasks (e.g. tracks_analyzed). + try: + duration_s = None + if task_row['start_time'] is not None: + end = task_row['end_time'] if task_row['end_time'] is not None else time.time() + duration_s = max(0.0, float(end) - float(task_row['start_time'])) + final_status = TASK_STATUS_SUCCESS if original_status == TASK_STATUS_SUCCESS else TASK_STATUS_REVOKED + record_task_history( + task_id, task_row['task_type'], final_status, + duration_s, details=original_details_dict, + ) + except Exception as e_hist: + logger.debug(f"history record skipped during archive of {task_id}: {e_hist}") + + if original_status == TASK_STATUS_SUCCESS: + archival_reason = "New main task started, old successful task archived." + else: + archival_reason = f"New main task started, stale task (status: {original_status}) has been archived." + + archived_details = { + "log": [f"[Archived] {archival_reason}. Original summary: {original_status_message}"], + "original_status_before_archival": original_status, + "archival_reason": archival_reason + } + archived_details_json = json.dumps(archived_details) + + with db.cursor() as update_cur: + # First, delete all child tasks to prevent DB bloat and avoid counting old tasks + update_cur.execute( + "DELETE FROM task_status WHERE parent_task_id = %s", + (task_id,) + ) + children_deleted = update_cur.rowcount + deleted_children_count += children_deleted + + if children_deleted > 0: + logger.info(f"Deleted {children_deleted} child tasks for parent task {task_id}") + + # Then archive the parent task + update_cur.execute( + "UPDATE task_status SET status = %s, details = %s, progress = 100, timestamp = NOW() WHERE task_id = %s AND status = %s", + (TASK_STATUS_REVOKED, archived_details_json, task_id, original_status) + ) + archived_count += 1 + + if archived_count > 0: + db.commit() + logger.info(f"Archived {archived_count} previous main tasks and deleted {deleted_children_count} child tasks.") + else: + logger.info("No previous main tasks found to clean up.") + except Exception: + db.rollback() + logger.exception("Error during the main task cleanup process") + finally: + cur.close() + +def get_active_main_task(task_type=None): + """Return the currently active main task. + + If task_type is provided, only return an active task of that type. + If task_type is None, return any active main task. + """ + db = get_db() + cur = db.cursor(cursor_factory=DictCursor) + non_terminal_statuses = (TASK_STATUS_PENDING, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS) + + if task_type: + cur.execute(""" + SELECT task_id, task_type, status, details + FROM task_status + WHERE task_type = %s AND status IN %s AND parent_task_id IS NULL + ORDER BY timestamp DESC + LIMIT 1 + """, (task_type, non_terminal_statuses)) + else: + cur.execute(""" + SELECT task_id, task_type, status, details + FROM task_status + WHERE status IN %s AND parent_task_id IS NULL + ORDER BY timestamp DESC + LIMIT 1 + """, (non_terminal_statuses,)) + + active_task = cur.fetchone() + cur.close() + return dict(active_task) if active_task else None + +def get_child_tasks_from_db(parent_task_id): + """Fetches all child tasks for a given parent_task_id from the database.""" + conn = get_db() # This now calls the function within this file + cur = conn.cursor(cursor_factory=DictCursor) + # MODIFIED: Select the 'details' column as well for the final check. + cur.execute("SELECT task_id, status, sub_type_identifier, details FROM task_status WHERE parent_task_id = %s", (parent_task_id,)) + tasks = cur.fetchall() + cur.close() + # DictCursor returns a list of dictionary-like objects, convert to plain dicts + return [dict(row) for row in tasks] + + + +# --------------------------------------------------------------------------- +# Song Alchemy -- anchors and radios (CRUD) +# --------------------------------------------------------------------------- + +def save_alchemy_anchor(name, centroid): + """Save a named anchor centroid into DB.""" + if not name or not centroid or not isinstance(centroid, list): + raise ValueError('Anchor name and centroid list are required.') + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + centroid_json = json.dumps(centroid) + cur.execute( + "INSERT INTO alchemy_anchors (name, centroid) VALUES (%s, %s) " + "ON CONFLICT (name) DO UPDATE SET centroid = EXCLUDED.centroid, created_at = NOW() " + "RETURNING id, name, created_at", + (name, centroid_json) + ) + row = cur.fetchone() + conn.commit() + return dict(row) if row else None + except Exception: + conn.rollback() + logger.exception(f"Failed to save alchemy anchor '{name}'") + return None + finally: + cur.close() + +def get_alchemy_anchors(): + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + cur.execute("SELECT id, name, created_at FROM alchemy_anchors ORDER BY created_at DESC") + rows = cur.fetchall() + return [dict(row) for row in rows] + except Exception: + logger.exception("Failed to load alchemy anchors") + return [] + finally: + cur.close() + +def delete_alchemy_anchor(anchor_id): + conn = get_db() + cur = conn.cursor() + try: + cur.execute("DELETE FROM alchemy_anchors WHERE id = %s", (anchor_id,)) + conn.commit() + return cur.rowcount > 0 + except Exception: + conn.rollback() + logger.exception(f"Failed to delete alchemy anchor id={anchor_id}") + return False + finally: + cur.close() + +def get_alchemy_anchor_by_id(anchor_id): + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + cur.execute("SELECT id, name, centroid, created_at FROM alchemy_anchors WHERE id = %s", (anchor_id,)) + row = cur.fetchone() + if not row: + return None + anchor = dict(row) + if isinstance(anchor.get('centroid'), str): + try: + anchor['centroid'] = json.loads(anchor['centroid']) + except Exception: + anchor['centroid'] = None + return anchor + except Exception: + logger.exception(f"Failed to fetch alchemy anchor id={anchor_id}") + return None + finally: + cur.close() + +def update_alchemy_anchor_name(anchor_id, name): + if not name or not isinstance(name, str): + return None + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + cur.execute( + "UPDATE alchemy_anchors SET name = %s WHERE id = %s RETURNING id, name", + (name.strip(), anchor_id) + ) + row = cur.fetchone() + conn.commit() + if not row: + return None + return dict(row) + except Exception: + conn.rollback() + logger.exception(f"Failed to rename alchemy anchor id={anchor_id}") + return None + finally: + cur.close() + +def get_alchemy_radios(): + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + cur.execute( + "SELECT r.id, r.anchor_id, a.name, r.temperature, r.n_results, r.enabled " + "FROM alchemy_radios r JOIN alchemy_anchors a ON a.id = r.anchor_id " + "ORDER BY a.name" + ) + rows = cur.fetchall() + return [dict(row) for row in rows] + except Exception: + logger.exception("Failed to load alchemy radios") + return [] + finally: + cur.close() + +def create_alchemy_radio(anchor_id, temperature, n_results, enabled=True): + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + cur.execute( + "INSERT INTO alchemy_radios (anchor_id, temperature, n_results, enabled) " + "VALUES (%s, %s, %s, %s) RETURNING id, anchor_id, temperature, n_results, enabled", + (anchor_id, temperature, n_results, bool(enabled)) + ) + row = cur.fetchone() + conn.commit() + return dict(row) if row else None + except Exception: + conn.rollback() + logger.exception(f"Failed to create alchemy radio for anchor_id={anchor_id}") + return None + finally: + cur.close() + +def update_alchemy_radio(radio_id, temperature, n_results, enabled): + conn = get_db() + cur = conn.cursor(cursor_factory=DictCursor) + try: + cur.execute( + "UPDATE alchemy_radios SET temperature = %s, n_results = %s, enabled = %s " + "WHERE id = %s RETURNING id, anchor_id, temperature, n_results, enabled", + (temperature, n_results, bool(enabled), radio_id) + ) + row = cur.fetchone() + conn.commit() + return dict(row) if row else None + except Exception: + conn.rollback() + logger.exception(f"Failed to update alchemy radio id={radio_id}") + return None + finally: + cur.close() + +def delete_alchemy_radio(radio_id): + conn = get_db() + cur = conn.cursor() + try: + cur.execute("DELETE FROM alchemy_radios WHERE id = %s", (radio_id,)) + conn.commit() + return cur.rowcount > 0 + except Exception: + conn.rollback() + logger.exception(f"Failed to delete alchemy radio id={radio_id}") + return False + finally: + cur.close() + + + +# --------------------------------------------------------------------------- +# Map and artist projection persistence +# --------------------------------------------------------------------------- + +def save_map_projection(index_name, id_map, projection_array): + """ + Save a precomputed 2D projection into the map_projection_data table. + projection_array: numpy array of shape (N,2), dtype=float32 + id_map: JSON-serializable list/dict mapping rows to item_ids + """ + conn = get_db() + try: + blob = projection_array.astype(np.float32).tobytes() + if not blob: + logger.info(f"Map projection '{index_name}' has no data; clearing existing store.") + with conn.cursor() as cur: + cur.execute( + "DELETE FROM map_projection_data WHERE index_name = %s OR index_name LIKE %s ESCAPE '\\'", + (index_name, index_name.replace('_', r'\_') + r"\_%\_%"), + ) + conn.commit() + return + embedding_dim = projection_array.shape[1] if projection_array.ndim == 2 else 0 + from tasks.index_build_helpers import store_voyager_index_segmented + store_voyager_index_segmented( + conn, + target_table="map_projection_data", + index_name=index_name, + index_bytes=blob, + id_map=id_map, + embedding_dimension=embedding_dim, + binary_column="projection_data", + ) + conn.commit() + try: + id_count = len(id_map) if hasattr(id_map, '__len__') else None + logger.info(f"Saved map projection '{index_name}' to DB: {len(blob)} bytes, ids={id_count}") + except Exception: + logger.debug("Saved map projection but failed to compute size/id_count for log.") + except Exception: + conn.rollback() + logger.exception("Failed to save map projection") + raise + +def load_artist_projection(index_name='artist_map', force_reload=False): + """Load precomputed artist component projection from DB. + Returns (artist_component_map, numpy_array) or (None, None). + artist_component_map format: [{'artist_id': '...', 'component_idx': 0, 'weight': 0.3}, ...] + """ + global ARTIST_PROJECTION_CACHE + # Try cache first (unless force_reload is True) + if not force_reload and ARTIST_PROJECTION_CACHE and ARTIST_PROJECTION_CACHE.get('index_name') == index_name: + logger.info(f"Artist projection '{index_name}' already loaded in cache. Skipping reload.") + return ARTIST_PROJECTION_CACHE.get('component_map'), ARTIST_PROJECTION_CACHE.get('projection') + + logger.info(f"Attempting to load artist projection '{index_name}' from database into memory...") + conn = get_db() + cur = conn.cursor() + try: + cur.execute("SELECT projection_data, artist_component_map_json FROM artist_component_projection WHERE index_name = %s", (index_name,)) + row = cur.fetchone() + if not row: + logger.warning(f"Artist projection '{index_name}' not found in the database. Cache will be empty.") + return None, None + proj_blob, component_map_json = row[0], row[1] + proj = np.frombuffer(proj_blob, dtype=np.float32) + # infer shape as (-1,2) if length divisible by 2 + if proj.size % 2 == 0: + proj = proj.reshape((-1, 2)) + component_map = json.loads(component_map_json) + ARTIST_PROJECTION_CACHE = {'index_name': index_name, 'component_map': component_map, 'projection': proj} + logger.info(f"Artist projection '{index_name}' with {len(component_map)} components loaded successfully into memory.") + return component_map, proj + except Exception: + logger.exception("Failed to load artist projection") + return None, None + finally: + cur.close() + +def save_artist_projection(index_name, component_map, projections): + """Save artist component projection to database. + component_map: [{'artist_id': '...', 'component_idx': 0, 'weight': 0.3}, ...] + projections: numpy array of shape (N, 2) + """ + conn = get_db() + cur = conn.cursor() + try: + component_map_json = json.dumps(component_map) + proj_blob = projections.astype(np.float32).tobytes() + cur.execute("INSERT INTO artist_component_projection (index_name, projection_data, artist_component_map_json) VALUES (%s, %s, %s) ON CONFLICT (index_name) DO UPDATE SET projection_data = EXCLUDED.projection_data, artist_component_map_json = EXCLUDED.artist_component_map_json, created_at = CURRENT_TIMESTAMP", (index_name, proj_blob, component_map_json)) + conn.commit() + logger.info(f"Saved artist projection '{index_name}' with {len(component_map)} components to database.") + except Exception: + conn.rollback() + logger.exception("Failed to save artist projection") + finally: + cur.close() + + + +# --------------------------------------------------------------------------- +# Playlists +# --------------------------------------------------------------------------- + +def update_playlist_table(playlists): # Removed db_path + conn = get_db() # This now calls the function within this file + cur = conn.cursor() + try: + # Clear all previous conceptual playlists to reflect only the current run. + cur.execute("DELETE FROM playlist") + for name, cluster in playlists.items(): + for item_id, title, author in cluster: + cur.execute("INSERT INTO playlist (playlist_name, item_id, title, author) VALUES (%s, %s, %s, %s) ON CONFLICT (playlist_name, item_id) DO NOTHING", (name, item_id, title, author)) + conn.commit() + except Exception: + conn.rollback() + logger.exception("Error updating playlist table") + finally: + cur.close() diff --git a/lyrics/lyrics_transcriber.py b/lyrics/lyrics_transcriber.py index 45faa563..3c55253a 100644 --- a/lyrics/lyrics_transcriber.py +++ b/lyrics/lyrics_transcriber.py @@ -379,7 +379,6 @@ def _fetch_from_configured_api( url, headers={'Accept': 'application/json'}, ) - import socket ctx = None try: import ssl diff --git a/native-build/linux/embedded_pg.py b/native-build/linux/embedded_pg.py index 05819da6..06f54cd7 100644 --- a/native-build/linux/embedded_pg.py +++ b/native-build/linux/embedded_pg.py @@ -79,15 +79,33 @@ def _initialized(data_dir): return False +def _has_cluster_data(data_dir): + """True if data_dir holds an initialized PostgreSQL cluster (never auto-delete it). + + Keyed on ``global/pg_control``: written at the END of initdb and required for + the server to start, so its presence proves a complete cluster with real data. + A half-built dir (interrupted initdb, no pg_control) holds no usable data and + is still cleared normally, so this guard never blocks first-run self-heal. + """ + return os.path.exists(os.path.join(data_dir, "global", "pg_control")) + + def _reset_data_dir(data_dir): """Empty a partially-initialized data dir so initdb can retry. initdb refuses a non-empty target, so a half-built cluster (interrupted initdb, no completion marker) would otherwise wedge every start forever. - Only ever called when :func:`_initialized` is False, so a complete cluster's - data is never touched.""" + Only ever called when :func:`_initialized` is False; additionally refuses to + delete a dir that still holds a real cluster, so a transient mis-detection + can never destroy existing data (it surfaces an error instead).""" if not (os.path.isdir(data_dir) and os.listdir(data_dir)): return + if _has_cluster_data(data_dir): + raise RuntimeError( + f"Refusing to wipe {data_dir}: it contains an existing PostgreSQL " + "cluster (global/pg_control present). Back it up or remove it " + "manually if you really want a fresh start." + ) logger.warning("Clearing incomplete PostgreSQL data dir %s before re-init", data_dir) for entry in os.listdir(data_dir): target = os.path.join(data_dir, entry) diff --git a/native-build/linux/launcher.py b/native-build/linux/launcher.py index 94de9175..e5b90eb2 100644 --- a/native-build/linux/launcher.py +++ b/native-build/linux/launcher.py @@ -49,7 +49,7 @@ def _run_flask(): import app as app_module waitress.serve( app_module.app, - host="127.0.0.1", + host="0.0.0.0", port=8000, threads=8, max_request_body_size=6 * 1024 * 1024 * 1024, diff --git a/native-build/macos/launcher.py b/native-build/macos/launcher.py index 195e018e..b52fcdb6 100644 --- a/native-build/macos/launcher.py +++ b/native-build/macos/launcher.py @@ -26,7 +26,7 @@ def _run_flask(): import app as app_module waitress.serve( app_module.app, - host="127.0.0.1", + host="0.0.0.0", port=8000, threads=8, max_request_body_size=6 * 1024 * 1024 * 1024, diff --git a/native-build/windows/db_backend.py b/native-build/windows/db_backend.py index da45e45a..9d81c9ed 100644 --- a/native-build/windows/db_backend.py +++ b/native-build/windows/db_backend.py @@ -22,6 +22,7 @@ is never reachable without the password and there is no trust window. """ +import importlib import logging import os import subprocess @@ -45,7 +46,7 @@ def _check_pgserver(): global _USE_PGSERVER if _USE_PGSERVER is None: try: - import pgserver + importlib.import_module("pgserver") _USE_PGSERVER = True logger.info("Using pgserver for embedded PostgreSQL") except ImportError: @@ -150,16 +151,34 @@ def _is_trust(line): logger.exception("Could not upgrade legacy PostgreSQL auth; leaving as-is") +def _has_cluster_data(data_dir): + """True if data_dir holds an initialized PostgreSQL cluster (never auto-delete it). + + Keyed on ``global/pg_control``: written at the END of initdb and required for + the server to start, so its presence proves a complete cluster with real data. + A half-built dir (interrupted initdb, no pg_control) holds no usable data and + is still cleared normally, so this guard never blocks first-run self-heal. + """ + return os.path.exists(os.path.join(data_dir, "global", "pg_control")) + + def _clear_stale_data_dir(data_dir): """``initdb`` refuses to run against a non-empty target directory. A crash mid-init -- or a partial cleanup of an earlier failed start -- can leave an un-initialized data dir behind (e.g. just a ``log/`` subdir) that has no ``PG_VERSION`` yet still makes the fresh ``initdb`` fail, bricking - every subsequent start. Wipe such leftovers before initializing. + every subsequent start. Wipe such leftovers before initializing -- but never + a dir that still holds a real cluster: surface an error instead of deleting. """ if not (os.path.isdir(data_dir) and os.listdir(data_dir)): return + if _has_cluster_data(data_dir): + raise RuntimeError( + f"Refusing to wipe {data_dir}: it contains an existing PostgreSQL " + "cluster (global/pg_control present). Back it up or remove it " + "manually if you really want a fresh start." + ) import shutil logger.warning("Clearing incomplete PostgreSQL data dir %s before init", data_dir) shutil.rmtree(data_dir, ignore_errors=True) diff --git a/native-build/windows/embedded_pg.py b/native-build/windows/embedded_pg.py index b2bcd17e..c10beb88 100644 --- a/native-build/windows/embedded_pg.py +++ b/native-build/windows/embedded_pg.py @@ -53,9 +53,26 @@ def _initialized(data_dir): return False +def _has_cluster_data(data_dir): + """True if data_dir holds an initialized PostgreSQL cluster (never auto-delete it). + + Keyed on ``global/pg_control``: written at the END of initdb and required for + the server to start, so its presence proves a complete cluster with real data. + A half-built dir (interrupted initdb, no pg_control) holds no usable data and + is still cleared normally, so this guard never blocks first-run self-heal. + """ + return os.path.exists(os.path.join(data_dir, "global", "pg_control")) + + def _reset_data_dir(data_dir): if not (os.path.isdir(data_dir) and os.listdir(data_dir)): return + if _has_cluster_data(data_dir): + raise RuntimeError( + f"Refusing to wipe {data_dir}: it contains an existing PostgreSQL " + "cluster (global/pg_control present). Back it up or remove it " + "manually if you really want a fresh start." + ) logger.warning("Clearing incomplete PostgreSQL data dir %s before re-init", data_dir) for entry in os.listdir(data_dir): target = os.path.join(data_dir, entry) diff --git a/native-build/windows/env.py b/native-build/windows/env.py index af17ea7d..eaef7bbe 100644 --- a/native-build/windows/env.py +++ b/native-build/windows/env.py @@ -18,7 +18,6 @@ """ import os -import sys from urllib.parse import quote from windows import paths diff --git a/native-build/windows/launcher.py b/native-build/windows/launcher.py index 0203027f..07f9a0d5 100644 --- a/native-build/windows/launcher.py +++ b/native-build/windows/launcher.py @@ -62,7 +62,6 @@ def _win_wait4(pid, options): import os import runpy import signal -import subprocess import sys import threading import time @@ -90,7 +89,7 @@ def _run_flask(): import app as app_module waitress.serve( app_module.app, - host="127.0.0.1", + host="0.0.0.0", port=8000, threads=8, max_request_body_size=6 * 1024 * 1024 * 1024, diff --git a/rq_worker.py b/rq_worker.py index 7701cbe6..23a2060b 100644 --- a/rq_worker.py +++ b/rq_worker.py @@ -30,7 +30,6 @@ try: from app_helper import redis_conn from app_logging import configure_logging - import config from config import APP_VERSION, TEMP_DIR except ImportError as e: print(f"Error importing worker dependencies: {e}") diff --git a/sanitization.py b/sanitization.py new file mode 100644 index 00000000..e190a996 --- /dev/null +++ b/sanitization.py @@ -0,0 +1,111 @@ +# sanitization.py +"""Untrusted-input sanitization helpers. + +A dependency-free leaf module (standard library + numpy only) so any layer can +clean a value before it reaches the database or a JSON column without pulling in +project modules. Every data-sanitization helper in the project lives here: + +- ``sanitize_string_for_db`` -- strip NUL + control characters from a string. +- ``sanitize_db_field`` -- the above plus a length cap and whitespace strip, + for a single named DB column. +- ``sanitize_json_for_db`` -- recursively sanitize strings inside dict/list/tuple. +- ``sanitize_for_json`` -- convert numpy scalars/arrays to JSON-native types. +""" +import logging +import re +from typing import Optional + +import numpy as np + +logger = logging.getLogger(__name__) + + +def sanitize_string_for_db(value: Optional[str]) -> Optional[str]: + """Remove NUL bytes (0x00) and control characters from a string before a DB write. + + PostgreSQL TEXT/VARCHAR columns reject strings containing NUL bytes (0x00), + which can appear in corrupted metadata from music files. Returns ``None`` for + ``None``; non-strings are coerced with ``str()``. + """ + if value is None: + return None + + if not isinstance(value, str): + value = str(value) + + # Remove NUL bytes (0x00) + value = value.replace('\x00', '') + + # Remove other control characters (0x01-0x1F except tab, newline, carriage return) + value = re.sub(r'[\x01-\x08\x0B-\x0C\x0E-\x1F]', '', value) + + return value + + +def sanitize_db_field(s, max_length=1000, field_name="field"): + """Sanitize a single string column value for PostgreSQL insertion. + + Like :func:`sanitize_string_for_db` but also caps the length at + ``max_length`` and strips surrounding whitespace. ``field_name`` is only used + in the truncation log message. + """ + if s is None: + return None + + if not isinstance(s, str): + try: + s = str(s) + except Exception: + logger.warning(f"Could not convert {field_name} to string, using empty string") + return "" + + # Remove NUL byte (0x00) -- PostgreSQL cannot store it. + s = s.replace('\x00', '') + + # Keep only printable characters plus space/tab/newline. + s = ''.join(char for char in s if char.isprintable() or char in '\n\t ') + + if len(s) > max_length: + logger.warning(f"{field_name} truncated from {len(s)} to {max_length} characters") + s = s[:max_length] + + return s.strip() + + +def sanitize_json_for_db(value): + """Recursively sanitize strings inside a JSON-serializable value. + + Applies :func:`sanitize_string_for_db` to every string found inside dicts, + lists and tuples. Non-string scalars are returned as-is. Use this before + ``json.dumps(...)`` when the result is going into a Postgres jsonb column. + """ + if isinstance(value, str): + return sanitize_string_for_db(value) + if isinstance(value, dict): + return {k: sanitize_json_for_db(v) for k, v in value.items()} + if isinstance(value, list): + return [sanitize_json_for_db(v) for v in value] + if isinstance(value, tuple): + return tuple(sanitize_json_for_db(v) for v in value) + return value + + +def sanitize_for_json(obj): + """Recursively convert numpy arrays and numpy numeric types to native Python + types so the object is JSON serializable. + """ + if isinstance(obj, dict): + return {k: sanitize_for_json(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [sanitize_for_json(elem) for elem in obj] + elif isinstance(obj, np.ndarray): + return obj.tolist() + # Handle numpy numeric types which are not JSON serializable by default + elif isinstance(obj, (np.int_, np.intc, np.intp, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64)): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.bool_): + return bool(obj) + else: + return obj diff --git a/ssrf_guard.py b/ssrf_guard.py new file mode 100644 index 00000000..532c7c75 --- /dev/null +++ b/ssrf_guard.py @@ -0,0 +1,49 @@ +# ssrf_guard.py +"""SSRF guard for user-supplied outbound URLs. + +A dependency-free leaf module (standard library only), so any blueprint, helper, +or task can validate a URL without pulling in the database / task-queue / Flask +layers. Keep it that way -- nothing here may import a project module. +""" +import ipaddress +import socket +from urllib.parse import urlparse + + +def validate_outbound_url(url): + """SSRF guard for user-supplied outbound HTTP(S) URLs. + + Returns ``(True, None)`` when the URL is safe to fetch, else + ``(False, reason)``. + """ + if not url: + return False, 'URL is required' + try: + parsed = urlparse(str(url)) + except Exception: + return False, 'Invalid URL' + if parsed.scheme not in ('http', 'https'): + return False, 'Only http and https URLs are supported' + host = parsed.hostname + if not host: + return False, 'URL host is required' + try: + addrinfo = socket.getaddrinfo( + host, parsed.port or (443 if parsed.scheme == 'https' else 80), + type=socket.SOCK_STREAM, + ) + except Exception: + return False, 'Could not resolve host' + for entry in addrinfo: + try: + ip_obj = ipaddress.ip_address(entry[4][0]) + except ValueError: + return False, 'Resolved host to invalid IP' + if ( + ip_obj.is_link_local + or ip_obj.is_multicast + or ip_obj.is_reserved + or ip_obj.is_unspecified + ): + return False, 'Target host resolves to a disallowed IP address' + return True, None diff --git a/tasks/ai/api.py b/tasks/ai/api.py index 2b45c756..9d7012d7 100644 --- a/tasks/ai/api.py +++ b/tasks/ai/api.py @@ -34,11 +34,10 @@ import ftfy -from config import MAX_SONGS_IN_AI_PROMPT +import config from tasks.ai.providers import ( gemini as ai_api_gemini, mistral as ai_api_mistral, - ollama as ai_api_ollama, openai as ai_api_openai, ) from tasks.ai.prompts import build_mcp_system_prompt @@ -48,6 +47,9 @@ VALID_PROVIDERS = {"OLLAMA", "OPENAI", "GEMINI", "MISTRAL", "NONE"} +# --------------------------------------------------------------------------- + + # --------------------------------------------------------------------------- # Provider detection / validation # --------------------------------------------------------------------------- @@ -174,10 +176,11 @@ def generate_text(prompt: str, ai_config: Dict, *, skip_delay: bool = False, if provider == "NONE": return "AI Naming Skipped" if provider == "OLLAMA": - return ai_api_ollama.generate_text( + return ai_api_openai.generate_text( ai_config["ollama_url"], ai_config["ollama_model"], prompt, + api_key="no-key-needed", skip_delay=skip_delay, temperature=temperature, max_tokens=max_tokens, @@ -248,9 +251,7 @@ def call_with_tools( system_prompt = build_mcp_system_prompt(tools, library_context) if provider == "OLLAMA": - # Ollama builds its own JSON-output prompt internally; system_prompt is - # ignored because the Ollama prompt contains the system text already. - return ai_api_ollama.call_with_tools( + return ai_api_openai.call_with_tools_ollama( ai_config["ollama_url"], ai_config["ollama_model"], user_message, @@ -322,18 +323,19 @@ def get_ai_playlist_name( MAX_LENGTH = 40 # Truncate song list to avoid token-limit issues - songs_for_prompt = song_list[:MAX_SONGS_IN_AI_PROMPT] + max_songs = config.MAX_SONGS_IN_AI_PROMPT + songs_for_prompt = song_list[:max_songs] formatted_song_list = "\n".join( [ f"- {song.get('title', 'Unknown Title')} by {song.get('author', 'Unknown Artist')}" for song in songs_for_prompt ] ) - if len(song_list) > MAX_SONGS_IN_AI_PROMPT: + if len(song_list) > max_songs: logger.info( "Truncated song list from %d to %d songs for AI prompt to avoid token limits", len(song_list), - MAX_SONGS_IN_AI_PROMPT, + max_songs, ) full_prompt = prompt_template.format(song_list_sample=formatted_song_list) diff --git a/tasks/ai/prompts.py b/tasks/ai/prompts.py index a4191a62..9892a524 100644 --- a/tasks/ai/prompts.py +++ b/tasks/ai/prompts.py @@ -14,14 +14,14 @@ build_ollama_tool_calling_prompt -- Ollama-specific JSON-output framing build_tool_calls_schema(tools) -- shared JSON Schema for tool_calls (used by every transport) build_ai_brainstorm_prompt(...) -- ai_brainstorm MCP tool prompt - get_dynamic_genres -- helper, exposed for tests + _get_dynamic_genres -- helper, exposed for tests """ from typing import Dict, List, Optional import config -# --- Clustering / playlist naming --------------------------------------------- +# --- Clustering / playlist naming prompt -------------------------------------- creative_prompt_template = ( "You are an expert music collector and MUST give a title to this playlist.\n" @@ -31,35 +31,33 @@ "No special fonts or emojis.\n" "* BAD EXAMPLES: 'Ambient Electronic Space - Electric Soundscapes - Emotional Waves' (Too long/descriptive)\n" "* BAD EXAMPLES: 'Blues Rock Fast Tracks' (Too direct/literal, not evocative enough)\n" - "* BAD EXAMPLES: '\U0001D5DD\U0001D5C2\U0001D5C8 \U0001D5C2\U0001D5CB\U0001D5C8\U0001D5C7\U0001D5C2 \U0001D5C9\U0001D5CB\U0001D5C8\U0001D5C7\U0001D5C2' (Non-standard characters)\n\n" + "* BAD EXAMPLES: '\\U0001D5DD\\U0001D5C2\\U0001D5C8 \\U0001D5C2\\U0001D5CB\\U0001D5C8\\U0001D5C7\\U0001D5C2 \\U0001D5C9\\U0001D5CB\\U0001D5C8\\U0001D5C7\\U0001D5C2' (Non-standard characters)\n\n" "CRITICAL: Your response MUST be ONLY the single playlist name. No explanations, no 'Playlist Name:', no numbering, no extra text or formatting whatsoever.\n\n" "This is the playlist:\n{song_list_sample}\n\n" ) -# --- MCP system prompt (used by all providers when calling with tools) -------- - -_FALLBACK_GENRES = ( - "rock, pop, metal, jazz, electronic, dance, alternative, indie, punk, blues, " - "hard rock, heavy metal, hip-hop, funk, country, soul" -) +# --- Constants shared with the rest of the AI subsystem ----------------------- +INTENT_CLASSES = ["seed", "text", "knowledge", "metadata"] -def get_dynamic_genres(library_context: Optional[Dict]) -> str: - """Return genre list from library context, falling back to defaults.""" - if library_context and library_context.get('top_genres'): - return ', '.join(library_context['top_genres'][:10]) - return _FALLBACK_GENRES +PRIMARY_INTENTS = ["seed", "text", "knowledge"] -VOICE_VOCAB = ["female vocalists", "female vocalist", "male vocalists"] +# --- MCP system prompt -------------------------------------------------------- -INTENT_CLASSES = ["seed", "text", "knowledge", "metadata"] -PRIMARY_INTENTS = ["seed", "text", "knowledge"] +def _get_dynamic_genres(library_context: Optional[Dict]) -> str: + """Return genre list from library context, falling back to defaults.""" + if library_context and library_context.get('top_genres'): + return ', '.join(library_context['top_genres'][:10]) + return config.AI_FALLBACK_GENRES -def build_mcp_system_prompt(tools: List[Dict], library_context: Optional[Dict] = None) -> str: +def build_mcp_system_prompt( + tools: List[Dict], + library_context: Optional[Dict] = None, +) -> str: """Build the canonical MCP system prompt used by ALL providers.""" tool_names = {t['name'] for t in tools} has_seed = 'seed_search' in tool_names @@ -81,8 +79,10 @@ def build_mcp_system_prompt(tools: List[Dict], library_context: Optional[Dict] = ) if has_knowledge: tool_lines.append( - "- knowledge_lookup(user_request): cultural/historical world-knowledge fallback " - "('Grammy winners 2020', 'songs sampled by Daft Punk'). LAST RESORT." + "- knowledge_lookup(user_request): popularity / 'best of' / cultural requests " + "('best rap of the 90s', 'top festival anthems', 'Grammy winners 2020'). Turns the " + "request into a grounded library search (genre/year/energy + sound descriptions + " + "seed artists); never invents song titles." ) tool_lines.append( "- search_database(genres?, voices?, moods?, year_min?, year_max?, min_rating?, scale?, " @@ -91,8 +91,8 @@ def build_mcp_system_prompt(tools: List[Dict], library_context: Optional[Dict] = ) tools_block = "\n".join(tool_lines) - genres_line = get_dynamic_genres(library_context) - voices_line = ", ".join(VOICE_VOCAB) + genres_line = _get_dynamic_genres(library_context) + voices_line = ", ".join(config.VOICE_VOCAB) moods_line = ", ".join(config.OTHER_FEATURE_LABELS) prompt = f"""You are a music playlist router. Return ONLY a JSON object with one or more tool calls. Put EVERY intent in this one response. @@ -263,32 +263,42 @@ def build_tool_calls_schema(tools: List[Dict]) -> Dict: # --- Free-text MCP prompts (brainstorm) --------------------------------------- def build_ai_brainstorm_prompt(user_request: str) -> str: - return f"""You are a music expert with extensive knowledge of songs, artists, and music history. + """Build the grounded-recipe prompt for the knowledge_lookup tool. + + The model does NOT know the library, so it must not name songs. It translates + the request into a search RECIPE -- metadata filters, "how it should sound" + descriptions (for audio similarity search), and a few seed artists -- which + the tool then runs against the real library and fuses. This trades the small + model's weak song recall for its strong understanding/categorisation. + """ + genres_line = ", ".join(config.STRATIFIED_GENRES) + moods_line = ", ".join(config.OTHER_FEATURE_LABELS) + voices_line = ", ".join(config.VOICE_VOCAB) + return f"""You are a music expert. Turn the request into a RECIPE used to search a music library. +You do NOT know which songs are in the library, so you MUST NOT name any songs. Describe and categorise only; the library does the finding. User request: "{user_request}" -TASK: Use your knowledge to suggest 25-35 specific songs (with exact artist names) that match this request. - -Think about: -- If they want songs similar to an artist \u2192 suggest songs by that artist AND similar artists -- If they want a genre/mood \u2192 suggest famous songs in that genre/mood -- If they want popular/radio hits \u2192 suggest well-known mainstream songs -- If they want a time period \u2192 suggest songs from that era -- If they want a vibe \u2192 suggest songs that match that feeling - -CRITICAL REQUIREMENTS: -1. Return ONLY a JSON array of objects -2. Each object MUST have "title" and "artist" fields -3. Be specific with exact song titles and artist names (as they appear in databases) -4. Include variety - different artists when possible -5. Format: [{{"title": "Song Name", "artist": "Artist Name"}}, ...] -6. NO explanations, NO numbering, ONLY the JSON array - -Example format: -[ - {{"title": "All the Small Things", "artist": "blink-182"}}, - {{"title": "Basket Case", "artist": "Green Day"}}, - {{"title": "American Idiot", "artist": "Green Day"}} -] - -Suggest songs for "{user_request}" now:""" +Return ONE JSON object with EXACTLY this shape: +{{"filters": {{"genres": [], "moods": [], "voices": [], "year_min": null, "year_max": null, "energy_min": null, "energy_max": null, "tempo_min": null, "tempo_max": null}}, "sound_descriptions": [], "seed_artists": [], "lyric_themes": []}} + +FIELD GUIDE (leave a field empty/null when the request does not imply it -- never invent constraints): +- filters.genres: 0+ values, chosen ONLY from: {genres_line} +- filters.moods: 0+ values, chosen ONLY from: {moods_line} +- filters.voices: 0+ values, chosen ONLY from: {voices_line} +- filters.year_min / year_max: 4-digit years. A decade like "90s" -> 1990 and 1999. "90s and 2000s" -> 1990 and 2009. +- filters.energy_min / energy_max: numbers 0.0 (calm) to 1.0 (intense). +- filters.tempo_min / tempo_max: BPM, 40 to 200. +- sound_descriptions: 2 to {config.AI_BRAINSTORM_SOUND_DESCRIPTIONS_MAX} vivid phrases describing HOW the ideal songs SOUND (instruments, production, era, energy, vibe). This is the most important field. NOT song names. +- seed_artists: up to {config.AI_BRAINSTORM_SEED_ARTISTS_MAX} well-known ARTISTS that exemplify the request. Artists ONLY, never songs. Omit if none are obvious. +- lyric_themes: 0 to {config.AI_BRAINSTORM_LYRIC_THEMES_MAX} short phrases ONLY when the request is about a TOPIC the lyrics should cover (e.g. "heartbreak", "summer roadtrip"). + +RULES: +- NEVER output a song title anywhere. +- genres / moods / voices MUST come from the lists above, or be left empty. +- Output ONLY the JSON object. No markdown fences, no comments, no extra text. + +EXAMPLE -- request "100 of the best rap songs from the 90s and 2000s": +{{"filters": {{"genres": ["Hip-Hop"], "moods": [], "voices": [], "year_min": 1990, "year_max": 2009, "energy_min": 0.5, "energy_max": 1.0, "tempo_min": null, "tempo_max": null}}, "sound_descriptions": ["gritty 90s east coast boom bap hip hop with hard-hitting drums and jazzy samples", "glossy early 2000s mainstream rap with heavy bass and crossover hooks"], "seed_artists": ["Nas", "Jay-Z", "2Pac", "Eminem"], "lyric_themes": []}} + +Now produce the JSON recipe for "{user_request}":""" diff --git a/tasks/ai/providers/ollama.py b/tasks/ai/providers/ollama.py deleted file mode 100644 index 80c7a162..00000000 --- a/tasks/ai/providers/ollama.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Ollama transport. - -Ollama lacks native function calling, so `call_with_tools` builds a JSON-output -prompt (via `tasks.ai.prompts.build_ollama_tool_calling_prompt`) and parses the -model's JSON response. `generate_text` delegates to the OpenAI-compatible -streaming code path (since Ollama also exposes /api/generate streaming SSE). -""" -import json -import logging -import re -from typing import Dict, List, Optional - -import httpx - -import config -from tasks.ai.providers import openai as ai_api_openai -from tasks.ai.prompts import build_ollama_tool_calling_prompt, build_tool_calls_schema - -logger = logging.getLogger(__name__) - - -def generate_text( - ollama_url: str, - model_name: str, - full_prompt: str, - *, - skip_delay: bool = False, - temperature: Optional[float] = None, - max_tokens: Optional[int] = None, -) -> str: - """Generate freeform text from an Ollama /api/generate endpoint. - - Reuses the OpenAI-compatible transport because the streaming code path is - shared (it auto-detects Ollama format from the URL). - """ - return ai_api_openai.generate_text( - ollama_url, model_name, full_prompt, api_key="no-key-needed", - skip_delay=skip_delay, temperature=temperature, max_tokens=max_tokens, - ) - - -def call_with_tools( - ollama_url: str, - model_name: str, - user_message: str, - tools: List[Dict], - log_messages: List[str], - library_context: Optional[Dict] = None, -) -> Dict: - """Prompt-based tool calling for Ollama (no native function calling). - - Returns ``{"tool_calls": [...]}`` on success, ``{"error": "..."}`` on failure. - """ - try: - prompt = build_ollama_tool_calling_prompt(user_message, tools, library_context) - - schema = build_tool_calls_schema(tools) - payload = { - "model": model_name, - "prompt": prompt, - "stream": False, - "format": schema, - "think": False, - # Cap how much the model can think/generate so it can't run on - # forever. A tool call needs only a few hundred tokens. - "options": {"temperature": 0, "num_predict": 1024}, - } - - timeout = config.AI_REQUEST_TIMEOUT_SECONDS - log_messages.append(f"Using timeout: {timeout} seconds for Ollama request") - # Single bounded call: the httpx read timeout aborts it at `timeout` - # seconds so it can never run forever. NO retry -- if the model returns - # nothing usable, we error out and the chat pipeline falls back (the user - # still gets a playlist) rather than making a second multi-minute call. - with httpx.Client(timeout=timeout) as client: - response = client.post(ollama_url, json=payload) - response.raise_for_status() - result = response.json() - - if "response" not in result: - return {"error": "Invalid Ollama response"} - - response_text = result["response"] - - cleaned = "" - try: - cleaned = response_text.strip() - cleaned = re.sub(r".*?", "", cleaned, flags=re.DOTALL).strip() - if "" in cleaned: - cleaned = ( - cleaned.split("")[-1].strip() - if "" in cleaned - else re.sub(r".*", "", cleaned, flags=re.DOTALL).strip() - ) - - log_messages.append(f"Ollama raw response (first 300 chars): {cleaned[:300]}") - - if "```json" in cleaned: - cleaned = cleaned.split("```json")[1].split("```")[0] - elif "```" in cleaned: - cleaned = cleaned.split("```")[1].split("```")[0] - cleaned = cleaned.strip() - - if cleaned.startswith("{") and '"type"' in cleaned and '"array"' in cleaned: - log_messages.append( - "\u26a0\ufe0f Ollama returned schema instead of tool calls, using fallback" - ) - return {"error": "Ollama returned schema definition instead of tool calls"} - - log_messages.append(f"Attempting to parse: {cleaned[:200]}") - parsed = json.loads(cleaned) - - if isinstance(parsed, dict) and "tool_calls" in parsed: - tool_calls = parsed["tool_calls"] - log_messages.append( - f"\u2713 Extracted tool_calls array with {len(tool_calls) if isinstance(tool_calls, list) else 1} items" - ) - elif isinstance(parsed, list): - tool_calls = parsed - log_messages.append("\u26a0\ufe0f Got array directly (expected object with tool_calls field)") - elif isinstance(parsed, dict) and "name" in parsed: - tool_calls = [parsed] - log_messages.append( - "\u26a0\ufe0f Got single tool call object (expected object with tool_calls array)" - ) - elif isinstance(parsed, dict) and "tool" in parsed and "arguments" in parsed: - tool_calls = [{"name": parsed["tool"], "arguments": parsed["arguments"]}] - log_messages.append( - "\u26a0\ufe0f Remapped {'tool','arguments'} -> {'name','arguments'} format" - ) - else: - log_messages.append( - f"\u26a0\ufe0f Unexpected JSON structure: {type(parsed)}, keys: {list(parsed.keys()) if isinstance(parsed, dict) else 'N/A'}" - ) - return {"error": "Ollama response missing 'tool_calls' field"} - - if not isinstance(tool_calls, list): - tool_calls = [tool_calls] - - valid_calls = [] - for tc in tool_calls: - if isinstance(tc, dict) and "name" in tc: - if "arguments" not in tc: - tc["arguments"] = {} - args = tc["arguments"] - keys_to_remove = [] - for k, v in args.items(): - if v is None or v == "" or v == [] or v == {}: - keys_to_remove.append(k) - elif k in ("tempo_min", "tempo_max", "energy_min", "min_rating") and v == 0: - keys_to_remove.append(k) - for k in keys_to_remove: - log_messages.append( - f" \U0001f9f9 Stripped empty/default arg '{k}={args[k]}' from {tc['name']}" - ) - del args[k] - valid_calls.append(tc) - else: - log_messages.append(f"\u26a0\ufe0f Skipping invalid tool call: {tc}") - - if not valid_calls: - return {"error": "No valid tool calls found in Ollama response"} - - log_messages.append(f"\u2705 Ollama returned {len(valid_calls)} valid tool calls") - return {"tool_calls": valid_calls} - - except json.JSONDecodeError: - logger.exception("JSON decode error while parsing Ollama tool response") - log_messages.append("\u274c Failed to parse Ollama JSON response.") - log_messages.append(f"Attempted to parse: {cleaned[:300]}") - return { - "error": "Failed to parse Ollama JSON response.", - "raw_response": response_text[:200], - } - except Exception: - logger.exception("Failed to parse Ollama response") - log_messages.append("Failed to parse Ollama response.") - log_messages.append(f"Response was: {response_text[:200]}") - return {"error": "Failed to parse Ollama tool calls", "raw_response": response_text} - - except httpx.ReadTimeout: - timeout = config.AI_REQUEST_TIMEOUT_SECONDS - logger.warning(f"Ollama request timed out after {timeout} seconds") - log_messages.append( - f"\u23f1\ufe0f Ollama request timed out after {timeout} seconds. Your model or hardware may be too slow." - ) - log_messages.append( - "\U0001f4a1 Solution: Set AI_REQUEST_TIMEOUT_SECONDS environment variable to a higher value (e.g., 600 for 10 minutes)" - ) - return { - "error": f"Ollama timed out after {timeout} seconds. Increase AI_REQUEST_TIMEOUT_SECONDS for slower hardware or larger models." - } - except httpx.TimeoutException: - timeout = config.AI_REQUEST_TIMEOUT_SECONDS - logger.warning("Ollama request timed out", exc_info=True) - log_messages.append(f"\u23f1\ufe0f Ollama request timed out after {timeout} seconds.") - log_messages.append( - "\U0001f4a1 Solution: Set AI_REQUEST_TIMEOUT_SECONDS environment variable to a higher value" - ) - return { - "error": f"Ollama timed out after {timeout} seconds. Increase AI_REQUEST_TIMEOUT_SECONDS for slower hardware or larger models." - } - except Exception: - logger.exception("Error calling Ollama with tools") - return {"error": "Ollama service is currently unavailable."} diff --git a/tasks/ai/providers/openai.py b/tasks/ai/providers/openai.py index 749c69f2..168babb0 100644 --- a/tasks/ai/providers/openai.py +++ b/tasks/ai/providers/openai.py @@ -1,15 +1,17 @@ -"""OpenAI-compatible transport (OpenAI, OpenRouter, anything speaking /v1/chat/completions). +"""OpenAI-compatible transport (OpenAI, OpenRouter, Ollama, anything speaking /v1/chat/completions or /api/generate). -Two public functions: - generate_text(...) -- single-prompt streaming completion (used for playlist naming, brainstorm, etc.) - call_with_tools(...) -- non-streaming chat with tool/function calling +Three public functions: + generate_text(...) -- single-prompt streaming completion (all providers) + call_with_tools(...) -- native function/tool calling (OpenAI, OpenRouter, etc.) + call_with_tools_ollama(...) -- prompt-based JSON tool calling (Ollama, no native support) These transports only handle HTTP plumbing. All business prompts come from -`tasks/ai_prompts.py`. +`tasks/ai/prompts.py`. """ import json import logging import os +import re import time from typing import Dict, List, Optional @@ -20,6 +22,8 @@ logger = logging.getLogger(__name__) +THINK_END_TAG = "" + def _is_ollama_format_url(server_url: str) -> bool: """Detect Ollama endpoints from the URL path (issue #467 fix preserved).""" @@ -54,8 +58,8 @@ def generate_text( max_completion_tokens, then drop max_completion_tokens), and content extraction-before-finish-reason ordering for OpenRouter compatibility. - NOTE: Detects Ollama vs OpenAI from the URL for backward compatibility with - `tasks/ai_api_ollama.generate_text` which delegates here. + NOTE: Detects Ollama vs OpenAI from the URL; the Ollama branch in + tasks/ai/api.py calls this directly (Ollama has no separate transport). """ is_ollama_format = _is_ollama_format_url(server_url) is_openai_format = not is_ollama_format @@ -160,7 +164,7 @@ def generate_text( logger.debug("Could not decode JSON line from stream: %s", line_str) continue - thought_enders = ["", "[/INST]", "[/THOUGHT]"] + thought_enders = [THINK_END_TAG, "[/INST]", "[/THOUGHT]"] extracted_text = full_raw_response_content.strip() for end_tag in thought_enders: if end_tag in extracted_text: @@ -394,3 +398,185 @@ def _post(p): except Exception: logger.exception("Error calling OpenAI with tools") return {"error": "OpenAI service is currently unavailable."} + + +# --------------------------------------------------------------------------- +# Ollama prompt-based tool calling (Ollama lacks native function calling) +# --------------------------------------------------------------------------- + +def call_with_tools_ollama( + ollama_url: str, + model_name: str, + user_message: str, + tools: List[Dict], + log_messages: List[str], + library_context: Optional[Dict] = None, +) -> Dict: + """Prompt-based tool calling for Ollama via /api/generate with structured output. + + Ollama has no native function/tool calling. We build a JSON-output prompt + (via ``tasks.ai.prompts.build_ollama_tool_calling_prompt``) and set + ``format`` to the tool-calls JSON Schema so the model is forced to emit + valid JSON that we parse into ``{"tool_calls": [...]}``. + + Returns ``{"tool_calls": [...]}`` on success, ``{"error": "..."}`` on failure. + """ + from tasks.ai.prompts import build_ollama_tool_calling_prompt, build_tool_calls_schema # noqa: E402 + + try: + prompt = build_ollama_tool_calling_prompt(user_message, tools, library_context) + + schema = build_tool_calls_schema(tools) + payload = { + "model": model_name, + "prompt": prompt, + "stream": False, + "format": schema, + "think": False, + # Cap generation so a model can't run forever. A tool call needs + # only a few hundred tokens. + "options": {"temperature": 0, "num_predict": 1024}, + } + + timeout = config.AI_REQUEST_TIMEOUT_SECONDS + log_messages.append(f"Using timeout: {timeout} seconds for Ollama request") + # Single bounded call: the httpx read timeout aborts it at `timeout` + # seconds so it can never run forever. NO retry -- if the model returns + # nothing usable, we error out and the chat pipeline falls back (the user + # still gets a playlist) rather than making a second multi-minute call. + with httpx.Client(timeout=timeout) as client: + response = client.post(ollama_url, json=payload) + response.raise_for_status() + result = response.json() + + if "response" not in result: + return {"error": "Invalid Ollama response"} + + response_text = result["response"] + + cleaned = "" + try: + cleaned = response_text.strip() + cleaned = re.sub(r".*?", "", cleaned, flags=re.DOTALL).strip() + if "" in cleaned: + cleaned = ( + cleaned.split(THINK_END_TAG)[-1].strip() + if THINK_END_TAG in cleaned + else re.sub(r".*", "", cleaned, flags=re.DOTALL).strip() + ) + + log_messages.append(f"Ollama raw response (first 300 chars): {cleaned[:300]}") + + if "```json" in cleaned: + cleaned = cleaned.split("```json")[1].split("```")[0] + elif "```" in cleaned: + cleaned = cleaned.split("```")[1].split("```")[0] + cleaned = cleaned.strip() + + if cleaned.startswith("{") and '"type"' in cleaned and '"array"' in cleaned: + log_messages.append( + "\u26a0\ufe0f Ollama returned schema instead of tool calls, using fallback" + ) + return {"error": "Ollama returned schema definition instead of tool calls"} + + log_messages.append(f"Attempting to parse: {cleaned[:200]}") + parsed = json.loads(cleaned) + + if isinstance(parsed, dict) and "tool_calls" in parsed: + tool_calls = parsed["tool_calls"] + log_messages.append( + f"\u2713 Extracted tool_calls array with {len(tool_calls) if isinstance(tool_calls, list) else 1} items" + ) + elif isinstance(parsed, list): + tool_calls = parsed + log_messages.append("\u26a0\ufe0f Got array directly (expected object with tool_calls field)") + elif isinstance(parsed, dict) and "name" in parsed: + tool_calls = [parsed] + log_messages.append( + "\u26a0\ufe0f Got single tool call object (expected object with tool_calls array)" + ) + elif isinstance(parsed, dict) and "tool" in parsed and "arguments" in parsed: + tool_calls = [{"name": parsed["tool"], "arguments": parsed["arguments"]}] + log_messages.append( + "\u26a0\ufe0f Remapped {'tool','arguments'} -> {'name','arguments'} format" + ) + else: + log_messages.append( + f"\u26a0\ufe0f Unexpected JSON structure: {type(parsed)}, keys: {list(parsed.keys()) if isinstance(parsed, dict) else 'N/A'}" + ) + return {"error": "Ollama response missing 'tool_calls' field"} + + if not isinstance(tool_calls, list): + tool_calls = [tool_calls] + + valid_calls = [] + for tc in tool_calls: + if isinstance(tc, dict) and "name" in tc: + if "arguments" not in tc: + tc["arguments"] = {} + elif not isinstance(tc["arguments"], dict): + log_messages.append( + f"Coerced non-dict arguments for tool '{tc['name']}' to empty dict" + ) + tc["arguments"] = {} + args = tc["arguments"] + keys_to_remove = [] + for k, v in args.items(): + if (v is None or v == "" or v == [] or v == {}) or ( + k in ("tempo_min", "tempo_max", "energy_min", "min_rating") and v == 0 + ): + keys_to_remove.append(k) + for k in keys_to_remove: + log_messages.append( + f" \U0001f9f9 Stripped empty/default arg '{k}={args[k]}' from {tc['name']}" + ) + del args[k] + valid_calls.append(tc) + else: + log_messages.append(f"\u26a0\ufe0f Skipping invalid tool call: {tc}") + + if not valid_calls: + return {"error": "No valid tool calls found in Ollama response"} + + log_messages.append(f"\u2705 Ollama returned {len(valid_calls)} valid tool calls") + return {"tool_calls": valid_calls} + + except json.JSONDecodeError: + logger.exception("JSON decode error while parsing Ollama tool response") + log_messages.append("\u274c Failed to parse Ollama JSON response.") + log_messages.append(f"Attempted to parse: {cleaned[:300]}") + return { + "error": "Failed to parse Ollama JSON response.", + "raw_response": response_text[:200], + } + except Exception: + logger.exception("Failed to parse Ollama response") + log_messages.append("Failed to parse Ollama response.") + log_messages.append(f"Response was: {response_text[:200]}") + return {"error": "Failed to parse Ollama tool calls", "raw_response": response_text} + + except httpx.ReadTimeout: + timeout = config.AI_REQUEST_TIMEOUT_SECONDS + logger.warning(f"Ollama request timed out after {timeout} seconds") + log_messages.append( + f"\u23f1\ufe0f Ollama request timed out after {timeout} seconds. Your model or hardware may be too slow." + ) + log_messages.append( + "\U0001f4a1 Solution: Set AI_REQUEST_TIMEOUT_SECONDS environment variable to a higher value (e.g., 600 for 10 minutes)" + ) + return { + "error": f"Ollama timed out after {timeout} seconds. Increase AI_REQUEST_TIMEOUT_SECONDS for slower hardware or larger models." + } + except httpx.TimeoutException: + timeout = config.AI_REQUEST_TIMEOUT_SECONDS + logger.warning("Ollama request timed out", exc_info=True) + log_messages.append(f"\u23f1\ufe0f Ollama request timed out after {timeout} seconds.") + log_messages.append( + "\U0001f4a1 Solution: Set AI_REQUEST_TIMEOUT_SECONDS environment variable to a higher value" + ) + return { + "error": f"Ollama timed out after {timeout} seconds. Increase AI_REQUEST_TIMEOUT_SECONDS for slower hardware or larger models." + } + except Exception: + logger.exception("Error calling Ollama with tools") + return {"error": "Ollama service is currently unavailable."} diff --git a/tasks/ai/tool_impl.py b/tasks/ai/tool_impl.py index c2795479..d290a456 100644 --- a/tasks/ai/tool_impl.py +++ b/tasks/ai/tool_impl.py @@ -949,149 +949,304 @@ def _lyrics_search_sync(query: str, get_songs: int) -> Dict: return {"songs": [], "message": f"lyrics_search error: {str(e)[:200]}"} -def _ai_brainstorm_sync(user_request: str, ai_config: Dict, get_songs: int) -> Dict: - """Use AI to brainstorm songs from world knowledge for any free-form request.""" - from tasks.ai.api import generate_text as _ai_generate_text - from tasks.ai.prompts import build_ai_brainstorm_prompt - - get_songs = int(get_songs) if get_songs is not None else 100 - - db_conn = get_db_connection() - log_messages = [] +def _extract_json_object(raw: str) -> Optional[Dict]: + """Best-effort recovery of a single JSON object from a model response. + Strips markdown code fences and any leading ``...`` preamble, + then parses the outermost ``{...}`` span. Returns the parsed dict, or ``None`` + when no JSON object can be recovered. + """ + if not raw: + return None + text = raw.strip() + if "" in text: + text = text.split("")[-1] + if "```json" in text: + text = text.split("```json", 1)[1].split("```", 1)[0] + elif "```" in text: + text = text.split("```", 1)[1].split("```", 1)[0] + text = text.strip() try: - log_messages.append(f"Using AI knowledge to brainstorm songs for: {user_request}") + whole = json.loads(text) + return whole if isinstance(whole, dict) else None + except (ValueError, TypeError): + pass + start = text.find("{") + end = text.rfind("}") + if start == -1 or end == -1 or end <= start: + return None + try: + parsed = json.loads(text[start:end + 1]) + except (ValueError, TypeError): + return None + return parsed if isinstance(parsed, dict) else None - prompt = build_ai_brainstorm_prompt(user_request) - raw_response = _ai_generate_text(prompt, ai_config, skip_delay=True) +def _clamp_recipe(recipe: Dict) -> Dict: + """Normalise a raw brainstorm recipe to safe, library-valid values. + + Clamps genres/moods/voices to the known vocab (case/punctuation-insensitive), + coerces numeric ranges into bounds and repairs reversed min/max, and caps the + list fields by their config limits. Always returns the full set of keys so the + executor never has to defend against missing fields. + """ + import config + + def _norm(s): + return re.sub(r"[^a-z0-9]", "", str(s).lower()) + + def _clamp_to_vocab(values, vocab): + canon = {_norm(v): v for v in vocab} + out, seen = [], set() + for v in values or []: + key = _norm(v) + if key and key in canon and key not in seen: + out.append(canon[key]) + seen.add(key) + return out - if raw_response.startswith("Error:"): - return {"songs": [], "message": f"AI Error: {raw_response}"} + def _as_list(v): + if isinstance(v, list): + return v + if v in (None, ""): + return [] + return [v] + def _num(v, lo, hi): try: - cleaned = raw_response.strip() - - if "```json" in cleaned: - cleaned = cleaned.split("```json")[1].split("```")[0] - elif "```" in cleaned: - cleaned = cleaned.split("```")[1].split("```")[0] + n = float(v) + except (TypeError, ValueError): + return None + return max(lo, min(hi, n)) + + def _clean_strings(values, cap): + out, seen = [], set() + for v in values or []: + s = str(v).strip() + key = s.lower() + if s and key not in seen: + out.append(s) + seen.add(key) + if len(out) >= cap: + break + return out - cleaned = cleaned.strip() + raw_filters = recipe.get("filters") if isinstance(recipe.get("filters"), dict) else {} - if "[" in cleaned and "]" in cleaned: - start = cleaned.find("[") - end = cleaned.rfind("]") + 1 - cleaned = cleaned[start:end] + year_min = _num(raw_filters.get("year_min"), 1900, 2100) + year_max = _num(raw_filters.get("year_max"), 1900, 2100) + year_min = int(year_min) if year_min is not None else None + year_max = int(year_max) if year_max is not None else None + if year_min is not None and year_max is not None and year_min > year_max: + year_min, year_max = year_max, year_min - cleaned = cleaned.replace("'\'", '"') + energy_min = _num(raw_filters.get("energy_min"), 0.0, 1.0) + energy_max = _num(raw_filters.get("energy_max"), 0.0, 1.0) + if energy_min is not None and energy_max is not None and energy_min > energy_max: + energy_min, energy_max = energy_max, energy_min - song_list = json.loads(cleaned) + tempo_min = _num(raw_filters.get("tempo_min"), config.TEMPO_MIN_BPM, config.TEMPO_MAX_BPM) + tempo_max = _num(raw_filters.get("tempo_max"), config.TEMPO_MIN_BPM, config.TEMPO_MAX_BPM) + if tempo_min is not None and tempo_max is not None and tempo_min > tempo_max: + tempo_min, tempo_max = tempo_max, tempo_min - if not isinstance(song_list, list): - raise ValueError("Response is not a JSON array") + seed_artists = ( + _clean_strings(_as_list(recipe.get("seed_artists")), config.AI_BRAINSTORM_SEED_ARTISTS_MAX) + if config.AI_BRAINSTORM_USE_ARTIST_SEEDS else [] + ) - log_messages.append(f"AI suggested {len(song_list)} songs") - except Exception as e: - log_messages.append(f"Failed to parse AI response: {str(e)}") - log_messages.append(f"Raw AI response (first 500 chars): {raw_response[:500]}") - return {"songs": [], "message": "\n".join(log_messages)} - - found_songs = [] - seen_ids = set() + return { + "filters": { + "genres": _clamp_to_vocab(_as_list(raw_filters.get("genres")), config.STRATIFIED_GENRES), + "moods": _clamp_to_vocab(_as_list(raw_filters.get("moods")), config.OTHER_FEATURE_LABELS), + "voices": _clamp_to_vocab(_as_list(raw_filters.get("voices")), config.VOICE_VOCAB), + "year_min": year_min, + "year_max": year_max, + "energy_min": energy_min, + "energy_max": energy_max, + "tempo_min": tempo_min, + "tempo_max": tempo_max, + }, + "sound_descriptions": _clean_strings( + _as_list(recipe.get("sound_descriptions")), config.AI_BRAINSTORM_SOUND_DESCRIPTIONS_MAX), + "seed_artists": seed_artists, + "lyric_themes": _clean_strings( + _as_list(recipe.get("lyric_themes")), config.AI_BRAINSTORM_LYRIC_THEMES_MAX), + } - def _normalize(s: str) -> str: - return re.sub(r"[\s\-\u2010\u2011\u2012\u2013\u2014/'\".,!?()]", '', s).lower() - def _escape_like(s: str) -> str: - return s.replace('%', r'\%').replace('_', r'\_') +def _ai_brainstorm_sync(user_request: str, ai_config: Dict, get_songs: int) -> Dict: + """Grounded brainstorm: the model emits a search RECIPE, not song titles. + + The recipe (metadata filters + "how it sounds" descriptions + seed artists + + lyric themes) is executed against the real library through the existing + grounded channels (CLAP audio search, artist similarity, lyrics search, + metadata filter) and the results are fused. Small models recall specific songs + poorly but understand and categorise requests well, so this keeps the catalog + external to the model's weights (issue #643). Grounding happens inside this + tool, so the planner still returns the result as-is. + """ + import config + from tasks.ai.api import generate_text as _ai_generate_text + from tasks.ai.prompts import build_ai_brainstorm_prompt - valid_items = [(item.get('title', ''), item.get('artist', '')) - for item in song_list - if item.get('title') and item.get('artist')] + get_songs = int(get_songs) if get_songs is not None else 200 + log_messages = [f"Brainstorming a grounded search recipe for: {user_request}"] + + prompt = build_ai_brainstorm_prompt(user_request) + raw_response = _ai_generate_text(prompt, ai_config, skip_delay=True, max_tokens=1500) + + if raw_response.startswith("Error:"): + logger.warning("Brainstorm AI call failed: %s", raw_response) + return {"songs": [], "message": "AI brainstorm failed; check the container logs."} + + parsed = _extract_json_object(raw_response) + if parsed is None: + logger.warning("Brainstorm recipe parse failed. Raw response (first 2000 chars): %s", raw_response[:2000]) + return {"songs": [], "message": "AI brainstorm could not produce a recipe; check the container logs."} + + recipe = _clamp_recipe(parsed) + filt = recipe["filters"] + + log_messages.append( + "Recipe: genres={g} moods={m} voices={v} year={y0}-{y1} energy={e0}-{e1} | " + "descriptions={nd} artists={na} lyric_themes={nl}".format( + g=filt["genres"] or "-", m=filt["moods"] or "-", v=filt["voices"] or "-", + y0=filt["year_min"] if filt["year_min"] is not None else "any", + y1=filt["year_max"] if filt["year_max"] is not None else "any", + e0=filt["energy_min"] if filt["energy_min"] is not None else "any", + e1=filt["energy_max"] if filt["energy_max"] is not None else "any", + nd=len(recipe["sound_descriptions"]), na=len(recipe["seed_artists"]), + nl=len(recipe["lyric_themes"]), + ) + ) + + found_songs: List[Dict] = [] + seen_ids: set = set() + seen_keys: set = set() + + def _add_one(s): + iid = s.get("item_id") + if not iid or iid in seen_ids: + return False + key = (s.get("title", "").strip().lower(), s.get("artist", "").strip().lower()) + if key in seen_keys: + return False + found_songs.append({ + "item_id": iid, + "title": s.get("title", ""), + "artist": s.get("artist", ""), + "album": s.get("album", ""), + }) + seen_ids.add(iid) + seen_keys.add(key) + return True + + def _add_batch(songs, channel): + added = 0 + for s in songs or []: + if len(found_songs) >= get_songs: + break + if _add_one(s): + added += 1 + if added: + log_messages.append(f" {channel}: +{added} (pool {len(found_songs)})") + return added + + def _energy_to_raw(value): + return config.ENERGY_MIN + float(value) * (config.ENERGY_MAX - config.ENERGY_MIN) + + def _year_gate(songs, year_min, year_max): + """Keep only songs whose release year is in range. Sound/artist channels + match on audio/similarity and cannot honor a release-year constraint, so a + request like '90s rap' would otherwise leak any-era songs. No-op when no + year is set.""" + if year_min is None and year_max is None: + return songs or [] + ids = [s.get("item_id") for s in (songs or []) if s.get("item_id")] + if not ids: + return [] + gated = _database_genre_query_sync( + get_songs=len(ids), year_min=year_min, year_max=year_max, candidate_item_ids=ids, + ) + return gated.get("songs") or [] + + def _run_filter(year_min, year_max, use_scored): + return _database_genre_query_sync( + genres=filt["genres"] or None, + get_songs=get_songs, + moods=(filt["moods"] or None) if use_scored else None, + tempo_min=filt["tempo_min"] if use_scored else None, + tempo_max=filt["tempo_max"] if use_scored else None, + energy_min=_energy_to_raw(filt["energy_min"]) if (use_scored and filt["energy_min"] is not None) else None, + energy_max=_energy_to_raw(filt["energy_max"]) if (use_scored and filt["energy_max"] is not None) else None, + year_min=year_min, + year_max=year_max, + voices=filt["voices"] or None, + score_threshold=config.AI_BRAINSTORM_GENRE_SCORE_THRESHOLD, + ).get("songs") or [] + + has_filter = bool( + filt["genres"] or filt["moods"] or filt["voices"] + or filt["year_min"] is not None or filt["year_max"] is not None + or filt["energy_min"] is not None or filt["energy_max"] is not None + or filt["tempo_min"] is not None or filt["tempo_max"] is not None + ) + relax_anchor = bool(filt["genres"] or filt["year_min"] is not None or filt["year_max"] is not None) + ymin, ymax = filt["year_min"], filt["year_max"] - stage2_items = [] - with db_conn.cursor(cursor_factory=DictCursor) as cur: - if valid_items: - values_params = [] - for title, artist in valid_items: - values_params.extend([title.lower(), artist.lower()]) - values_clause = ', '.join(['(%s, %s)'] * len(valid_items)) - cur.execute(f""" - SELECT item_id, title, author, album - FROM public.score - WHERE (LOWER(title), LOWER(author)) IN (VALUES {values_clause}) - """, values_params) - exact_rows = cur.fetchall() - - exact_match_map = {} - for row in exact_rows: - key = (row['title'].lower(), row['author'].lower()) - if key not in exact_match_map: - exact_match_map[key] = row - - stage2_items = [] - for title, artist in valid_items: - key = (title.lower(), artist.lower()) - result = exact_match_map.get(key) - if result and result['item_id'] not in seen_ids: - found_songs.append({ - "item_id": result['item_id'], - "title": result['title'], - "artist": result['author'], - "album": result.get('album', '') - }) - seen_ids.add(result['item_id']) - elif not result: - stage2_items.append((title, artist)) - - if stage2_items: - or_conditions = [] - fuzzy_params = [] - fuzzy_lookup_order = [] - for title, artist in stage2_items: - title_norm = _normalize(title) - artist_norm = _normalize(artist) - if title_norm and artist_norm: - or_conditions.append("""( - LOWER(REPLACE(REPLACE(REPLACE(REPLACE(REPLACE(title, ' ', ''), '-', ''), '''', ''), '.', ''), ',', '')) - LIKE LOWER(%s) - AND LOWER(REPLACE(REPLACE(REPLACE(REPLACE(REPLACE(author, ' ', ''), '-', ''), '''', ''), '.', ''), ',', '')) - LIKE LOWER(%s) - )""") - fuzzy_params.extend([f"%{_escape_like(title_norm)}%", f"%{_escape_like(artist_norm)}%"]) - fuzzy_lookup_order.append((title_norm, artist_norm)) - - if or_conditions: - where_clause = ' OR '.join(or_conditions) - cur.execute(f""" - SELECT item_id, title, author, album - FROM public.score - WHERE {where_clause} - ORDER BY LENGTH(title) + LENGTH(author) - """, fuzzy_params) - fuzzy_rows = cur.fetchall() - - for row in fuzzy_rows: - if row['item_id'] not in seen_ids: - db_title_norm = _normalize(row['title']) - db_artist_norm = _normalize(row['author']) - for t_norm, a_norm in fuzzy_lookup_order: - if t_norm in db_title_norm and a_norm in db_artist_norm: - found_songs.append({ - "item_id": row['item_id'], - "title": row['title'], - "artist": row['author'], - "album": row.get('album', '') - }) - seen_ids.add(row['item_id']) - fuzzy_lookup_order.remove((t_norm, a_norm)) - break - - found_songs = found_songs[:get_songs] - - log_messages.append(f"Found {len(found_songs)} songs in database (from {len(song_list)} AI suggestions)") - - return {"songs": found_songs, "ai_suggestions": len(song_list), "message": "\n".join(log_messages)} - finally: - db_conn.close() + try: + channels = [] + for i, desc in enumerate(recipe["sound_descriptions"]): + songs = _year_gate(_text_search_sync(desc, None, None, get_songs).get("songs"), ymin, ymax) + if songs: + channels.append((f"audio#{i + 1}", songs)) + for art in recipe["seed_artists"]: + raw = _artist_similarity_api_sync(art, config.AI_BRAINSTORM_SIMILAR_ARTISTS_PER_SEED, get_songs) + songs = _year_gate(raw.get("songs"), ymin, ymax) + if songs: + channels.append((f"artist:{art}", songs)) + for i, theme in enumerate(recipe["lyric_themes"]): + songs = _year_gate(_lyrics_search_sync(theme, get_songs).get("songs"), ymin, ymax) + if songs: + channels.append((f"lyrics#{i + 1}", songs)) + if has_filter: + fsongs = _run_filter(ymin, ymax, use_scored=True) + if fsongs: + channels.append(("filter", fsongs)) + + cursors = [0] * len(channels) + added_per = [0] * len(channels) + progressing = True + while len(found_songs) < get_songs and progressing: + progressing = False + for ci, (_label, songs) in enumerate(channels): + while cursors[ci] < len(songs): + s = songs[cursors[ci]] + cursors[ci] += 1 + if _add_one(s): + added_per[ci] += 1 + progressing = True + break + if len(found_songs) >= get_songs: + break + for ci, (label, _songs) in enumerate(channels): + if added_per[ci]: + log_messages.append(f" {label}: +{added_per[ci]}") + + floor = min(get_songs, config.AI_BRAINSTORM_POOL_FLOOR) + if len(found_songs) < floor and relax_anchor: + log_messages.append(f" pool under floor ({len(found_songs)} < {floor}); relaxing") + pad = config.AI_BRAINSTORM_RELAX_YEAR_PAD + rmin = (ymin - pad) if ymin is not None else None + rmax = (ymax + pad) if ymax is not None else None + _add_batch(_run_filter(rmin, rmax, use_scored=False), "relax:filter") + if len(found_songs) < floor and filt["genres"]: + relaxed = _text_search_sync(", ".join(filt["genres"]) + " music", None, None, get_songs).get("songs") + _add_batch(_year_gate(relaxed, rmin, rmax), "relax:audio") + except Exception: + logger.exception("Brainstorm channel execution failed") + + found_songs = found_songs[:get_songs] + log_messages.append(f"Brainstorm fused {len(found_songs)} library songs") + return {"songs": found_songs, "message": "\n".join(log_messages)} diff --git a/tasks/ai/tools.py b/tasks/ai/tools.py index c519e80a..2b0e706e 100644 --- a/tasks/ai/tools.py +++ b/tasks/ai/tools.py @@ -327,10 +327,11 @@ def get_mcp_tools() -> List[Dict]: { "name": "knowledge_lookup", "description": ( - "World-knowledge fallback. USE ONLY when the library can't surface the answer " - "via seed_search or search_database. Good for: 'Grammy winners 2020', " - "'#1 hits of 1985', 'songs sampled by Daft Punk', 'best festival anthems'. " - "Returns AI-suggested songs that are then matched against the library." + "Popularity / 'best of' / cultural requests that need world knowledge to " + "interpret: 'best rap of the 90s', '#1 hits of 1985', 'best festival anthems', " + "'Grammy winners 2020'. The model turns the request into a grounded search " + "recipe (genre/year/energy filters + 'how it sounds' descriptions + seed " + "artists) that is run against THIS library and fused. It never invents song titles." ), "inputSchema": { "type": "object", diff --git a/tasks/analysis.py b/tasks/analysis.py index 66f3e2a9..61429d02 100644 --- a/tasks/analysis.py +++ b/tasks/analysis.py @@ -11,12 +11,11 @@ import time import logging import uuid -import traceback import gc import platform import librosa -import onnxruntime as ort # re-exported: tests patch `tasks.analysis.ort.InferenceSession` +import onnxruntime as ort # noqa: F401 re-exported: tests patch `tasks.analysis.ort.InferenceSession` # RQ import from rq import get_current_job, Retry @@ -48,11 +47,12 @@ from flask_app import app from app_helper import ( redis_conn, rq_queue_default, get_db, save_task_status, - get_task_info_from_db, get_child_tasks_from_db, + get_task_info_from_db, build_and_store_map_projection, build_and_store_artist_projection, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED, ) +from database import get_child_tasks_from_db from error import error_manager from error.error_dictionary import ( @@ -69,7 +69,7 @@ # tests depend on (``run_inference``, ``_find_onnx_name``, ``sigmoid``). # Helpers consumed only inside this file go through ``_ah.`` instead. from . import analysis_helper as _ah -from .analysis_helper import ( +from .analysis_helper import ( # noqa: F401 DEFINED_TENSOR_NAMES, _find_onnx_name, # re-export: tests do `from tasks.analysis import _find_onnx_name` run_inference, # re-export: tests do `from tasks.analysis import run_inference` diff --git a/tasks/analysis_helper.py b/tasks/analysis_helper.py index 421f39c1..6453058c 100644 --- a/tasks/analysis_helper.py +++ b/tasks/analysis_helper.py @@ -11,12 +11,13 @@ from .memory_utils import cleanup_onnx_session, comprehensive_memory_cleanup -# `app_helper` and `app_helper_artist` are safe at module top: they have no -# import cycle back into this module. Optional ML modules -# (.clap_analyzer / lyrics.lyrics_transcriber) stay inline inside the -# per-feature helpers so workers without those models can still import this -# module. -from app_helper import ( +# `database` and `app_helper_artist` are safe at module top: they have no +# import cycle back into this module, and importing the DB primitives directly +# (rather than via the app_helper facade) keeps this helper decoupled from the +# blueprint layer. Optional ML modules (.clap_analyzer / lyrics.lyrics_transcriber) +# stay inline inside the per-feature helpers so workers without those models can +# still import this module. +from database import ( get_db, get_clap_embedding, save_track_analysis_and_embedding, @@ -565,7 +566,7 @@ def run_lyrics_for_track(item, path, track_audio, track_sr, track_name_full, if track_audio is None or track_audio.size == 0 or track_sr is None: raise RuntimeError("Failed to load audio for lyrics analysis") else: - def audio_loader(): + def audio_loader(): # noqa: F811 p = download_fn() if download_fn is not None else None if not p: raise RuntimeError("Failed to download audio for lyrics ASR") diff --git a/tasks/clap_analyzer.py b/tasks/clap_analyzer.py index fb5dbb18..fee9ac8c 100644 --- a/tasks/clap_analyzer.py +++ b/tasks/clap_analyzer.py @@ -733,7 +733,7 @@ def is_clap_available() -> bool: # --- CLAP-based Other Features (replaces mood-specific ONNX models) --- def get_or_cache_other_feature_text_embeddings(redis_conn) -> Optional[dict]: - """Return CLAP text embeddings for OTHER_FEATURE_LABELS, cached in Redis. + """Return CLAP text embeddings for config.OTHER_FEATURE_LABELS, cached in Redis. This function now includes the same optimised caching strategy used in devel‑DCLAP: @@ -780,7 +780,7 @@ def get_or_cache_other_feature_text_embeddings(redis_conn) -> Optional[dict]: logger.warning(f"Failed to read CLAP text embeddings from Redis: {e}") # Compute text embeddings for each label - logger.info(f"Computing CLAP text embeddings for OTHER_FEATURE_LABELS: {config.OTHER_FEATURE_LABELS}") + logger.info(f"Computing CLAP text embeddings for config.OTHER_FEATURE_LABELS: {config.OTHER_FEATURE_LABELS}") try: embeddings = get_text_embeddings_batch(config.OTHER_FEATURE_LABELS) if embeddings is None: diff --git a/tasks/cleaning.py b/tasks/cleaning.py index 0f61a365..a44f888d 100644 --- a/tasks/cleaning.py +++ b/tasks/cleaning.py @@ -3,7 +3,6 @@ import time import logging import uuid -import traceback from collections import defaultdict # RQ import @@ -29,7 +28,8 @@ def identify_and_clean_orphaned_albums_task(): This combines identification and deletion into a single automated process. """ from flask_app import app - from app_helper import redis_conn, get_db, save_task_status, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE + from app_helper import redis_conn, get_db, save_task_status + from config import TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE current_job = get_current_job(redis_conn) current_task_id = current_job.id if current_job else str(uuid.uuid4()) diff --git a/tasks/clustering.py b/tasks/clustering.py index f03395d6..ecbbf401 100644 --- a/tasks/clustering.py +++ b/tasks/clustering.py @@ -19,11 +19,31 @@ from psycopg2.extras import DictCursor # Import configuration -from config import MAX_SONGS_PER_CLUSTER, MOOD_LABELS, STRATIFIED_GENRES, MUTATION_KMEANS_COORD_FRACTION, MUTATION_INT_ABS_DELTA, MUTATION_FLOAT_ABS_DELTA, TOP_N_ELITES, EXPLOITATION_START_FRACTION, EXPLOITATION_PROBABILITY_CONFIG, SAMPLING_PERCENTAGE_CHANGE_PER_RUN, ITERATIONS_PER_BATCH_JOB, MAX_CONCURRENT_BATCH_JOBS, MIN_PLAYLIST_SIZE_FOR_TOP_N, CLUSTERING_BATCH_TIMEOUT_MINUTES, CLUSTERING_MAX_FAILED_BATCHES, CLUSTERING_CLEANING +from config import ( + MAX_SONGS_PER_CLUSTER, MOOD_LABELS, STRATIFIED_GENRES, + MUTATION_KMEANS_COORD_FRACTION, MUTATION_INT_ABS_DELTA, + MUTATION_FLOAT_ABS_DELTA, TOP_N_ELITES, EXPLOITATION_START_FRACTION, + EXPLOITATION_PROBABILITY_CONFIG, SAMPLING_PERCENTAGE_CHANGE_PER_RUN, + ITERATIONS_PER_BATCH_JOB, MAX_CONCURRENT_BATCH_JOBS, + MIN_PLAYLIST_SIZE_FOR_TOP_N, CLUSTERING_BATCH_TIMEOUT_MINUTES, + CLUSTERING_MAX_FAILED_BATCHES, CLUSTERING_CLEANING, + TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, + TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED, +) from error import error_manager from error.error_dictionary import ERR_CLUSTERING_FAILED +# App helper functions +from app_helper import ( + save_task_status, redis_conn, get_task_info_from_db, + get_db, rq_queue_default, +) +from database import update_playlist_table, get_child_tasks_from_db + +# JSON sanitization (numpy scalars/arrays -> native types) shared project-wide +from sanitization import sanitize_for_json + # Import AI naming function and prompt template # (used by clustering_helper._try_ai_name_playlist, imported there) # Import media server functions @@ -44,18 +64,11 @@ select_top_n_diverse_playlists ) -# we want to maintain np.float_ for backwards compatibility but it was removed in numpy 2.0 -# the check below in sanitize_for_json causes an AttributeError that crashes the clustering algo -# since it tries to access np.float_ so we monkeypatch np.float_ to point to np.float64 -if not np.__dict__.get('float_'): - np.float_ = np.float64 - logger = logging.getLogger(__name__) def batch_task_failure_handler(job, connection, type, value, tb): """A failure handler for the clustering batch sub-task, executed by the worker.""" from flask_app import app - from app_helper import save_task_status, TASK_STATUS_FAILURE with app.app_context(): task_id = getattr(job, 'id', None) or getattr(job, 'get_id', lambda: None)() parent_id = job.kwargs.get('parent_task_id') @@ -86,28 +99,6 @@ def batch_task_failure_handler(job, connection, type, value, tb): ) app.logger.error(f"Clustering batch task {task_id} (parent: {parent_id}) failed permanently. DB status updated.\n{tb_formatted}") -def _sanitize_for_json(obj): - """ - Recursively converts numpy arrays and numpy numeric types to native Python types - to ensure the object is JSON serializable. - """ - if isinstance(obj, dict): - return {k: _sanitize_for_json(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [_sanitize_for_json(elem) for elem in obj] - elif isinstance(obj, np.ndarray): - return obj.tolist() - # Handle numpy numeric types which are not JSON serializable by default - elif isinstance(obj, (np.int_, np.intc, np.intp, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64)): - return int(obj) - elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)): - return float(obj) - elif isinstance(obj, np.bool_): - return bool(obj) - else: - return obj - - def run_clustering_batch_task( batch_id_str, start_run_idx, num_iterations_in_batch, genre_to_lightweight_track_data_map_json, @@ -134,9 +125,6 @@ def run_clustering_batch_task( """ # --- Local imports to prevent circular dependency --- from flask_app import app - from app_helper import (redis_conn, save_task_status, get_task_info_from_db, - TASK_STATUS_PROGRESS, TASK_STATUS_REVOKED, TASK_STATUS_FAILURE, - TASK_STATUS_SUCCESS) current_job = get_current_job(redis_conn) current_task_id = current_job.id if current_job else str(uuid.uuid4()) @@ -227,7 +215,7 @@ def _log_and_update(message, progress, details=None, state=TASK_STATUS_PROGRESS) # *** FIX: Sanitize the result to make it JSON-serializable before logging/returning *** if best_result_in_batch: - best_result_in_batch = _sanitize_for_json(best_result_in_batch) + best_result_in_batch = sanitize_for_json(best_result_in_batch) final_details = { "best_score_in_batch": best_score_in_batch, @@ -276,7 +264,6 @@ def run_clustering_task( """ # --- Local imports to prevent circular dependency --- from flask_app import app - from app_helper import redis_conn, get_db, save_task_status, get_task_info_from_db, update_playlist_table, get_child_tasks_from_db, TASK_STATUS_STARTED, TASK_STATUS_PROGRESS, TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED current_job = get_current_job(redis_conn) current_task_id = current_job.id if current_job else str(uuid.uuid4()) @@ -668,8 +655,6 @@ def _monitor_and_process_batches(state_dict, parent_task_id, initial_check=False CRITICAL: This prevents the main task from hanging at 4980/5000 runs by implementing timeouts and forced progress tracking. """ - from app_helper import redis_conn, get_child_tasks_from_db, TASK_STATUS_SUCCESS, TASK_STATUS_FAILURE, TASK_STATUS_REVOKED, TASK_STATUS_STARTED - current_time = time.time() timeout_seconds = CLUSTERING_BATCH_TIMEOUT_MINUTES * 60 processed_jobs = state_dict.get("processed_job_ids", set()) @@ -833,7 +818,6 @@ def _monitor_and_process_batches(state_dict, parent_task_id, initial_check=False def _launch_batch_job(state_dict, parent_task_id, batch_idx, total_runs, genre_map, target_per_genre, *args): """Constructs and enqueues a single batch job.""" - from app_helper import rq_queue_default # Local import to avoid circular dependency issues at top-level # Unpack all the parameters passed via *args ( diff --git a/tasks/clustering_gpu.py b/tasks/clustering_gpu.py index 26e1ac79..4e7ba26a 100644 --- a/tasks/clustering_gpu.py +++ b/tasks/clustering_gpu.py @@ -64,7 +64,7 @@ def check_gpu_available(): try: import cupy as cp - import cuml + import cuml # noqa: F401 # Try to create a small array on GPU to verify it works test_array = cp.array([1, 2, 3]) _ = test_array.sum() diff --git a/tasks/clustering_helper.py b/tasks/clustering_helper.py index c14598c8..4a38a84e 100644 --- a/tasks/clustering_helper.py +++ b/tasks/clustering_helper.py @@ -26,7 +26,7 @@ logger.debug("GPU clustering module not available, using CPU only") # RQ imports for safe result fetching -from rq.job import Job +from rq.job import Job, JobStatus from rq.exceptions import NoSuchJobError from config import (STRATIFIED_GENRES, OTHER_FEATURE_LABELS, MOOD_LABELS, MAX_DISTANCE, @@ -36,13 +36,18 @@ LN_MOOD_PURITY_EMBEDING_STATS, LN_OTHER_FEATURES_DIVERSITY_STATS, LN_OTHER_FEATURES_PURITY_STATS, OTHER_FEATURE_PREDOMINANCE_THRESHOLD_FOR_PURITY, - USE_GPU_CLUSTERING) + USE_GPU_CLUSTERING, TASK_STATUS_SUCCESS) from .commons import score_vector # Import AI naming for playlist helpers from tasks.ai.api import get_ai_playlist_name from tasks.ai.prompts import creative_prompt_template +# Low-level DB / queue primitives, imported directly rather than via the +# app_helper facade (keeps this helper decoupled from the blueprint layer). +from database import get_tracks_by_ids, get_score_data_by_ids, get_task_info_from_db +from taskqueue import redis_conn + # --- Playlist Naming & Shuffling Helpers --- @@ -175,8 +180,6 @@ def _perform_single_clustering_iteration( def _prepare_iteration_data(item_ids, active_mood_labels, use_embeddings, log_prefix, run_idx): """Fetches track data, creates feature/embedding vectors, and ensures alignment.""" - # Local imports to prevent circular dependency - from app_helper import get_tracks_by_ids, get_score_data_by_ids logger.info(f"{log_prefix} Iteration {run_idx}: Fetching data for {len(item_ids)} tracks. Use embeddings: {use_embeddings}") rows = get_tracks_by_ids(item_ids) if use_embeddings else get_score_data_by_ids(item_ids) # These functions are now imported locally @@ -718,8 +721,6 @@ def get_job_result_safely(job_id, parent_task_id, task_type="child task"): """ # Local imports to prevent circular dependency from flask_app import app - from rq.job import JobStatus - from app_helper import redis_conn, get_task_info_from_db, TASK_STATUS_SUCCESS try: job = Job.fetch(job_id, connection=redis_conn) diff --git a/tasks/memory_utils.py b/tasks/memory_utils.py index 06cb3aad..c10d2f20 100644 --- a/tasks/memory_utils.py +++ b/tasks/memory_utils.py @@ -1,12 +1,13 @@ """ -Memory management and data sanitization utilities for AudioMuse AI. +Memory management utilities for AudioMuse AI. -This module provides utilities to address two critical issues: -1. PostgreSQL NUL byte errors from corrupted metadata -2. ONNX Runtime GPU memory allocation failures from fragmentation +Addresses ONNX Runtime GPU memory allocation failures from fragmentation: +force CUDA cache clearing, explicit session disposal, allocation-error handling +with retry, and periodic session recycling. + +String/JSON sanitization helpers now live in the top-level ``sanitization`` module. Key functions: -- sanitize_string_for_db: Remove NULL bytes and control characters before DB writes - cleanup_cuda_memory: Force CUDA cache clearing and garbage collection - cleanup_onnx_session: Explicit session disposal with immediate GC - handle_onnx_memory_error: Detect allocation errors, trigger cleanup, enable retry @@ -15,71 +16,11 @@ import gc import logging -import re from typing import Optional, Dict, Any, Callable logger = logging.getLogger(__name__) -def sanitize_string_for_db(value: Optional[str]) -> Optional[str]: - """ - Remove NULL bytes (0x00) and control characters from strings before database writes. - - PostgreSQL TEXT/VARCHAR columns reject strings containing NULL bytes (0x00), which - can appear in corrupted metadata from music files. This function sanitizes strings - to prevent database insertion errors. - - Args: - value: String to sanitize (can be None) - - Returns: - Sanitized string with NULL bytes and control characters removed, or None if input is None - - Examples: - >>> sanitize_string_for_db("Tyler, The Creator\x00YoungBoy") - "Tyler, The CreatorYoungBoy" - >>> sanitize_string_for_db(None) - None - >>> sanitize_string_for_db("") - "" - """ - if value is None: - return None - - if not isinstance(value, str): - # Convert to string if not already - value = str(value) - - # Remove NULL bytes (0x00) - value = value.replace('\x00', '') - - # Remove other control characters (0x01-0x1F except newline, tab, carriage return) - # Keep: \t (0x09), \n (0x0A), \r (0x0D) - # Remove: 0x01-0x08, 0x0B-0x0C, 0x0E-0x1F - value = re.sub(r'[\x01-\x08\x0B-\x0C\x0E-\x1F]', '', value) - - return value - - -def sanitize_json_for_db(value): - """Recursively sanitize strings inside a JSON-serializable value. - - Applies :func:`sanitize_string_for_db` to every string found inside - dicts, lists and tuples. Non-string scalars are returned as-is. Use this - before ``json.dumps(...)`` when the result is going into a Postgres - jsonb column. - """ - if isinstance(value, str): - return sanitize_string_for_db(value) - if isinstance(value, dict): - return {k: sanitize_json_for_db(v) for k, v in value.items()} - if isinstance(value, list): - return [sanitize_json_for_db(v) for v in value] - if isinstance(value, tuple): - return tuple(sanitize_json_for_db(v) for v in value) - return value - - def cleanup_cuda_memory(force: bool = False) -> bool: """ Force CUDA cache clearing and garbage collection to free GPU memory. diff --git a/tasks/provider_migration_tasks.py b/tasks/provider_migration_tasks.py index 6f51c37c..b0160017 100644 --- a/tasks/provider_migration_tasks.py +++ b/tasks/provider_migration_tasks.py @@ -36,7 +36,7 @@ import re import time -from tasks.memory_utils import sanitize_string_for_db as _sanitize_text +from sanitization import sanitize_string_for_db as _sanitize_text logger = logging.getLogger(__name__) diff --git a/tasks/radio_manager.py b/tasks/radio_manager.py index 170a7b11..0493a940 100644 --- a/tasks/radio_manager.py +++ b/tasks/radio_manager.py @@ -16,7 +16,7 @@ def run_radio_playlists(): Falls back to create_playlist for unsupported backends. """ - from app_helper import get_alchemy_radios + from database import get_alchemy_radios radios = [r for r in get_alchemy_radios() if r.get('enabled')] logger.info(f"Radio playlist run started for {len(radios)} enabled radios.") diff --git a/tasks/sem_grove_manager.py b/tasks/sem_grove_manager.py index a0b007bf..579981d6 100644 --- a/tasks/sem_grove_manager.py +++ b/tasks/sem_grove_manager.py @@ -338,7 +338,7 @@ def _merge_batches(): def _load_sem_grove_index_from_db() -> bool: """Load the SemGrove merged Voyager index from the DB into the global cache.""" try: - import voyager # type: ignore + import voyager # type: ignore # noqa: F401 except ImportError: logger.warning("Voyager unavailable; cannot load SemGrove index.") return False diff --git a/tasks/song_alchemy.py b/tasks/song_alchemy.py index a8b985ce..f3286864 100644 --- a/tasks/song_alchemy.py +++ b/tasks/song_alchemy.py @@ -7,7 +7,7 @@ _project_to_2d, _project_with_discriminant, ) -from app_helper import get_score_data_by_ids, load_map_projection +from database import get_score_data_by_ids, load_map_projection import config logger = logging.getLogger(__name__) @@ -124,7 +124,7 @@ def _compute_centroid_from_items(items: List[dict]) -> np.ndarray: weights.append(weight) elif item_type == 'anchor': - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id anchor = get_alchemy_anchor_by_id(item_id) if anchor and anchor.get('centroid') and isinstance(anchor.get('centroid'), list): vectors.append(np.array(anchor['centroid'], dtype=float)) @@ -157,6 +157,8 @@ def song_alchemy(add_items=None, subtract_items=None, add_ids=None, subtract_ids Returns list of song detail dicts (using get_score_data_by_ids mapping) """ + from app_helper_artist import get_artist_name_by_id + if n_results is None: n_results = config.ALCHEMY_DEFAULT_N_RESULTS n_results = min(n_results, config.ALCHEMY_MAX_N_RESULTS) @@ -293,7 +295,7 @@ def song_alchemy(add_items=None, subtract_items=None, add_ids=None, subtract_ids # Add anchors as individual points add_anchor_items = [item for item in add_items if item.get('type') == 'anchor'] if add_anchor_items: - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id for item in add_anchor_items: anchor_id = item['id'] anchor = get_alchemy_anchor_by_id(anchor_id) @@ -322,7 +324,6 @@ def song_alchemy(add_items=None, subtract_items=None, add_ids=None, subtract_ids logger.info(f"Retrieved {len(gmm_vecs)} GMM components for artist {artist_id}") for comp_idx, (vec, weight) in enumerate(zip(gmm_vecs, gmm_weights)): # Store metadata for artist component - from app_helper_artist import get_artist_name_by_id artist_name = artist_id resolved = get_artist_name_by_id(artist_id) if resolved: @@ -355,7 +356,7 @@ def song_alchemy(add_items=None, subtract_items=None, add_ids=None, subtract_ids # Add anchors as individual points subtract_anchor_items = [item for item in subtract_items if item.get('type') == 'anchor'] if subtract_anchor_items: - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id for item in subtract_anchor_items: anchor_id = item['id'] anchor = get_alchemy_anchor_by_id(anchor_id) @@ -384,7 +385,6 @@ def song_alchemy(add_items=None, subtract_items=None, add_ids=None, subtract_ids logger.info(f"Retrieved {len(gmm_vecs)} GMM components for artist {artist_id}") for comp_idx, (vec, weight) in enumerate(zip(gmm_vecs, gmm_weights)): # Store metadata for artist component - from app_helper_artist import get_artist_name_by_id artist_name = artist_id resolved = get_artist_name_by_id(artist_id) if resolved: @@ -444,7 +444,7 @@ def song_alchemy(add_items=None, subtract_items=None, add_ids=None, subtract_ids # Load precomputed artist component projections artist_comp_to_coord = {} try: - from app_helper import ARTIST_PROJECTION_CACHE + from database import ARTIST_PROJECTION_CACHE if ARTIST_PROJECTION_CACHE: component_map = ARTIST_PROJECTION_CACHE.get('component_map', []) projection = ARTIST_PROJECTION_CACHE.get('projection') @@ -600,7 +600,7 @@ def _centroid_from_member_coords(items, is_add=True): vec = get_vector_by_id(item_id) elif isinstance(pid, str) and pid.startswith('__add_anchor__'): anchor_id = pid.replace('__add_anchor__', '') - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id anchor = get_alchemy_anchor_by_id(anchor_id) if anchor and anchor.get('centroid') and isinstance(anchor['centroid'], list): vec = np.array(anchor['centroid'], dtype=float) @@ -608,7 +608,7 @@ def _centroid_from_member_coords(items, is_add=True): vec = None elif isinstance(pid, str) and pid.startswith('__sub_anchor__'): anchor_id = pid.replace('__sub_anchor__', '') - from app_helper import get_alchemy_anchor_by_id + from database import get_alchemy_anchor_by_id anchor = get_alchemy_anchor_by_id(anchor_id) if anchor and anchor.get('centroid') and isinstance(anchor['centroid'], list): vec = np.array(anchor['centroid'], dtype=float) diff --git a/tasks/voyager_manager.py b/tasks/voyager_manager.py index 17420d1c..3a3718a4 100644 --- a/tasks/voyager_manager.py +++ b/tasks/voyager_manager.py @@ -1037,7 +1037,7 @@ def find_nearest_neighbors_by_vector(query_vector: np.ndarray, n: int = 100, eli if voyager_index is None or id_map is None: raise RuntimeError("Voyager index is not loaded in memory.") - from app_helper import get_db, get_score_data_by_ids + from app_helper import get_db db_conn = get_db() # If caller didn't supply eliminate_duplicates explicitly (None), use configured default diff --git a/test/integration/test.py b/test/integration/test.py index f8b6999f..5301a0c3 100644 --- a/test/integration/test.py +++ b/test/integration/test.py @@ -255,8 +255,6 @@ def test_map_visualization(): elapsed = time.time() - start_time print(f"[TIMING] Map Visualization test completed in {elapsed:.2f} seconds") -import pytest - #pytest -v -s test.py -k test_annoy_similarity_and_playlist def test_annoy_similarity_and_playlist(): start_time = time.time() diff --git a/test/integration/test_analysis_integration.py b/test/integration/test_analysis_integration.py index 2ea829b8..b85beaa0 100644 --- a/test/integration/test_analysis_integration.py +++ b/test/integration/test_analysis_integration.py @@ -15,7 +15,6 @@ import sys import types from pathlib import Path -import importlib import json import pytest diff --git a/test/integration/test_brainstorm_integration.py b/test/integration/test_brainstorm_integration.py new file mode 100644 index 00000000..220485dd --- /dev/null +++ b/test/integration/test_brainstorm_integration.py @@ -0,0 +1,143 @@ +"""Real-Postgres integration test for the grounded AI brainstorm (#643). + +Drives ``_ai_brainstorm_sync`` against a live ``score`` table. Only the LLM +transport is stubbed (it returns a fixed recipe); the retrieval + fusion runs for +real, so this proves the recipe -> SQL filter channel -> fused pool path surfaces +the correct real rows and excludes the non-matching ones (wrong genre, wrong year, +or genre confidence below threshold). + +Database selection mirrors test_app_endpoints_integration.py: + * AUDIOMUSE_TEST_DATABASE_URL -- a throwaway DB the test fully owns, or + * an ephemeral instance via the optional ``pgserver`` package, or + * the module is skipped. + +Run locally: + pip install pgserver + pytest test/integration/test_brainstorm_integration.py -m integration -s -v --tb=short +""" +import importlib.util +import json +import os +import sys +import tempfile +import types +from unittest.mock import Mock, patch + +import pytest + +_REPO_ROOT = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +try: + import psycopg2 +except Exception: # pragma: no cover - psycopg2 is in test/requirements.txt + psycopg2 = None + + +_SCORE_DDL = ( + "CREATE TABLE score (item_id TEXT PRIMARY KEY, title TEXT, author TEXT, " + "album TEXT, album_artist TEXT, tempo REAL, key TEXT, scale TEXT, " + "mood_vector TEXT, energy REAL, other_features TEXT, year INTEGER, " + "rating INTEGER, file_path TEXT)" +) + +# (item_id, title, author, mood_vector, year). Only r1/r2 satisfy genre rock>=0.3 +# AND year in 1990..1999. +_SEED_ROWS = [ + ('r1', 'Rock One', 'Band A', 'rock:0.82,pop:0.20', 1995), + ('r2', 'Rock Two', 'Band B', 'rock:0.55,indie:0.30', 1992), + ('r3', 'Pop One', 'Band C', 'pop:0.90', 1995), + ('r4', 'Rock Late', 'Band D', 'rock:0.70', 2010), + ('r5', 'Rock Faint', 'Band E', 'rock:0.10,pop:0.60', 1994), +] + + +def _import_module(mod_name, relative_path): + """Load a module by file path, bypassing tasks/__init__.py (mirrors unit conftest).""" + mod_path = os.path.normpath(os.path.join(_REPO_ROOT, relative_path)) + if mod_name not in sys.modules: + spec = importlib.util.spec_from_file_location(mod_name, mod_path) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + return sys.modules[mod_name] + + +@pytest.fixture(scope='session') +def pg_dsn(): + if psycopg2 is None: + pytest.skip("psycopg2 not importable") + dsn = os.environ.get('AUDIOMUSE_TEST_DATABASE_URL') + if dsn: + try: + psycopg2.connect(dsn).close() + except Exception as e: + pytest.skip(f"AUDIOMUSE_TEST_DATABASE_URL not reachable: {e}") + yield dsn + return + try: + import pgserver + except Exception: + pytest.skip( + "No test database. Set AUDIOMUSE_TEST_DATABASE_URL to a disposable " + "DB, or `pip install pgserver` for an ephemeral local instance." + ) + data_dir = tempfile.mkdtemp(prefix='audiomuse_pg_') + server = pgserver.get_server(data_dir) + try: + yield server.get_uri() + finally: + server.cleanup() + + +@pytest.fixture +def brainstorm_db(pg_dsn): + conn = psycopg2.connect(pg_dsn) + conn.autocommit = True + with conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS score CASCADE") + cur.execute(_SCORE_DDL) + for item_id, title, author, mood_vector, year in _SEED_ROWS: + cur.execute( + "INSERT INTO score (item_id, title, author, mood_vector, year) " + "VALUES (%s, %s, %s, %s, %s)", + (item_id, title, author, mood_vector, year), + ) + conn.close() + yield pg_dsn + + +def _fake_ai_module(recipe_obj): + mod = types.ModuleType('tasks.ai.api') + mod.generate_text = Mock(return_value=json.dumps(recipe_obj)) + return mod + + +@pytest.mark.integration +class TestBrainstormGroundedRetrievalRealDb: + def test_filter_channel_surfaces_only_matching_rows(self, brainstorm_db, monkeypatch): + _import_module('tasks.mcp_helper', 'tasks/mcp_helper.py') + _import_module('tasks.ai.prompts', 'tasks/ai/prompts.py') + tool_impl = _import_module('tasks.ai.tool_impl', 'tasks/ai/tool_impl.py') + import config as cfg + + recipe = { + "filters": {"genres": ["rock"], "year_min": 1990, "year_max": 1999}, + "sound_descriptions": [], + "seed_artists": [], + "lyric_themes": [], + } + + # Floor of 1 keeps the relax pass from firing so the assertion isolates the + # primary filter channel against the real rows. + monkeypatch.setattr(cfg, 'AI_BRAINSTORM_POOL_FLOOR', 1) + monkeypatch.setattr(tool_impl, 'get_db_connection', lambda: psycopg2.connect(brainstorm_db)) + + with patch.dict(sys.modules, {'tasks.ai.api': _fake_ai_module(recipe)}): + result = tool_impl._ai_brainstorm_sync("best rock of the 90s", {"provider": "OLLAMA"}, 50) + + ids = {s["item_id"] for s in result["songs"]} + assert ids == {"r1", "r2"} + for song in result["songs"]: + assert song["item_id"] and song["title"] and song["artist"] diff --git a/test/integration/test_gpu_status.py b/test/integration/test_gpu_status.py index f32ecd3d..6da13c77 100644 --- a/test/integration/test_gpu_status.py +++ b/test/integration/test_gpu_status.py @@ -115,7 +115,7 @@ def test_cuml_clustering(): # Test cuML KMeans cuml_ok = False try: - from cuml.cluster import KMeans as cuKMeans + from cuml.cluster import KMeans as cuKMeans # noqa: F401 print_result("cuML KMeans", True, "Import successful") cuml_ok = True except Exception as e: @@ -123,14 +123,14 @@ def test_cuml_clustering(): # Test cuML DBSCAN try: - from cuml.cluster import DBSCAN as cuDBSCAN + from cuml.cluster import DBSCAN as cuDBSCAN # noqa: F401 print_result("cuML DBSCAN", True, "Import successful") except Exception as e: print_result("cuML DBSCAN", False, str(e)) # Test cuML PCA try: - from cuml.decomposition import PCA as cuPCA + from cuml.decomposition import PCA as cuPCA # noqa: F401 print_result("cuML PCA", True, "Import successful") except Exception as e: print_result("cuML PCA", False, str(e)) diff --git a/test/unit/conftest.py b/test/unit/conftest.py index 08762e32..cc27c7ee 100644 --- a/test/unit/conftest.py +++ b/test/unit/conftest.py @@ -114,3 +114,42 @@ def config_restore(): yield for attr, val in saved.items(): setattr(cfg, attr, val) + + +# --------------------------------------------------------------------------- +# Import-architecture report (terminal summary) +# --------------------------------------------------------------------------- + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + """Print the import-architecture report (layer table, max-chain confirmation, + and the recap of chains at the ceiling) whenever the architecture gate ran. + + Uses the terminal reporter so the report shows on every run -- pass or fail -- + without needing ``-s``. A PR that deepens the eager import graph will see the + new chains listed here (and the depth test will fail with the same recap). + """ + # The gate may be imported bare ("test_import_architecture") or + # package-qualified ("test.unit.test_import_architecture") depending on the + # presence of __init__.py files, so match by suffix. + mod = next( + (m for name, m in list(sys.modules.items()) + if name == "test_import_architecture" or name.endswith(".test_import_architecture")), + None, + ) + if mod is None: + return + ran = any( + "test_import_architecture" in getattr(rep, "nodeid", "") + for key in ("passed", "failed", "error") + for rep in terminalreporter.stats.get(key, []) + ) + if not ran: + return + try: + lines = mod.architecture_report() + except Exception as exc: # never let the report break the run + terminalreporter.write_line(f"[architecture] report unavailable: {exc}") + return + terminalreporter.section("Import-architecture report", "=") + for line in lines: + terminalreporter.write_line(line) diff --git a/test/unit/test_ai.py b/test/unit/test_ai.py index 85167863..6e9546d3 100644 --- a/test/unit/test_ai.py +++ b/test/unit/test_ai.py @@ -81,7 +81,6 @@ def _ensure_google_genai_stub(): google_mod.__path__ = [] sys.modules['google'] = google_mod from unittest.mock import MagicMock as _MM - from unittest.mock import MagicMock as _MM genai_mod = types.ModuleType('google.genai') genai_mod.Client = _MM genai_types = types.ModuleType('google.genai.types') @@ -116,7 +115,6 @@ def _ensure_mistralai_stub(): for _name, _relpath in ( ('tasks.ai.prompts', 'tasks/ai/prompts.py'), ('tasks.ai.providers.openai', 'tasks/ai/providers/openai.py'), - ('tasks.ai.providers.ollama', 'tasks/ai/providers/ollama.py'), ('tasks.ai.providers.gemini', 'tasks/ai/providers/gemini.py'), ('tasks.ai.providers.mistral', 'tasks/ai/providers/mistral.py'), ('tasks.ai.api', 'tasks/ai/api.py'), @@ -124,12 +122,11 @@ def _ensure_mistralai_stub(): _load_submodule(_name, _relpath) -from unittest.mock import Mock, patch, call +from unittest.mock import Mock, patch import requests import json from tasks.ai.api import clean_playlist_name, get_ai_playlist_name from tasks.ai.providers.openai import generate_text as get_openai_compatible_playlist_name -from tasks.ai.providers.ollama import generate_text as get_ollama_playlist_name from tasks.ai.providers.gemini import generate_text as get_gemini_playlist_name from tasks.ai.providers.mistral import generate_text as get_mistral_playlist_name from tasks.ai.prompts import creative_prompt_template @@ -833,28 +830,26 @@ def test_ultra_minimal_fallback_requires_proper_error_code(self, mock_post, mock class TestGetOllamaPlaylistName: - """Tests for Ollama-specific wrapper function""" + """Tests for Ollama-format generate_text (now handled directly by openai.py)""" - @patch('tasks.ai.providers.openai.generate_text') - def test_calls_openai_compatible_with_correct_params(self, mock_func): - """Test that Ollama wrapper calls underlying function correctly""" - mock_func.return_value = "Test Playlist" + @patch('tasks.ai.providers.openai.requests.post') + def test_calls_with_ollama_format_url(self, mock_post): + """Test that generate_text handles Ollama /api/generate URLs correctly""" + # Simulate a successful streaming response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.iter_lines.return_value = [ + b'{"response":"Test Playlist","done":true}' + ] + mock_post.return_value = mock_response - result = get_ollama_playlist_name( - ollama_url="http://localhost:11434/api/generate", + result = get_openai_compatible_playlist_name( + server_url="http://localhost:11434/api/generate", model_name="deepseek-r1:1.5b", - full_prompt="test prompt" - ) - - mock_func.assert_called_once_with( - "http://localhost:11434/api/generate", - "deepseek-r1:1.5b", - "test prompt", + full_prompt="test prompt", api_key="no-key-needed", - skip_delay=False, - temperature=None, - max_tokens=None ) + assert result == "Test Playlist" diff --git a/test/unit/test_app_alchemy_anchor.py b/test/unit/test_app_alchemy_anchor.py index 10b5c1e3..fcd75775 100644 --- a/test/unit/test_app_alchemy_anchor.py +++ b/test/unit/test_app_alchemy_anchor.py @@ -19,21 +19,21 @@ def client(app): class TestCreateAnchorValidation: - @patch('app_helper.save_alchemy_anchor') + @patch('database.save_alchemy_anchor') def test_whitespace_only_name_returns_400(self, mock_save, client): response = client.post('/api/anchors', json={'name': ' ', 'centroid': [0.1, 0.2]}) assert response.status_code == 400 assert response.get_json() == {'error': 'Anchor name is required'} mock_save.assert_not_called() - @patch('app_helper.save_alchemy_anchor') + @patch('database.save_alchemy_anchor') def test_non_list_centroid_returns_400(self, mock_save, client): response = client.post('/api/anchors', json={'name': 'My Anchor', 'centroid': 'not-a-list'}) assert response.status_code == 400 assert response.get_json() == {'error': 'Anchor centroid is required and must be a list'} mock_save.assert_not_called() - @patch('app_helper.save_alchemy_anchor') + @patch('database.save_alchemy_anchor') def test_empty_list_centroid_returns_400(self, mock_save, client): response = client.post('/api/anchors', json={'name': 'My Anchor', 'centroid': []}) assert response.status_code == 400 @@ -42,7 +42,7 @@ def test_empty_list_centroid_returns_400(self, mock_save, client): class TestRenameAnchorValidation: - @patch('app_helper.update_alchemy_anchor_name') + @patch('database.update_alchemy_anchor_name') def test_whitespace_only_name_returns_400(self, mock_update, client): response = client.put('/api/anchors/7', json={'name': ' '}) assert response.status_code == 400 diff --git a/test/unit/test_app_alchemy_radio.py b/test/unit/test_app_alchemy_radio.py index 2d478df3..0fd14d89 100644 --- a/test/unit/test_app_alchemy_radio.py +++ b/test/unit/test_app_alchemy_radio.py @@ -20,7 +20,7 @@ def client(app): class TestListRadios: - @patch('app_helper.get_alchemy_radios') + @patch('database.get_alchemy_radios') def test_returns_radio_list(self, mock_get, client): mock_get.return_value = [ {'id': 1, 'anchor_id': 5, 'name': 'Chill', 'temperature': 1.0, 'n_results': 100, 'enabled': True}, @@ -36,42 +36,42 @@ def test_returns_radio_list(self, mock_get, client): class TestCreateRadioValidation: - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_missing_anchor_id_returns_400(self, mock_create, client): response = client.post('/api/radios', json={'temperature': 1.0, 'n_results': 100}) assert response.status_code == 400 assert response.get_json() == {'error': 'Radio anchor is required'} mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_missing_temperature_returns_400(self, mock_create, client): response = client.post('/api/radios', json={'anchor_id': 5, 'n_results': 100}) assert response.status_code == 400 assert response.get_json() == {'error': 'Radio temperature is required'} mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_missing_n_results_returns_400(self, mock_create, client): response = client.post('/api/radios', json={'anchor_id': 5, 'temperature': 1.0}) assert response.status_code == 400 assert response.get_json() == {'error': 'Radio number of results is required'} mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_non_numeric_temperature_returns_400(self, mock_create, client): response = client.post('/api/radios', json={'anchor_id': 5, 'temperature': 'hot', 'n_results': 100}) assert response.status_code == 400 assert response.get_json() == {'error': 'Radio temperature must be a number'} mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_negative_temperature_returns_400(self, mock_create, client): response = client.post('/api/radios', json={'anchor_id': 5, 'temperature': -1, 'n_results': 100}) assert response.status_code == 400 assert response.get_json() == {'error': 'Radio temperature must be 0 or greater'} mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_non_finite_temperature_returns_400(self, mock_create, client): for bad_value in ('NaN', 'Infinity', 'inf'): response = client.post('/api/radios', @@ -81,14 +81,14 @@ def test_non_finite_temperature_returns_400(self, mock_create, client): assert response.get_json() == {'error': 'Radio temperature must be a finite number'} mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_n_results_out_of_range_returns_400(self, mock_create, client): for bad_value in (0, config.ALCHEMY_MAX_N_RESULTS + 1): response = client.post('/api/radios', json={'anchor_id': 5, 'temperature': 1.0, 'n_results': bad_value}) assert response.status_code == 400 mock_create.assert_not_called() - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_valid_payload_creates_radio(self, mock_create, client): mock_create.return_value = {'id': 1, 'anchor_id': 5, 'temperature': 1.0, 'n_results': 100, 'enabled': True} response = client.post('/api/radios', json={'anchor_id': 5, 'temperature': 1.0, 'n_results': 100}) @@ -96,7 +96,7 @@ def test_valid_payload_creates_radio(self, mock_create, client): assert response.get_json() == {'radio': {'id': 1, 'anchor_id': 5, 'temperature': 1.0, 'n_results': 100, 'enabled': True}} mock_create.assert_called_once_with(5, 1.0, 100, True) - @patch('app_helper.create_alchemy_radio') + @patch('database.create_alchemy_radio') def test_duplicate_anchor_returns_400(self, mock_create, client): mock_create.return_value = None response = client.post('/api/radios', json={'anchor_id': 5, 'temperature': 1.0, 'n_results': 100}) @@ -105,19 +105,19 @@ def test_duplicate_anchor_returns_400(self, mock_create, client): class TestUpdateRadio: - @patch('app_helper.update_alchemy_radio') + @patch('database.update_alchemy_radio') def test_missing_settings_returns_400(self, mock_update, client): response = client.put('/api/radios/3', json={'enabled': False}) assert response.status_code == 400 mock_update.assert_not_called() - @patch('app_helper.update_alchemy_radio') + @patch('database.update_alchemy_radio') def test_unknown_radio_returns_404(self, mock_update, client): mock_update.return_value = None response = client.put('/api/radios/3', json={'temperature': 0.5, 'n_results': 20, 'enabled': False}) assert response.status_code == 404 - @patch('app_helper.update_alchemy_radio') + @patch('database.update_alchemy_radio') def test_valid_update_saves_disabled_state(self, mock_update, client): mock_update.return_value = {'id': 3, 'anchor_id': 5, 'temperature': 0.5, 'n_results': 20, 'enabled': False} response = client.put('/api/radios/3', json={'temperature': 0.5, 'n_results': 20, 'enabled': False}) @@ -126,13 +126,13 @@ def test_valid_update_saves_disabled_state(self, mock_update, client): class TestDeleteRadio: - @patch('app_helper.delete_alchemy_radio') + @patch('database.delete_alchemy_radio') def test_unknown_radio_returns_404(self, mock_delete, client): mock_delete.return_value = False response = client.delete('/api/radios/9') assert response.status_code == 404 - @patch('app_helper.delete_alchemy_radio') + @patch('database.delete_alchemy_radio') def test_delete_returns_ok(self, mock_delete, client): mock_delete.return_value = True response = client.delete('/api/radios/9') @@ -165,7 +165,7 @@ def _radio(self, radio_id, anchor_id, name, temperature=1.0, n_results=100, enab return {'id': radio_id, 'anchor_id': anchor_id, 'name': name, 'temperature': temperature, 'n_results': n_results, 'enabled': enabled} - @patch('app_helper.get_alchemy_radios') + @patch('database.get_alchemy_radios') @patch('tasks.radio_manager.create_or_replace_playlist') @patch('tasks.radio_manager.song_alchemy') def test_creates_playlists_for_enabled_radios_only(self, mock_alchemy, mock_upsert, @@ -186,7 +186,7 @@ def test_creates_playlists_for_enabled_radios_only(self, mock_alchemy, mock_upse assert summary['playlists_created'] == 1 assert summary['radios_enabled'] == 1 - @patch('app_helper.get_alchemy_radios') + @patch('database.get_alchemy_radios') @patch('tasks.radio_manager.create_or_replace_playlist') @patch('tasks.radio_manager.song_alchemy') def test_upserts_after_generation(self, mock_alchemy, mock_upsert, @@ -204,7 +204,7 @@ def test_upserts_after_generation(self, mock_alchemy, mock_upsert, names = [c[0] for c in order_tracker.mock_calls] assert names.index('alchemy') < names.index('upsert') - @patch('app_helper.get_alchemy_radios') + @patch('database.get_alchemy_radios') @patch('tasks.radio_manager.create_or_replace_playlist') @patch('tasks.radio_manager.song_alchemy') def test_one_failing_radio_does_not_block_others(self, mock_alchemy, mock_upsert, @@ -223,7 +223,7 @@ def test_one_failing_radio_does_not_block_others(self, mock_alchemy, mock_upsert assert summary['playlists_created'] == 1 assert summary['failed'] == ['Broken'] - @patch('app_helper.get_alchemy_radios') + @patch('database.get_alchemy_radios') @patch('tasks.radio_manager.create_or_replace_playlist') @patch('tasks.radio_manager.song_alchemy') def test_radio_with_no_results_creates_no_playlist(self, mock_alchemy, mock_upsert, diff --git a/test/unit/test_app_analysis.py b/test/unit/test_app_analysis.py index 0caa966f..a39d2cc0 100644 --- a/test/unit/test_app_analysis.py +++ b/test/unit/test_app_analysis.py @@ -29,7 +29,7 @@ class TestCleaningPage: def test_cleaning_page_returns_html(self, client): """Test that /cleaning returns HTML content""" - with patch('flask.render_template') as mock_render: + with patch('app_analysis.render_template') as mock_render: mock_render.return_value = "Cleaning Page" response = client.get('/cleaning') @@ -47,12 +47,12 @@ class TestStartAnalysisEndpoint: @pytest.fixture(autouse=True) def patch_active_analysis_task(self): - with patch('app_helper.get_active_main_task', return_value=None) as mock_active_task: + with patch('app_analysis.get_active_main_task', return_value=None) as mock_active_task: yield mock_active_task - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_successful_analysis_start_with_defaults( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -81,9 +81,9 @@ def test_successful_analysis_start_with_defaults( save_call_args = mock_save_status.call_args[0] assert save_call_args[1] == "main_analysis" - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') @patch('app_analysis.NUM_RECENT_ALBUMS', 5) @patch('app_analysis.TOP_N_MOODS', 10) def test_analysis_start_uses_config_defaults( @@ -108,9 +108,9 @@ def test_analysis_start_uses_config_defaults( # Check that args tuple contains (num_recent_albums, top_n_moods) assert call_kwargs['args'] == (5, 10) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_analysis_start_with_custom_params( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -136,9 +136,9 @@ def test_analysis_start_with_custom_params( call_kwargs = mock_queue.enqueue.call_args[1] assert call_kwargs['args'] == (10, 15) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_analysis_enqueue_task_parameters( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -162,9 +162,9 @@ def test_analysis_enqueue_task_parameters( assert call_args[1]['description'] == "Main Music Analysis" assert call_args[1]['job_timeout'] == -1 # No timeout - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_analysis_handles_missing_json( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -183,9 +183,9 @@ def test_analysis_handles_missing_json( # Should still work with defaults assert response.status_code == 202 - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_analysis_saves_pending_status( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -215,12 +215,12 @@ class TestStartCleaningEndpoint: @pytest.fixture(autouse=True) def patch_active_cleaning_task(self): - with patch('app_helper.get_active_main_task', return_value=None) as mock_active_task: + with patch('app_analysis.get_active_main_task', return_value=None) as mock_active_task: yield mock_active_task - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_successful_cleaning_start( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -244,9 +244,9 @@ def test_successful_cleaning_start( # Verify task status was saved mock_save_status.assert_called_once() - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_cleaning_enqueue_task_parameters( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -267,9 +267,9 @@ def test_cleaning_enqueue_task_parameters( assert call_args[1]['description'] == "Database Cleaning (Identify and Delete Orphaned Albums)" assert call_args[1]['job_timeout'] == -1 # No timeout - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_cleaning_saves_pending_status( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -288,9 +288,9 @@ def test_cleaning_saves_pending_status( call_args = mock_save_status.call_args[0] assert call_args[1] == "cleaning" - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_cleaning_cleans_up_previous_tasks( self, mock_save_status, mock_cleanup, mock_queue, client ): @@ -307,10 +307,10 @@ def test_cleaning_cleans_up_previous_tasks( # Verify cleanup was called before enqueueing new task mock_cleanup.assert_called_once() - @patch('app_helper.get_active_main_task', return_value={'task_id': 'existing-cleaning-123', 'status': 'STARTED'}) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.get_active_main_task', return_value={'task_id': 'existing-cleaning-123', 'status': 'STARTED'}) + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_cleaning_blocks_when_active_task_exists( self, mock_save_status, mock_cleanup, mock_queue, mock_get_active, client ): @@ -327,10 +327,10 @@ def test_cleaning_blocks_when_active_task_exists( class TestEndpointErrorHandling: """Tests for error handling in endpoints""" - @patch('app_helper.get_active_main_task', return_value=None) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.get_active_main_task', return_value=None) + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_analysis_handles_enqueue_failure( self, mock_save_status, mock_cleanup, mock_queue, mock_get_active, client ): @@ -341,10 +341,10 @@ def test_analysis_handles_enqueue_failure( with pytest.raises(Exception): client.post('/api/analysis/start', json={}) - @patch('app_helper.get_active_main_task', return_value={'task_id': 'existing-cleaning-123', 'status': 'STARTED', 'task_type': 'cleaning'}) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.get_active_main_task', return_value={'task_id': 'existing-cleaning-123', 'status': 'STARTED', 'task_type': 'cleaning'}) + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_analysis_blocks_when_another_batch_is_active( self, mock_save_status, mock_cleanup, mock_queue, mock_get_active, client ): @@ -358,9 +358,9 @@ def test_analysis_blocks_when_another_batch_is_active( mock_cleanup.assert_not_called() mock_queue.enqueue.assert_not_called() - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_analysis.rq_queue_high') + @patch('app_analysis.clean_up_previous_main_tasks') + @patch('app_analysis.save_task_status') def test_cleaning_handles_enqueue_failure( self, mock_save_status, mock_cleanup, mock_queue, client ): diff --git a/test/unit/test_app_clustering.py b/test/unit/test_app_clustering.py index 3ca794e7..e01c1a66 100644 --- a/test/unit/test_app_clustering.py +++ b/test/unit/test_app_clustering.py @@ -20,10 +20,10 @@ def client(app): class TestStartClusteringEndpoint: - @patch('app_helper.get_active_main_task', return_value=None) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_clustering.get_active_main_task', return_value=None) + @patch('app_clustering.rq_queue_high') + @patch('app_clustering.clean_up_previous_main_tasks') + @patch('app_clustering.save_task_status') def test_successful_clustering_start_with_no_active_task( self, mock_save_status, mock_cleanup, mock_queue, mock_get_active, client ): @@ -42,10 +42,10 @@ def test_successful_clustering_start_with_no_active_task( mock_cleanup.assert_called_once() mock_save_status.assert_called_once() - @patch('app_helper.get_active_main_task', return_value={'task_id': 'existing-clustering-123', 'status': 'STARTED'}) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_clustering.get_active_main_task', return_value={'task_id': 'existing-clustering-123', 'status': 'STARTED'}) + @patch('app_clustering.rq_queue_high') + @patch('app_clustering.clean_up_previous_main_tasks') + @patch('app_clustering.save_task_status') def test_clustering_blocks_when_active_task_exists( self, mock_save_status, mock_cleanup, mock_queue, mock_get_active, client ): @@ -58,10 +58,10 @@ def test_clustering_blocks_when_active_task_exists( mock_cleanup.assert_not_called() mock_queue.enqueue.assert_not_called() - @patch('app_helper.get_active_main_task', return_value={'task_id': 'existing-cleaning-123', 'status': 'STARTED', 'task_type': 'cleaning'}) - @patch('app_helper.rq_queue_high') - @patch('app_helper.clean_up_previous_main_tasks') - @patch('app_helper.save_task_status') + @patch('app_clustering.get_active_main_task', return_value={'task_id': 'existing-cleaning-123', 'status': 'STARTED', 'task_type': 'cleaning'}) + @patch('app_clustering.rq_queue_high') + @patch('app_clustering.clean_up_previous_main_tasks') + @patch('app_clustering.save_task_status') def test_clustering_blocks_when_another_batch_is_active( self, mock_save_status, mock_cleanup, mock_queue, mock_get_active, client ): diff --git a/test/unit/test_app_helper_task_note.py b/test/unit/test_app_helper_task_note.py index 822a11ac..7befa25c 100644 --- a/test/unit/test_app_helper_task_note.py +++ b/test/unit/test_app_helper_task_note.py @@ -109,7 +109,7 @@ def test_best_params_subset_size_preferred(self): 'num_playlists_created': 8, } result = _build_task_note('main_clustering', details, MagicMock()) - assert result == 'sampled: 500 • clusters: 8' + assert result == 'sampled: 500 | clusters: 8' def test_non_dict_best_params_falls_back_to_sampled_songs(self): details = {'best_params': 'oops', 'sampled_songs': 100} diff --git a/test/unit/test_app_logging.py b/test/unit/test_app_logging.py new file mode 100644 index 00000000..7bbefc37 --- /dev/null +++ b/test/unit/test_app_logging.py @@ -0,0 +1,108 @@ +"""Unit tests for app_logging's record sanitization. + +Covers the two jobs of ``_sanitize_log_text`` / ``LogSanitizingFilter``: +console-safety (emoji and non-Latin-1 symbol stripping for Windows code-pages) +and log-injection safety (CR/LF and other control characters are neutralised so +an attacker-controlled value cannot forge or split log lines -- CWE-117). + +Emoji and accented characters are built with ``chr()`` so this source file +stays pure ASCII. +""" +import logging + +from app_logging import _sanitize_log_text, LogSanitizingFilter, configure_logging + +_CHECK_MARK = chr(0x2705) +_MUSIC_NOTE = chr(0x1F3B5) +_ACCENTED = "caf" + chr(0xE9) + " se" + chr(0xF1) + "or " + chr(0xFC) + "ber" + + +class TestSanitizeLogText: + def test_removes_emoji_and_symbols(self): + assert _sanitize_log_text("done " + _CHECK_MARK) == "done" + assert _sanitize_log_text("track " + _MUSIC_NOTE + " ready") == "track ready" + + def test_newline_becomes_space(self): + assert _sanitize_log_text("hello\nworld") == "hello world" + + def test_crlf_collapses_to_single_space(self): + assert _sanitize_log_text("hello\r\nworld") == "hello world" + + def test_control_chars_become_space(self): + assert _sanitize_log_text("a\x00b\x07c\x7f") == "a b c" + + def test_unicode_line_separators_become_space(self): + for sep in (chr(0x85), chr(0x2028), chr(0x2029)): + assert _sanitize_log_text("a" + sep + "b") == "a b" + + def test_unicode_separator_cannot_forge_a_line(self): + forged = "user42" + chr(0x2028) + "[INFO]-[fake]-dropped all tables" + result = _sanitize_log_text(forged) + assert all(sep not in result for sep in (chr(0x85), chr(0x2028), chr(0x2029))) + assert len(result.splitlines()) == 1 + + def test_tab_is_preserved(self): + assert _sanitize_log_text("col1\tcol2") == "col1\tcol2" + + def test_latin1_accents_pass_through(self): + assert _sanitize_log_text(_ACCENTED) == _ACCENTED + + def test_log_injection_cannot_forge_a_line(self): + forged = "user42\n[INFO]-[fake]-dropped all tables" + result = _sanitize_log_text(forged) + assert "\n" not in result + assert "\r" not in result + assert result == "user42 [INFO]-[fake]-dropped all tables" + + def test_non_string_returned_unchanged(self): + assert _sanitize_log_text(123) == 123 + assert _sanitize_log_text(None) is None + + +class TestLogSanitizingFilter: + def _record(self, msg, args=None): + return logging.LogRecord( + name="test", level=logging.INFO, pathname=__file__, lineno=1, + msg=msg, args=args, exc_info=None, + ) + + def test_sanitizes_msg(self): + record = self._record("oops\ninjected " + _CHECK_MARK) + LogSanitizingFilter().filter(record) + assert record.msg == "oops injected" + + def test_sanitizes_tuple_args_leaving_non_strings(self): + record = self._record("%s %s", args=("a\nb", 7)) + LogSanitizingFilter().filter(record) + assert record.args == ("a b", 7) + + def test_sanitizes_dict_args(self): + record = self._record("%(x)s") + record.args = {"x": "p\nq", "n": 3} + LogSanitizingFilter().filter(record) + assert record.args == {"x": "p q", "n": 3} + + def test_filter_always_returns_true(self): + assert LogSanitizingFilter().filter(self._record("hi")) is True + + +class TestConfigureLogging: + def test_attaches_sanitizing_filter_once(self): + root = logging.getLogger() + saved = {handler: list(handler.filters) for handler in root.handlers} + try: + configure_logging() + configure_logging() + assert root.handlers + for handler in root.handlers: + count = sum(isinstance(f, LogSanitizingFilter) for f in handler.filters) + assert count == 1 + finally: + for handler in root.handlers: + if handler in saved: + handler.filters = saved[handler] + else: + handler.filters = [ + f for f in handler.filters + if not isinstance(f, LogSanitizingFilter) + ] diff --git a/test/unit/test_app_map_helpers.py b/test/unit/test_app_map_helpers.py index 1622c60a..81209d77 100644 --- a/test/unit/test_app_map_helpers.py +++ b/test/unit/test_app_map_helpers.py @@ -1,5 +1,3 @@ -import pytest - from app_map import _pick_top_mood, _round_coord, _sample_items diff --git a/test/unit/test_clustering.py b/test/unit/test_clustering.py index 33470293..773744f3 100644 --- a/test/unit/test_clustering.py +++ b/test/unit/test_clustering.py @@ -373,7 +373,7 @@ class TestSanitizeForJson: def test_sanitize_numpy_array(self): """Test numpy array conversion to list""" - from tasks.clustering import _sanitize_for_json + from sanitization import sanitize_for_json as _sanitize_for_json obj = np.array([1.0, 2.0, 3.0]) result = _sanitize_for_json(obj) @@ -383,7 +383,7 @@ def test_sanitize_numpy_array(self): def test_sanitize_numpy_integers(self): """Test numpy integer conversion""" - from tasks.clustering import _sanitize_for_json + from sanitization import sanitize_for_json as _sanitize_for_json obj = { 'int8': np.int8(42), @@ -400,7 +400,7 @@ def test_sanitize_numpy_integers(self): def test_sanitize_numpy_floats(self): """Test numpy float conversion""" - from tasks.clustering import _sanitize_for_json + from sanitization import sanitize_for_json as _sanitize_for_json obj = { 'float32': np.float32(3.14), @@ -415,7 +415,7 @@ def test_sanitize_numpy_floats(self): def test_sanitize_numpy_bool(self): """Test numpy bool conversion""" - from tasks.clustering import _sanitize_for_json + from sanitization import sanitize_for_json as _sanitize_for_json obj = {'flag': np.bool_(True)} result = _sanitize_for_json(obj) @@ -425,7 +425,7 @@ def test_sanitize_numpy_bool(self): def test_sanitize_nested_structures(self): """Test sanitization of nested dictionaries and lists""" - from tasks.clustering import _sanitize_for_json + from sanitization import sanitize_for_json as _sanitize_for_json obj = { 'array': np.array([1, 2, 3]), @@ -443,7 +443,7 @@ def test_sanitize_nested_structures(self): def test_sanitize_preserves_native_types(self): """Test that native Python types are preserved""" - from tasks.clustering import _sanitize_for_json + from sanitization import sanitize_for_json as _sanitize_for_json obj = { 'string': 'hello', diff --git a/test/unit/test_error_manager.py b/test/unit/test_error_manager.py index d328d951..e54e0bec 100644 --- a/test/unit/test_error_manager.py +++ b/test/unit/test_error_manager.py @@ -8,8 +8,6 @@ import os import sys -import pytest - REPO_ROOT = os.path.normpath( os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..') ) diff --git a/test/unit/test_import_architecture.py b/test/unit/test_import_architecture.py index d8c0fd18..a11fff09 100644 --- a/test/unit/test_import_architecture.py +++ b/test/unit/test_import_architecture.py @@ -2,14 +2,22 @@ Only module-level (eager) imports count here; function-level imports are the sanctioned escape hatch used across the codebase (mediaserver providers, -voyager/app_helper consumers, config's DB-override loader). Three invariants -keep the graph flat and acyclic so deep chains and cycles cannot creep back in: +voyager/app_helper consumers, config's DB-override loader). Six invariants +keep the graph flat, layered, and acyclic so deep chains and cycles cannot +creep back in: 1. Foundation modules stay leaves: they import nothing internal at module level. 2. No module-level import cycles, except the lyrics package init whose try/except fallback design is deliberately order-dependent. -3. No eager import chain may exceed MAX_CHAIN modules (app -> blueprint -> - hub/manager -> leaf). Anything deeper must use a function-level import. +3. No eager import chain may exceed MAX_CHAIN modules (e.g. app -> blueprint -> + helper -> service-facade -> leaf). Anything deeper must use a function-level + import. +4. Layers point downward: a module may import its own or any lower layer, never + a higher one (checked transitively, so indirect chains count too). +5. Forbidden edges: specific module-level dependencies are banned outright + (e.g. the database/queue layer must never import the app_helper facade). +6. Independence: route blueprints never import one another -- they compose only + through app.py. """ import ast @@ -29,6 +37,9 @@ "config", "tz_helper", "error.error_dictionary", + "ssrf_guard", + "sanitization", + "tasks.memory_utils", } ALLOWED_CYCLES = { @@ -36,7 +47,57 @@ frozenset({"error", "error.error_manager"}), } -MAX_CHAIN = 6 # allows for package __init__.py → submodule edges adding 1-2 phantom hops +# The honest eager-import floor for this codebase is 5, not 4: two intentional +# 3-deep facades are each imported one hop below the root task/blueprint modules. +# - error/ package: error -> error.error_manager -> error.error_dictionary +# - AI dispatch: tasks.ai.api -> {providers.openai, prompts} -> config +# Flattening either to reach 4 would dismantle a deliberate facade (merging the +# error dictionary/manager, or inlining the AI providers) and hurt readability, +# so 5 is the enforced ceiling. Anything deeper must use a function-level import. +MAX_CHAIN = 5 + +# --- Layered architecture (invariant 4) ------------------------------------ +# Ordered low -> high. A module may import its own layer or any LOWER layer, +# never a higher one (checked transitively). Only listed modules are +# constrained; cycle-participating modules (error/, lyrics/) are intentionally +# omitted and governed by the cycles test instead. +LAYERS = [ + {"config", "tz_helper", "error.error_dictionary", "ssrf_guard", "sanitization", "tasks.memory_utils"}, + {"database", "taskqueue", "tasks.ai.prompts", + "tasks.ai.providers.openai", "tasks.ai.providers.gemini", "tasks.ai.providers.mistral"}, + {"app_helper", "app_helper_artist", "tasks.ai.api"}, + {"tasks.clustering_helper", "tasks.analysis_helper"}, + {"tasks.clustering", "tasks.analysis"}, + {"app"}, +] + +# --- Forbidden module-level dependencies (invariant 5) --------------------- +# (importer, target): the importer must not reach the target via any eager +# import chain (direct or indirect). Keeps the data/queue layer and the AI +# transport modules from depending on the higher facade/domain layers. +FORBIDDEN_IMPORTS = [ + ("database", "app_helper"), + ("taskqueue", "app_helper"), + ("database", "app_helper_artist"), + ("taskqueue", "app_helper_artist"), + ("tasks.ai.prompts", "tasks.ai.api"), + ("tasks.ai.providers.openai", "tasks.ai.api"), + ("tasks.ai.providers.gemini", "tasks.ai.api"), + ("tasks.ai.providers.mistral", "tasks.ai.api"), + ("app_helper", "tasks.clustering"), + ("app_helper", "tasks.analysis"), +] + +# --- Independence groups (invariant 6) ------------------------------------- +# Members must not import one another (direct or indirect). Flask route +# blueprints are self-contained features that compose only through app.py. +INDEPENDENT_GROUPS = [ + {"app_chat", "app_clustering", "app_analysis", "app_cron", "app_voyager", + "app_sonic_fingerprint", "app_path", "app_external", "app_alchemy", "app_map", + "app_waveform", "app_artist_similarity", "app_clap_search", "app_lyrics", + "app_sem_grove", "app_backup", "app_provider_migration", "app_dashboard", + "app_users", "app_sync"}, +] def _collect_modules(): @@ -85,12 +146,21 @@ def _build_eager_graph(modules): targets = [alias.name for alias in node.names] elif isinstance(node, ast.ImportFrom): base = _resolve_relative(node.module or "", node.level, name, is_package) if node.level else (node.module or "") - targets = [base] + [f"{base}.{alias.name}" for alias in node.names if base] + # Track the imported names (base.name). The ancestor-package + # loop below also charges the base package itself, because + # importing base.name runs base/__init__.py first. + targets = [f"{base}.{alias.name}" for alias in node.names if base] if base else [alias.name for alias in node.names] else: continue for target in targets: parts = target.split(".") - for i in range(1, len(parts) + 1): + # Importing "a.b.c" (or "from a.b import c") runs a/__init__.py + # and a/b/__init__.py before binding the target, so the importer + # eagerly depends on every ancestor package that is a real module + # AND the deepest matching module. Empty package __init__ files + # are graph leaves and add no depth; only packages that import at + # module level (e.g. error/ and lyrics/) actually extend a chain. + for i in range(len(parts), 0, -1): candidate = ".".join(parts[:i]) if candidate in modules and candidate != name: graph[name].add(candidate) @@ -146,36 +216,134 @@ def _find_cycles(graph, modules): def _longest_chain(graph, modules): + """Longest eager-import chain (module count), with each module counted once. + + The graph may contain the allowed package-init cycles (error/, lyrics/), so a + naive longest-path walk could revisit a node and over-count. We first strip + cycle back-edges with a DFS to obtain a DAG, then take its longest path with + memoization. Stripping back-edges keeps every chain a simple path while + preserving the honest forward depth, e.g. + app -> app_setup -> error -> error.error_manager -> error.error_dictionary. + """ + # Strip back-edges (edges into the current DFS stack) so cycles cannot + # inflate a chain by revisiting a module. + color = {} # 0 = unvisited, 1 = on stack, 2 = done + dag = defaultdict(set) + + def _strip(u): + color[u] = 1 + for v in sorted(graph.get(u, ())): + c = color.get(v, 0) + if c == 1: + continue # back-edge into the DFS stack (a cycle): drop it + dag[u].add(v) + if c == 0: + _strip(v) + color[u] = 2 + + for root in sorted(modules): + if color.get(root, 0) == 0: + _strip(root) + + # Longest path on the resulting DAG (memoizable: no cycles remain). cache = {} - def depth_from(node, path): + def depth_from(node): if node in cache: return cache[node] best = (1, (node,)) - for succ in graph.get(node, ()): - if succ in path: - continue - sub_len, sub_chain = depth_from(succ, path | {node}) + for succ in dag.get(node, ()): + sub_len, sub_chain = depth_from(succ) if 1 + sub_len > best[0]: best = (1 + sub_len, (node,) + sub_chain) - if not any(s in path for s in graph.get(node, ())): - cache[node] = best + cache[node] = best return best overall = (0, ()) for node in modules: - candidate = depth_from(node, frozenset()) + candidate = depth_from(node) if candidate[0] > overall[0]: overall = candidate return overall +def _max_chains(graph, modules): + """Return ``(max_len, chains)`` -- every maximal simple eager chain whose + length equals the longest. Used for the human-readable recap so a PR that + deepens the graph shows exactly which chains now sit at the ceiling. + """ + best_len = 0 + chains = [] + + def depth_first(node, path): + nonlocal best_len, chains + extended = False + for succ in sorted(graph.get(node, ())): + if succ in path: + continue + extended = True + depth_first(succ, path + (succ,)) + if not extended: # maximal: cannot extend further + n = len(path) + if n > best_len: + best_len, chains = n, [path] + elif n == best_len: + chains.append(path) + + for node in sorted(modules): + depth_first(node, (node,)) + return best_len, sorted(set(chains)) + + @lru_cache(maxsize=1) def _graph(): modules = _collect_modules() return modules, _build_eager_graph(modules) +def architecture_report(): + """Human-readable diagnostic: the layer table + direction tally, the measured + max eager chain vs the ceiling, and every chain tied at the maximum. + + Rendered into the pytest terminal summary by ``test/unit/conftest.py`` so the + numbers show on every run; the depth gate reuses the chain recap in its + failure message. + """ + modules, graph = _graph() + level = {m: i for i, layer in enumerate(LAYERS) for m in layer} + + down = horiz = up = 0 + for src, dsts in graph.items(): + if src not in level: + continue + for dst in dsts: + if dst not in level: + continue + if level[dst] > level[src]: + up += 1 + elif level[dst] == level[src]: + horiz += 1 + else: + down += 1 + + max_len, chains = _max_chains(graph, modules) + + lines = ["Layers (L0 = foundation, ascending to the app entrypoint); " + "every dependency must point DOWN to a lower or equal layer:"] + for i, layer in enumerate(LAYERS): + lines.append(f" L{i}: " + ", ".join(sorted(layer))) + lines.append(f" layered edges: {down} downward (ok), {horiz} horizontal/same-layer, " + f"{up} upward (ILLEGAL)") + lines.append("") + status = "OK" if max_len <= MAX_CHAIN else "OVER CEILING" + lines.append(f"Max eager import chain: {max_len} modules " + f"(ceiling MAX_CHAIN={MAX_CHAIN}) -> {status}") + lines.append(f"Chains at depth {max_len} ({len(chains)}):") + for chain in chains: + lines.append(" " + " -> ".join(chain)) + return lines + + def test_foundation_modules_are_leaves(): _, graph = _graph() violations = {leaf: sorted(graph.get(leaf, ())) for leaf in LEAF_MODULES if graph.get(leaf)} @@ -196,9 +364,70 @@ def test_no_module_level_import_cycles(): def test_eager_import_chains_stay_shallow(): modules, graph = _graph() - length, chain = _longest_chain(graph, modules) - assert length <= MAX_CHAIN, ( - f"Eager import chain of {length} modules exceeds the maximum of " - f"{MAX_CHAIN}: {' -> '.join(chain)}. Convert one edge to a " - f"function-level import to flatten it." + length, _ = _longest_chain(graph, modules) + if length > MAX_CHAIN: + max_len, chains = _max_chains(graph, modules) + recap = "\n ".join(" -> ".join(c) for c in chains) + raise AssertionError( + f"Eager import chain of {length} modules exceeds the maximum of " + f"{MAX_CHAIN}. Convert one edge in each NEW chain below to a " + f"function-level import to flatten it.\n" + f"All chains at depth {max_len}:\n {recap}" + ) + + +def _reachable(graph, start): + """All modules eagerly reachable from ``start`` (direct + indirect).""" + seen = set() + stack = list(graph.get(start, ())) + while stack: + node = stack.pop() + if node in seen: + continue + seen.add(node) + stack.extend(graph.get(node, ())) + return seen + + +def test_layered_dependencies_point_downward(): + modules, graph = _graph() + level = {m: i for i, layer in enumerate(LAYERS) for m in layer} + unknown = sorted(m for m in level if m not in modules) + assert not unknown, f"LAYERS references modules that no longer exist: {unknown}" + + violations = [] + for src, src_level in level.items(): + for dst in _reachable(graph, src): + if dst in level and level[dst] > src_level: + violations.append(f"{src} (layer {src_level}) -> {dst} (layer {level[dst]})") + assert not violations, ( + "Lower layers must not import higher layers at module level (move the " + "dependency down, or push the import inside the function that uses it):\n " + + "\n ".join(sorted(violations)) + ) + + +def test_forbidden_imports(): + modules, graph = _graph() + violations = [] + for src, dst in FORBIDDEN_IMPORTS: + if src in modules and dst in _reachable(graph, src): + violations.append(f"{src} -> ... -> {dst}") + assert not violations, ( + "Forbidden module-level dependencies detected (these layers must not " + "depend on the higher ones):\n " + "\n ".join(violations) + ) + + +def test_independent_modules_do_not_cross_import(): + modules, graph = _graph() + violations = [] + for group in INDEPENDENT_GROUPS: + present = group & set(modules) + for src in sorted(present): + for dst in sorted(_reachable(graph, src) & present - {src}): + violations.append(f"{src} -> {dst}") + assert not violations, ( + "Independent modules must not import one another at module level " + "(compose them through app.py instead):\n " + "\n ".join(violations) ) diff --git a/test/unit/test_import_smoke.py b/test/unit/test_import_smoke.py new file mode 100644 index 00000000..ef502db2 --- /dev/null +++ b/test/unit/test_import_smoke.py @@ -0,0 +1,74 @@ +"""Smoke test: every project module must import cleanly. + +Complements the flake8 F821 (undefined-name) gate. F821 is static and cannot see +cross-module breakage like ``from x import name_that_no_longer_exists`` or +``import missing_module`` -- actually importing each module catches those plus +any other module-load-time error (a moved function leaving a dangling import, a +typo in a re-export, etc.). + +A psycopg2 / redis *connection* error means the import machinery already +succeeded (the module merely tried to reach a service while loading, e.g. +``app`` runs ``init_db()`` at import) and is treated as a pass. Genuinely-missing +optional native deps are skipped; everything else (ImportError, NameError, +SyntaxError, AttributeError) fails the test. +""" +import importlib +import os +from pathlib import Path + +import psycopg2 +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] + +# Mirrors test_import_architecture._collect_modules excludes (plus scripts/, which +# are standalone build entry points not meant to import in a server context). +EXCLUDED_DIRS = { + ".git", ".venv", ".venv-windows", "node_modules", "__pycache__", + "build", "dist", "pginstall", "native-build", "test", "scripts", +} + +# ImportErrors naming an optional/native dependency that may be absent in CI. +_OPTIONAL_DEPS = ("cuml", "cupy", "voyager", "faiss", "tensorflow") + + +def _discover_modules(): + modules = [] + for dirpath, dirnames, filenames in os.walk(REPO_ROOT): + dirnames[:] = [d for d in dirnames if d not in EXCLUDED_DIRS and not d.startswith(".")] + for filename in filenames: + if not filename.endswith(".py"): + continue + parts = list((Path(dirpath) / filename).relative_to(REPO_ROOT).parts) + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][:-3] + if parts: + modules.append(".".join(parts)) + return sorted(set(modules)) + + +MODULES = _discover_modules() + + +@pytest.mark.parametrize("modname", MODULES) +def test_module_imports_cleanly(modname): + """Importing the module must not raise an ImportError/NameError/etc. + + Reaching a DB/Redis connection during import is fine -- it proves the module + loaded; only a real load-time error should fail. + """ + try: + importlib.import_module(modname) + except psycopg2.OperationalError: + pytest.skip(f"{modname}: reached DB connection during import (machinery OK)") + except ImportError as exc: + msg = str(exc).lower() + if any(dep in msg for dep in _OPTIONAL_DEPS): + pytest.skip(f"{modname}: optional dependency absent ({exc})") + raise + except Exception as exc: # noqa: BLE001 -- surface the real failure + if "redis" in type(exc).__module__ and "connect" in str(exc).lower(): + pytest.skip(f"{modname}: reached Redis connection during import (machinery OK)") + raise diff --git a/test/unit/test_index_build_helpers.py b/test/unit/test_index_build_helpers.py index 641a3a62..f8fbc5a9 100644 --- a/test/unit/test_index_build_helpers.py +++ b/test/unit/test_index_build_helpers.py @@ -759,7 +759,7 @@ def test_rejects_batch_with_mismatched_ids_length(self): def test_skips_empty_batches_silently(self): try: - import voyager + import voyager # noqa: F401 except ImportError: pytest.skip("voyager not installed") rng = np.random.default_rng(5) diff --git a/test/unit/test_mcp_server.py b/test/unit/test_mcp_server.py index 17fa8d1a..56b987a9 100644 --- a/test/unit/test_mcp_server.py +++ b/test/unit/test_mcp_server.py @@ -3,7 +3,8 @@ Tests cover MCP helper + tool functions: - get_library_context(): Library statistics with caching - _database_genre_query_sync(): Genre regex matching, filters, relevance scoring -- _ai_brainstorm_sync(): Two-stage matching (exact + fuzzy normalized) +- _extract_json_object() / _clamp_recipe(): brainstorm recipe parse + vocab clamp +- _ai_brainstorm_sync(): grounded recipe -> fused retrieval channels (#643) - _song_similarity_api_sync(): Song lookup with exact/fuzzy fallback - Energy normalization in execute_mcp_tool() - Pre-execution validation (filterless search_database rejection) @@ -533,60 +534,94 @@ def test_rerouting_applied_in_database_query(self): # --------------------------------------------------------------------------- -# ai_brainstorm normalization patterns (unit-testable without DB) +# brainstorm recipe helpers (pure, no DB) # --------------------------------------------------------------------------- @pytest.mark.unit -class TestBrainstormNormalization: - """Test the normalization logic used in _ai_brainstorm_sync.""" +class TestExtractJsonObject: + """_extract_json_object: recover ONE JSON object from messy model output.""" - def _normalize(self, text): - """Reproduce the normalization from mcp_server.""" - return (text.lower() - .replace(' ', '') - .replace('-', '') - .replace("'", '') - .replace('.', '') - .replace(',', '')) + def _fn(self): + return _import_mcp_impl()._extract_json_object - def test_lowercase(self): - assert self._normalize("Hello") == "hello" + def test_plain_object(self): + assert self._fn()('{"a": 1}') == {"a": 1} - def test_remove_spaces(self): - assert self._normalize("The Beatles") == "thebeatles" + def test_fenced_object(self): + assert self._fn()('```json\n{"a": 1}\n```') == {"a": 1} - def test_remove_dashes(self): - assert self._normalize("up-beat") == "upbeat" + def test_think_preamble_stripped(self): + assert self._fn()('reasoning...\n{"a": 2}') == {"a": 2} - def test_remove_apostrophes(self): - assert self._normalize("Don't Stop") == "dontstop" + def test_object_embedded_in_prose(self): + assert self._fn()('Sure, here: {"a": 3} hope it helps') == {"a": 3} - def test_remove_periods(self): - assert self._normalize("Mr. Jones") == "mrjones" + def test_array_is_rejected(self): + assert self._fn()('[{"a": 1}]') is None - def test_remove_commas(self): - assert self._normalize("Hello, World") == "helloworld" + def test_garbage_returns_none(self): + assert self._fn()('no json here at all') is None - def test_ac_dc_normalization(self): - """AC/DC normalizes consistently (slash not removed but spaces/dots are).""" - result = self._normalize("AC DC") - assert result == "acdc" + def test_empty_returns_none(self): + assert self._fn()('') is None - def test_complex_normalization(self): - assert self._normalize("Don't Stop Me Now") == "dontstopmenow" - def test_both_title_and_artist_required(self): - """Demonstrate that matching requires BOTH title and artist.""" - title_norm = self._normalize("Bohemian Rhapsody") - artist_norm = self._normalize("Queen") - assert title_norm == "bohemianrhapsody" - assert artist_norm == "queen" +@pytest.mark.unit +class TestClampRecipe: + """_clamp_recipe: normalise a raw recipe to library-valid, bounded values.""" + + def _fn(self): + return _import_mcp_impl()._clamp_recipe + + def test_genres_clamped_to_vocab_case_and_punct_insensitive(self): + out = self._fn()({"filters": {"genres": ["hip hop", "ROCK", "not-a-genre"]}}) + assert out["filters"]["genres"] == ["Hip-Hop", "rock"] + + def test_moods_and_voices_clamped(self): + out = self._fn()({"filters": {"moods": ["Party", "bogus"], "voices": ["female vocalist"]}}) + assert out["filters"]["moods"] == ["party"] + assert out["filters"]["voices"] == ["female vocalist"] - def test_same_title_different_artist_not_equal(self): - """Same title with different artist should not be considered same.""" - t1 = self._normalize("Yesterday") + "|" + self._normalize("The Beatles") - t2 = self._normalize("Yesterday") + "|" + self._normalize("Some Cover Artist") - assert t1 != t2 + def test_year_range_reversed_is_swapped(self): + out = self._fn()({"filters": {"year_min": 2009, "year_max": 1990}}) + assert out["filters"]["year_min"] == 1990 + assert out["filters"]["year_max"] == 2009 + + def test_energy_clamped_then_swapped(self): + out = self._fn()({"filters": {"energy_min": 2.0, "energy_max": -1.0}}) + assert out["filters"]["energy_min"] == pytest.approx(0.0) + assert out["filters"]["energy_max"] == pytest.approx(1.0) + + def test_lists_deduped_and_capped(self): + import config as cfg + out = self._fn()({ + "sound_descriptions": ["a", "a", "b", "c", "d", "e"], + "seed_artists": ["X", "x", "Y", "Z", "W", "V"], + "lyric_themes": ["t1", "t2", "t3"], + }) + assert out["sound_descriptions"][:2] == ["a", "b"] + assert len(out["sound_descriptions"]) <= cfg.AI_BRAINSTORM_SOUND_DESCRIPTIONS_MAX + assert len(out["seed_artists"]) <= cfg.AI_BRAINSTORM_SEED_ARTISTS_MAX + assert len(out["lyric_themes"]) <= cfg.AI_BRAINSTORM_LYRIC_THEMES_MAX + + def test_missing_filters_yields_empty_defaults(self): + out = self._fn()({}) + f = out["filters"] + assert f["genres"] == [] and f["moods"] == [] and f["voices"] == [] + assert f["year_min"] is None and f["energy_max"] is None + assert out["sound_descriptions"] == [] and out["seed_artists"] == [] + + def test_non_list_fields_are_coerced(self): + out = self._fn()({"sound_descriptions": "just one", "filters": {"genres": "rock"}}) + assert out["sound_descriptions"] == ["just one"] + assert out["filters"]["genres"] == ["rock"] + + def test_seed_artists_suppressed_when_disabled(self): + mod = _import_mcp_impl() + import config as cfg + with patch.object(cfg, "AI_BRAINSTORM_USE_ARTIST_SEEDS", False): + out = mod._clamp_recipe({"seed_artists": ["Nas", "Jay-Z"]}) + assert out["seed_artists"] == [] # --------------------------------------------------------------------------- @@ -986,222 +1021,130 @@ def test_result_structure(self): @pytest.mark.unit class TestAiBrainstormSync: - """Tests for _ai_brainstorm_sync - AI knowledge brainstorming with two-stage matching.""" + """_ai_brainstorm_sync emits a grounded recipe and fuses retrieval channels (#643).""" - def _make_ai_module(self, response="[]"): + def _make_ai_module(self, response): mock_mod = MagicMock() mock_mod.generate_text = Mock(return_value=response) return mock_mod def _make_ai_config(self): - return { - "provider": "gemini", - "gemini_key": "fake-key", - "gemini_model": "gemini-pro", + return {"provider": "gemini", "gemini_key": "fake-key", "gemini_model": "gemini-pro"} + + def _recipe(self, **over): + base = { + "filters": {"genres": ["rock"]}, + "sound_descriptions": ["driving guitar rock"], + "seed_artists": [], + "lyric_themes": [], } + base.update(over) + return json.dumps(base) + + def _patch_channels(self, mod, audio=None, artist=None, lyrics=None, filt=None): + empty = {"songs": []} + return ( + patch.object(mod, '_text_search_sync', return_value=audio if audio is not None else empty), + patch.object(mod, '_artist_similarity_api_sync', return_value=artist if artist is not None else empty), + patch.object(mod, '_lyrics_search_sync', return_value=lyrics if lyrics is not None else empty), + patch.object(mod, '_database_genre_query_sync', return_value=filt if filt is not None else empty), + ) - def _setup_cursor(self): - cur = MagicMock() - cur.__enter__ = Mock(return_value=cur) - cur.__exit__ = Mock(return_value=False) - cur.fetchall = Mock(return_value=[]) - return cur - - def test_ai_error_response_returns_empty(self): - """AI returns 'Error: ...' -> result has empty songs.""" + def test_ai_error_returns_empty(self): + """AI transport error -> empty songs, no channels touched.""" mod = _import_mcp_impl() - cur = self._setup_cursor() - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - ai_mod = self._make_ai_module("Error: API rate limit exceeded") - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): result = mod._ai_brainstorm_sync("rock classics", self._make_ai_config(), 10) - assert result["songs"] == [] - assert "Error" in result["message"] - def test_valid_json_array_parsed(self): - """AI returns valid JSON array, DB finds matching rows.""" + def test_unparseable_returns_empty_without_traceback(self): + """No JSON object recoverable -> empty + generic message (never a traceback).""" mod = _import_mcp_impl() - cur = self._setup_cursor() - - ai_response = json.dumps([ - {"title": "Bohemian Rhapsody", "artist": "Queen"}, - {"title": "Stairway to Heaven", "artist": "Led Zeppelin"}, - ]) - ai_mod = self._make_ai_module(ai_response) - - cur.fetchall = Mock(return_value=[ - _make_dict_row({"item_id": "100", "title": "Bohemian Rhapsody", "author": "Queen"}), - _make_dict_row({"item_id": "101", "title": "Stairway to Heaven", "author": "Led Zeppelin"}), - ]) - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): - result = mod._ai_brainstorm_sync("classic rock", self._make_ai_config(), 10) - - assert len(result["songs"]) == 2 - titles = [s["title"] for s in result["songs"]] - assert "Bohemian Rhapsody" in titles + ai_mod = self._make_ai_module("here are some great rock songs, but no json") + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): + result = mod._ai_brainstorm_sync("rock", self._make_ai_config(), 10) + assert result["songs"] == [] + assert "Traceback" not in result["message"] - def test_markdown_code_blocks_stripped(self): - """AI response wrapped in ```json...``` is still parsed correctly.""" + def test_recipe_drives_channels_and_fuses(self): + """Recipe runs audio + artist + filter channels and unions their results.""" mod = _import_mcp_impl() - cur = self._setup_cursor() - - ai_response = '```json\n[{"title": "Hey Jude", "artist": "The Beatles"}]\n```' - ai_mod = self._make_ai_module(ai_response) - - cur.fetchall = Mock(return_value=[ - _make_dict_row({"item_id": "200", "title": "Hey Jude", "author": "The Beatles"}), - ]) - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): - result = mod._ai_brainstorm_sync("beatles hits", self._make_ai_config(), 10) - - assert len(result["songs"]) == 1 - assert result["songs"][0]["title"] == "Hey Jude" - - def test_stage1_exact_match(self): - """AI suggests song in DB with exact title+artist -> found via stage 1.""" + ai_mod = self._make_ai_module(self._recipe(seed_artists=["Nirvana"])) + p_audio, p_artist, p_lyrics, p_filt = self._patch_channels( + mod, + audio={"songs": [{"item_id": "1", "title": "A", "artist": "X"}]}, + artist={"songs": [{"item_id": "2", "title": "B", "artist": "Nirvana"}]}, + filt={"songs": [{"item_id": "3", "title": "C", "artist": "Z"}]}, + ) + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}), \ + p_audio as a, p_artist as ar, p_lyrics, p_filt as f: + result = mod._ai_brainstorm_sync("90s rock like Nirvana", self._make_ai_config(), 50) + ids = sorted(s["item_id"] for s in result["songs"]) + assert ids == ["1", "2", "3"] + assert a.called and ar.called and f.called + + def test_dedup_across_channels(self): + """The same item returned by several channels appears once.""" mod = _import_mcp_impl() - cur = self._setup_cursor() - - ai_response = json.dumps([{"title": "Creep", "artist": "Radiohead"}]) - ai_mod = self._make_ai_module(ai_response) - - cur.fetchall = Mock(return_value=[ - _make_dict_row({"item_id": "300", "title": "Creep", "author": "Radiohead"}), - ]) - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): - result = mod._ai_brainstorm_sync("90s alternative", self._make_ai_config(), 10) - + ai_mod = self._make_ai_module(self._recipe(seed_artists=["Nirvana"])) + dup = {"songs": [{"item_id": "1", "title": "A", "artist": "X"}]} + p_audio, p_artist, p_lyrics, p_filt = self._patch_channels(mod, audio=dup, artist=dup, filt=dup) + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}), p_audio, p_artist, p_lyrics, p_filt: + result = mod._ai_brainstorm_sync("x", self._make_ai_config(), 50) assert len(result["songs"]) == 1 - assert result["songs"][0]["item_id"] == "300" - def test_stage2_fuzzy_normalized_match(self): - """AI suggests 'Don't Stop Me Now' by 'Queen', DB has 'Dont Stop Me Now' -> fuzzy match.""" + def test_get_songs_cap_respected(self): + """The fused pool never exceeds get_songs.""" mod = _import_mcp_impl() - cur = self._setup_cursor() - - ai_response = json.dumps([{"title": "Don't Stop Me Now", "artist": "Queen"}]) - ai_mod = self._make_ai_module(ai_response) - - call_count = [0] - - def _fetchall_side_effect(): - call_count[0] += 1 - if call_count[0] == 1: - return [] - else: - return [_make_dict_row({ - "item_id": "400", - "title": "Dont Stop Me Now", - "author": "Queen" - })] - - cur.fetchall = Mock(side_effect=_fetchall_side_effect) - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): - result = mod._ai_brainstorm_sync("fun queen songs", self._make_ai_config(), 10) + ai_mod = self._make_ai_module(self._recipe()) + many = {"songs": [{"item_id": str(i), "title": f"T{i}", "artist": f"A{i}"} for i in range(100)]} + p_audio, p_artist, p_lyrics, p_filt = self._patch_channels(mod, audio=many) + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}), p_audio, p_artist, p_lyrics, p_filt: + result = mod._ai_brainstorm_sync("x", self._make_ai_config(), 10) + assert len(result["songs"]) == 10 - assert len(result["songs"]) == 1 - assert result["songs"][0]["item_id"] == "400" - - def test_normalize_logic(self): - """Verify _normalize strips spaces, dashes, apostrophes, periods, commas.""" - # Reproduce the normalization regex from _ai_brainstorm_sync - def _normalize(s): - return re.sub(r"[\s\-\u2010\u2011\u2012\u2013\u2014/'\".,!?()]", '', s).lower() - - assert _normalize("Don't Stop Me Now") == "dontstopmenow" - assert _normalize("Mr. Jones") == "mrjones" - assert _normalize("Hello, World") == "helloworld" - assert _normalize("up-beat") == "upbeat" - assert _normalize("rock & roll") == "rock&roll" - - def test_escape_like(self): - """_escape_like escapes % and _ characters.""" - def _escape_like(s): - return s.replace('%', r'\%').replace('_', r'\_') - - assert _escape_like("100%") == r"100\%" - assert _escape_like("under_score") == r"under\_score" - assert _escape_like("normal") == "normal" - - def test_float_get_songs_converted_to_int(self): - """Passing get_songs=50.0 (Gemini float) should not raise.""" + def test_float_get_songs_does_not_raise(self): + """Providers may send get_songs as a float (e.g. 50.0).""" mod = _import_mcp_impl() - cur = self._setup_cursor() - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - ai_response = json.dumps([{"title": "Song", "artist": "Artist"}]) - ai_mod = self._make_ai_module(ai_response) - - cur.fetchall = Mock(return_value=[]) - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): + ai_mod = self._make_ai_module(self._recipe()) + p_audio, p_artist, p_lyrics, p_filt = self._patch_channels(mod) + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}), p_audio, p_artist, p_lyrics, p_filt: result = mod._ai_brainstorm_sync("test", self._make_ai_config(), 50.0) - assert "songs" in result - def test_invalid_json_returns_empty(self): - """AI returns non-JSON text -> result has empty songs.""" + def test_year_gate_excludes_out_of_era_sound_results(self): + """A year in the recipe gates the sound channel: CLAP can surface any-era + songs, but only in-era ones survive (the #643 'best rap of the 90s' fix).""" mod = _import_mcp_impl() - cur = self._setup_cursor() - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - ai_mod = self._make_ai_module("Here are some great rock songs that you might enjoy!") - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): - result = mod._ai_brainstorm_sync("rock", self._make_ai_config(), 10) - - assert result["songs"] == [] - assert "parse" in result["message"].lower() or "Failed" in result["message"] - - def test_results_trimmed_to_get_songs(self): - """AI suggests 30 songs, get_songs=10 -> only 10 returned.""" - mod = _import_mcp_impl() - cur = self._setup_cursor() - - suggestions = [ - {"title": f"Song {i}", "artist": f"Artist {i}"} for i in range(30) - ] - ai_response = json.dumps(suggestions) - ai_mod = self._make_ai_module(ai_response) - - exact_rows = [ - _make_dict_row({"item_id": str(i), "title": f"Song {i}", "author": f"Artist {i}"}) - for i in range(30) - ] - cur.fetchall = Mock(return_value=exact_rows) - conn = _make_connection(cur) - conn.cursor = Mock(return_value=cur) - - with patch.object(mod, 'get_db_connection', return_value=conn), \ - patch.dict(sys.modules, {'tasks.ai.api': ai_mod}): - result = mod._ai_brainstorm_sync("test", self._make_ai_config(), 10) - - assert len(result["songs"]) <= 10 + ai_mod = self._make_ai_module(json.dumps({ + "filters": {"genres": ["rock"], "year_min": 1990, "year_max": 1999}, + "sound_descriptions": ["driving guitar rock"], + "seed_artists": [], + "lyric_themes": [], + })) + audio = {"songs": [ + {"item_id": "in", "title": "In Era", "artist": "X"}, + {"item_id": "out", "title": "Out Era", "artist": "Y"}, + ]} + in_era = {"in"} + + def _db(*args, **kwargs): + cids = kwargs.get("candidate_item_ids") + if cids: + return {"songs": [s for s in audio["songs"] if s["item_id"] in cids and s["item_id"] in in_era]} + return {"songs": []} + + with patch.dict(sys.modules, {'tasks.ai.api': ai_mod}), \ + patch.object(mod, '_text_search_sync', return_value=audio), \ + patch.object(mod, '_artist_similarity_api_sync', return_value={"songs": []}), \ + patch.object(mod, '_lyrics_search_sync', return_value={"songs": []}), \ + patch.object(mod, '_database_genre_query_sync', side_effect=_db): + result = mod._ai_brainstorm_sync("best rock of the 90s", self._make_ai_config(), 50) + + ids = {s["item_id"] for s in result["songs"]} + assert "in" in ids + assert "out" not in ids # --------------------------------------------------------------------------- diff --git a/test/unit/test_memory_cleanup.py b/test/unit/test_memory_cleanup.py index e0613824..e791872b 100644 --- a/test/unit/test_memory_cleanup.py +++ b/test/unit/test_memory_cleanup.py @@ -143,7 +143,7 @@ class TestAnalyzeAlbumMemoryCleanup: @patch('tasks.analysis.get_tracks_from_album') @patch('tasks.analysis.download_track') @patch('tasks.analysis.analyze_track') - @patch('app_helper.get_db') + @patch('tasks.analysis_helper.get_db') @patch('tasks.analysis.ort') @patch('tasks.analysis.cleanup_onnx_session') @patch('tasks.memory_utils.cleanup_cuda_memory') @@ -167,10 +167,10 @@ def test_cleanup_on_database_error( ] mock_download.return_value = "/tmp/track.mp3" - # Mock database to raise error - mock_conn = MagicMock() - mock_get_db.return_value = mock_conn - mock_conn.cursor.side_effect = OperationalError("Connection failed") + # Mock the database connection used during album analysis (the first DB + # touch is analysis_helper.get_existing_track_ids) to fail like a real + # connection error, so the OperationalError propagates. + mock_get_db.side_effect = OperationalError("Connection failed") # Mock ONNX sessions mock_ort.get_available_providers.return_value = ['CPUExecutionProvider'] diff --git a/test/unit/test_memory_utils.py b/test/unit/test_memory_utils.py index f5ac88b2..19990bd7 100644 --- a/test/unit/test_memory_utils.py +++ b/test/unit/test_memory_utils.py @@ -6,9 +6,11 @@ import pytest from unittest.mock import Mock, MagicMock -from tasks.memory_utils import ( +from sanitization import ( sanitize_string_for_db, sanitize_json_for_db, +) +from tasks.memory_utils import ( cleanup_cuda_memory, cleanup_onnx_session, comprehensive_memory_cleanup, diff --git a/test/unit/test_provider_migration_blueprint.py b/test/unit/test_provider_migration_blueprint.py index fd3181ca..02404031 100644 --- a/test/unit/test_provider_migration_blueprint.py +++ b/test/unit/test_provider_migration_blueprint.py @@ -345,7 +345,7 @@ def test_happy_path_enqueues_job(self, bp_mod, client, fake_db): # --------------------------------------------------------------------------- # SSRF guard on the user-supplied media-server URL (_validate_probe_url -> -# app_helper.validate_outbound_url). Self-hosted servers live on the LAN / +# ssrf_guard.validate_outbound_url). Self-hosted servers live on the LAN / # loopback, so those are accepted; cloud-metadata, link-local, multicast, # unspecified and non-HTTP(S) schemes are rejected. IP literals are used so the # checks never depend on DNS resolution. diff --git a/test/unit/test_security_ssrf.py b/test/unit/test_security_ssrf.py index 619417bd..0c502c64 100644 --- a/test/unit/test_security_ssrf.py +++ b/test/unit/test_security_ssrf.py @@ -1,4 +1,4 @@ -"""Direct unit matrix for app_helper.validate_outbound_url (SSRF guard). +"""Direct unit matrix for ssrf_guard.validate_outbound_url (SSRF guard). The migration blueprint exercises this guard only indirectly. Here the allow/deny contract is pinned directly so it is a first-class, fast regression @@ -12,7 +12,7 @@ import pytest -from app_helper import validate_outbound_url +from ssrf_guard import validate_outbound_url def _addrinfo(ip, port=80): @@ -35,7 +35,7 @@ def test_rejects_non_http_schemes(self, url): @pytest.mark.parametrize('url', ['http://8.8.8.8', 'https://8.8.8.8']) def test_accepts_http_and_https(self, url): - with patch('app_helper.socket.getaddrinfo', return_value=_addrinfo('8.8.8.8')): + with patch('ssrf_guard.socket.getaddrinfo', return_value=_addrinfo('8.8.8.8')): assert validate_outbound_url(url) == (True, None) @@ -59,7 +59,7 @@ class TestValidateOutboundUrlIpClasses: ('http://127.0.0.1:8096', '127.0.0.1'), ]) def test_loopback_allowed(self, url, ip): - with patch('app_helper.socket.getaddrinfo', return_value=_addrinfo(ip)): + with patch('ssrf_guard.socket.getaddrinfo', return_value=_addrinfo(ip)): assert validate_outbound_url(url) == (True, None) @pytest.mark.parametrize('url,ip', [ @@ -68,7 +68,7 @@ def test_loopback_allowed(self, url, ip): ('http://192.168.1.50:8096', '192.168.1.50'), ]) def test_rfc1918_allowed(self, url, ip): - with patch('app_helper.socket.getaddrinfo', return_value=_addrinfo(ip)): + with patch('ssrf_guard.socket.getaddrinfo', return_value=_addrinfo(ip)): assert validate_outbound_url(url) == (True, None) @pytest.mark.parametrize('url,ip', [ @@ -76,7 +76,7 @@ def test_rfc1918_allowed(self, url, ip): ('http://1.2.3.4:8096', '1.2.3.4'), ]) def test_public_ip_allowed(self, url, ip): - with patch('app_helper.socket.getaddrinfo', return_value=_addrinfo(ip)): + with patch('ssrf_guard.socket.getaddrinfo', return_value=_addrinfo(ip)): assert validate_outbound_url(url) == (True, None) @pytest.mark.parametrize('url,ip', [ @@ -87,26 +87,26 @@ def test_public_ip_allowed(self, url, ip): ('http://240.0.0.1', '240.0.0.1'), ]) def test_dangerous_ip_classes_rejected(self, url, ip): - with patch('app_helper.socket.getaddrinfo', return_value=_addrinfo(ip)): + with patch('ssrf_guard.socket.getaddrinfo', return_value=_addrinfo(ip)): ok, reason = validate_outbound_url(url) assert ok is False assert reason == 'Target host resolves to a disallowed IP address' def test_dns_name_resolving_to_metadata_rejected(self): - with patch('app_helper.socket.getaddrinfo', return_value=_addrinfo('169.254.169.254')): + with patch('ssrf_guard.socket.getaddrinfo', return_value=_addrinfo('169.254.169.254')): ok, reason = validate_outbound_url('http://metadata.internal/') assert ok is False assert reason == 'Target host resolves to a disallowed IP address' def test_unresolvable_host_rejected(self): - with patch('app_helper.socket.getaddrinfo', side_effect=socket.gaierror('nope')): + with patch('ssrf_guard.socket.getaddrinfo', side_effect=socket.gaierror('nope')): ok, reason = validate_outbound_url('http://does-not-exist.invalid/') assert ok is False assert reason == 'Could not resolve host' def test_resolved_invalid_ip_rejected(self): bad = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('not-an-ip', 80))] - with patch('app_helper.socket.getaddrinfo', return_value=bad): + with patch('ssrf_guard.socket.getaddrinfo', return_value=bad): ok, reason = validate_outbound_url('http://weird.host/') assert ok is False assert reason == 'Resolved host to invalid IP' diff --git a/test/unit/test_string_sanitization.py b/test/unit/test_string_sanitization.py index f9943193..71515e97 100644 --- a/test/unit/test_string_sanitization.py +++ b/test/unit/test_string_sanitization.py @@ -10,7 +10,7 @@ class TestSaveTrackStringSanitization: """Test string sanitization in save_track_analysis_and_embedding.""" - @patch('app_helper.get_db') + @patch('database.get_db') def test_sanitize_removes_nul_bytes(self, mock_get_db): """Test that NUL bytes are removed from all string fields.""" from app_helper import save_track_analysis_and_embedding @@ -54,7 +54,7 @@ def test_sanitize_removes_nul_bytes(self, mock_get_db): assert values[2] == "ArtistName" assert values[9] == "AlbumName" - @patch('app_helper.get_db') + @patch('database.get_db') def test_sanitize_removes_control_characters(self, mock_get_db): """Test that control characters are removed.""" from app_helper import save_track_analysis_and_embedding @@ -79,7 +79,7 @@ def test_sanitize_removes_control_characters(self, mock_get_db): assert values[1] == "SongTitle" assert values[2] == "ArtistName" - @patch('app_helper.get_db') + @patch('database.get_db') def test_sanitize_handles_none_values(self, mock_get_db): """Test that None values are handled correctly.""" from app_helper import save_track_analysis_and_embedding @@ -105,7 +105,7 @@ def test_sanitize_handles_none_values(self, mock_get_db): assert values[5] is None # scale assert values[8] is None # other_features - @patch('app_helper.get_db') + @patch('database.get_db') def test_sanitize_truncates_long_strings(self, mock_get_db): """Test that overly long strings are truncated.""" from app_helper import save_track_analysis_and_embedding @@ -134,7 +134,7 @@ def test_sanitize_truncates_long_strings(self, mock_get_db): assert len(values[2]) == 200 # author assert len(values[8]) == 2000 # other_features - @patch('app_helper.get_db') + @patch('database.get_db') def test_sanitize_strips_whitespace(self, mock_get_db): """Test that leading/trailing whitespace is stripped.""" from app_helper import save_track_analysis_and_embedding @@ -158,7 +158,7 @@ def test_sanitize_strips_whitespace(self, mock_get_db): assert values[4] == "C" assert values[5] == "major" - @patch('app_helper.get_db') + @patch('database.get_db') def test_sanitize_preserves_unicode(self, mock_get_db): """Test that Unicode characters are preserved.""" from app_helper import save_track_analysis_and_embedding