imagefind/scripts/live-0516-search-acceptance.py

#!/usr/bin/env python3
"""Read-only live acceptance for audio search and transcript APIs."""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

import httpx
from zhconv import convert


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True)
    auth = parser.add_mutually_exclusive_group(required=True)
    auth.add_argument("--password-stdin", action="store_true")
    auth.add_argument("--token-file", type=Path)
    args = parser.parse_args()
    password = sys.stdin.readline().rstrip("\r\n") if args.password_stdin else ""
    headers = {}
    if args.token_file is not None:
        headers["Authorization"] = f"Bearer {args.token_file.read_text().strip()}"
    client = httpx.Client(
        base_url=args.base_url.rstrip("/"), headers=headers, timeout=httpx.Timeout(30, connect=10)
    )
    try:
        if args.password_stdin:
            login = client.post(
                "/api/v1/auth/login", json={"password": password, "remember_device": False}
            )
            login.raise_for_status()
            csrf = str(login.json().get("csrf_token") or "")
            if not csrf:
                raise AssertionError("login response omitted CSRF token")
            client.headers["X-CSRF-Token"] = csrf
            password = ""

        status = client.get("/api/v1/status")
        status.raise_for_status()
        if status.json().get("version") != "0.5.45":
            raise AssertionError(f"unexpected version: {status.json().get('version')}")

        coverage_response = client.get("/api/v1/search/coverage")
        coverage_response.raise_for_status()
        coverage = coverage_response.json()
        required = {
            "total",
            "indexed",
            "searchable",
            "empty",
            "queued",
            "running",
            "failed",
            "percent",
            "current_model_percent",
        }
        if not required <= coverage.keys():
            raise AssertionError(f"coverage fields missing: {sorted(required - coverage.keys())}")

        videos_response = client.get("/api/v1/videos?limit=200")
        videos_response.raise_for_status()
        videos = videos_response.json()
        selected = None
        transcript = None
        selected_item = None
        fallback = None
        for video in videos:
            response = client.get(f"/api/v1/videos/{video['id']}/transcript?page=1&page_size=30")
            response.raise_for_status()
            candidate = response.json()
            if candidate.get("items"):
                fallback = fallback or (video, candidate, candidate["items"][0])
                chinese_item = next(
                    (
                        item
                        for item in candidate["items"]
                        if re.search(r"[\u3400-\u9fff]{2,}", str(item.get("raw_text") or ""))
                    ),
                    None,
                )
                if chinese_item:
                    selected, transcript, selected_item = video, candidate, chinese_item
                    break
        if selected is None and fallback:
            selected, transcript, selected_item = fallback

        checks: dict[str, object] = {
            "version": status.json()["version"],
            "coverage": {key: coverage[key] for key in sorted(required)},
            "transcript_video_found": bool(selected),
        }
        resources_response = client.get("/api/v1/system/resources")
        resources_response.raise_for_status()
        resources = resources_response.json()
        diagnostics_response = client.get("/api/v1/system/diagnostics")
        diagnostics_response.raise_for_status()
        diagnostics = diagnostics_response.json()
        models_response = client.get("/api/v1/models")
        models_response.raise_for_status()
        models = models_response.json()
        checks["runtime"] = {
            "cpu_percent": resources.get("cpu_percent"),
            "memory_available_bytes": resources.get("memory_available_bytes"),
            "io": resources.get("io"),
            "ai_lane": (resources.get("lanes") or {}).get("ai"),
            "database": diagnostics.get("database"),
            "event_loop_lag_ms": diagnostics.get("event_loop_lag_ms"),
            "request_p95_ms": diagnostics.get("request_p95_ms"),
            "audio_accelerator": ((models.get("accelerator") or {}).get("components") or {}).get("audio"),
            "audio_health": (models.get("component_health") or {}).get("audio"),
        }
        if selected and transcript:
            first = selected_item or transcript["items"][0]
            cjk_runs = re.findall(r"[\u3400-\u9fff]{2,}", str(first.get("raw_text") or ""))
            if cjk_runs:
                query = cjk_runs[0][: min(6, len(cjk_runs[0]))]
                variants = list(dict.fromkeys((query, convert(query, "zh-cn"), convert(query, "zh-tw"))))
                results = []
                for variant in variants:
                    audio = client.post(
                        "/api/v1/search",
                        json={"text": variant, "recognition_types": ["audio"], "limit": 50},
                    )
                    audio.raise_for_status()
                    audio_ids = {item.get("video_id") for item in audio.json().get("items", [])}
                    combined = client.post("/api/v1/search", json={"text": variant, "limit": 50})
                    combined.raise_for_status()
                    combined_ids = {item.get("video_id") for item in combined.json().get("items", [])}
                    results.append(
                        {
                            "variant_changed": variant != query,
                            "audio_match": selected["id"] in audio_ids,
                            "combined_match": selected["id"] in combined_ids,
                        }
                    )
                if not all(item["audio_match"] and item["combined_match"] for item in results):
                    raise AssertionError(f"live Chinese audio search mismatch: {results}")
                checks["chinese_query_length"] = len(query)
                checks["variant_checks"] = results
            checks["transcript"] = {
                "status": transcript.get("status"),
                "total": transcript.get("total"),
                "pages": transcript.get("pages"),
                "items_returned": len(transcript.get("items", [])),
            }
        print(json.dumps(checks, ensure_ascii=False, indent=2), flush=True)
    finally:
        password = ""
        try:
            if args.password_stdin:
                client.post("/api/v1/auth/logout")
        finally:
            client.close()


if __name__ == "__main__":
    main()