mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-14 22:11:27 +00:00
Merge pull request #9 from MrSleeps/feature/search-http-api
Add HTTP API for Jellyfin scene search
This commit is contained in:
@@ -41,3 +41,10 @@ JELLYFIN_SPECIALIST_TAGS=
|
||||
JELLYFIN_FORCE_SPECIALIST_TAG=scene-search:specialist
|
||||
JELLYFIN_FORCE_GENERAL_TAG=scene-search:general
|
||||
JELLYFIN_SKIP_TAG=scene-search:skip
|
||||
|
||||
# HTTP search API used by the Jellyfin plugin
|
||||
SCENE_API_HOST=127.0.0.1
|
||||
SCENE_API_PORT=8097
|
||||
# Optional but recommended. When set, /api/search and thumbnails require
|
||||
# Authorization: Bearer <key> or X-Scene-Search-Key: <key>.
|
||||
SCENE_API_KEY=
|
||||
|
||||
@@ -272,6 +272,56 @@ python -m scene_search.cli search "people walking down a street"
|
||||
|
||||
Search output shows the stored caption as well as the Meilisearch and SigLIP diagnostic scores.
|
||||
|
||||
## HTTP search API
|
||||
|
||||
The Jellyfin plugin will talk to the same hybrid search logic through a small local FastAPI service.
|
||||
|
||||
Configure it in `.env`:
|
||||
|
||||
```env
|
||||
SCENE_API_HOST=127.0.0.1
|
||||
SCENE_API_PORT=8097
|
||||
SCENE_API_KEY=replace-with-a-random-secret
|
||||
```
|
||||
|
||||
If Jellyfin and this project run on the same host, keeping the API bound to `127.0.0.1` is recommended.
|
||||
|
||||
Start the service:
|
||||
|
||||
```bash
|
||||
python -m scene_search.cli serve
|
||||
```
|
||||
|
||||
The health endpoint is deliberately cheap and does not load SigLIP:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8097/api/health
|
||||
```
|
||||
|
||||
The first real search loads SigLIP and keeps it resident for later requests:
|
||||
|
||||
```bash
|
||||
curl -H "Authorization: Bearer replace-with-a-random-secret" \
|
||||
"http://127.0.0.1:8097/api/search?q=woman%20standing%20outside"
|
||||
```
|
||||
|
||||
`/api/search` searches Jellyfin-backed documents only and returns two views of the same ranking:
|
||||
|
||||
- `items`: one de-duplicated result per Jellyfin item, intended for Jellyfin's normal search provider.
|
||||
- `scenes`: timestamped scene hits with captions and scores, intended for the later Scene Search page.
|
||||
|
||||
Each scene contains its Jellyfin item/media-source/library IDs, item type, timestamp, caption, tags, Meilisearch ranking score and SigLIP diagnostic score.
|
||||
|
||||
Cached Trickplay thumbnails are available through:
|
||||
|
||||
```text
|
||||
GET /api/frames/{frame_id}/thumbnail
|
||||
```
|
||||
|
||||
When `SCENE_API_KEY` is set, search and thumbnail endpoints accept either `Authorization: Bearer <key>` or `X-Scene-Search-Key: <key>`. `/api/health` remains unauthenticated for local service monitoring.
|
||||
|
||||
The API intentionally runs with one worker. Multiple workers would each load their own copy of the embedding model into GPU memory.
|
||||
|
||||
## Removing indexed sources
|
||||
|
||||
Old prototype clips and Jellyfin-backed items are stored with different source kinds.
|
||||
|
||||
@@ -5,3 +5,5 @@ python-dotenv>=1.0
|
||||
requests>=2.32
|
||||
torch>=2.4
|
||||
transformers>=4.57
|
||||
fastapi>=0.115
|
||||
uvicorn>=0.30
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hmac
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, Query
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from .config import settings
|
||||
from .database import connect, init_db
|
||||
from .search_service import SearchService
|
||||
|
||||
app = FastAPI(
|
||||
title="Jellyfin Scene Search API",
|
||||
version="1.0.0",
|
||||
)
|
||||
|
||||
_service: SearchService | None = None
|
||||
|
||||
|
||||
def get_search_service() -> SearchService:
|
||||
global _service
|
||||
if _service is None:
|
||||
_service = SearchService(
|
||||
db_path=settings.db_path,
|
||||
meili_url=settings.meili_url,
|
||||
meili_master_key=settings.meili_master_key,
|
||||
meili_index=settings.meili_index,
|
||||
model_id=settings.model_id,
|
||||
device=settings.device,
|
||||
semantic_ratio=settings.hybrid_semantic_ratio,
|
||||
)
|
||||
return _service
|
||||
|
||||
|
||||
def require_api_key(
|
||||
authorization: str | None = Header(default=None),
|
||||
x_scene_search_key: str | None = Header(default=None),
|
||||
) -> None:
|
||||
expected = settings.scene_api_key
|
||||
if not expected:
|
||||
return
|
||||
|
||||
supplied = x_scene_search_key
|
||||
if authorization and authorization.lower().startswith("bearer "):
|
||||
supplied = authorization[7:].strip()
|
||||
|
||||
if not supplied or not hmac.compare_digest(supplied, expected):
|
||||
raise HTTPException(status_code=401, detail="Invalid API key")
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
def health() -> dict:
|
||||
return get_search_service().health()
|
||||
|
||||
|
||||
@app.get("/api/search", dependencies=[Depends(require_api_key)])
|
||||
def search(
|
||||
q: str = Query(min_length=1, max_length=500),
|
||||
scene_limit: int = Query(default=50, ge=1, le=100),
|
||||
item_limit: int = Query(default=20, ge=1, le=50),
|
||||
) -> dict:
|
||||
try:
|
||||
return get_search_service().search(
|
||||
q,
|
||||
scene_limit=scene_limit,
|
||||
item_limit=item_limit,
|
||||
jellyfin_only=True,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=503, detail=str(exc)) from exc
|
||||
|
||||
|
||||
@app.get(
|
||||
"/api/frames/{frame_id}/thumbnail",
|
||||
dependencies=[Depends(require_api_key)],
|
||||
)
|
||||
def frame_thumbnail(frame_id: str) -> FileResponse:
|
||||
init_db(settings.db_path)
|
||||
with connect(settings.db_path) as con:
|
||||
row = con.execute(
|
||||
"""
|
||||
SELECT f.thumb_path
|
||||
FROM frames f
|
||||
JOIN media m ON m.id = f.media_id
|
||||
WHERE f.id = ? AND m.source_kind = 'jellyfin'
|
||||
""",
|
||||
(frame_id,),
|
||||
).fetchone()
|
||||
|
||||
if row is None:
|
||||
raise HTTPException(status_code=404, detail="Frame not found")
|
||||
|
||||
path = Path(str(row["thumb_path"]))
|
||||
if not path.is_file():
|
||||
raise HTTPException(status_code=404, detail="Thumbnail file not found")
|
||||
|
||||
return FileResponse(path, media_type="image/jpeg")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(
|
||||
"scene_search.api:app",
|
||||
host=settings.scene_api_host,
|
||||
port=settings.scene_api_port,
|
||||
workers=1,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+42
-73
@@ -29,6 +29,7 @@ from .jellyfin_index import (
|
||||
from .maintenance import list_media_by_source, remove_media_by_source
|
||||
from .meili import Meili
|
||||
from .routing import ExplicitContentRouter, apply_context_routing
|
||||
from .search_service import SearchService
|
||||
|
||||
|
||||
def format_ts(ms: int) -> str:
|
||||
@@ -650,88 +651,53 @@ def cmd_status(_: argparse.Namespace) -> None:
|
||||
print(f"Indexing: {stats.get('isIndexing', 'unknown')}")
|
||||
|
||||
def cmd_search(args: argparse.Namespace) -> None:
|
||||
init_db(settings.db_path)
|
||||
embedder = build_embedder()
|
||||
vector = embedder.embed_text(args.query)
|
||||
meili = Meili(settings.meili_url, settings.meili_master_key, settings.meili_index)
|
||||
result = meili.search_vector(
|
||||
vector,
|
||||
args.query,
|
||||
args.limit,
|
||||
settings.hybrid_semantic_ratio,
|
||||
service = SearchService(
|
||||
db_path=settings.db_path,
|
||||
meili_url=settings.meili_url,
|
||||
meili_master_key=settings.meili_master_key,
|
||||
meili_index=settings.meili_index,
|
||||
model_id=settings.model_id,
|
||||
device=settings.device,
|
||||
semantic_ratio=settings.hybrid_semantic_ratio,
|
||||
)
|
||||
|
||||
hits = result.get("hits", [])
|
||||
if not hits:
|
||||
result = service.search(
|
||||
args.query,
|
||||
scene_limit=args.limit,
|
||||
item_limit=args.limit,
|
||||
jellyfin_only=False,
|
||||
)
|
||||
scenes = result["scenes"]
|
||||
if not scenes:
|
||||
print("No hits")
|
||||
return
|
||||
|
||||
frame_ids = [hit.get("id") for hit in hits if hit.get("id")]
|
||||
frame_vectors: dict[str, object] = {}
|
||||
if frame_ids:
|
||||
placeholders = ",".join("?" for _ in frame_ids)
|
||||
with connect(settings.db_path) as con:
|
||||
rows = con.execute(
|
||||
f"""
|
||||
SELECT f.id, f.embedding_f16, m.embedding_dim
|
||||
FROM frames f
|
||||
JOIN media m ON m.id = f.media_id
|
||||
WHERE f.id IN ({placeholders})
|
||||
""",
|
||||
frame_ids,
|
||||
).fetchall()
|
||||
frame_vectors = {
|
||||
row["id"]: unpack_vector(row["embedding_f16"], row["embedding_dim"])
|
||||
for row in rows
|
||||
}
|
||||
|
||||
scored_vectors = [frame_vectors.get(hit.get("id")) for hit in hits]
|
||||
if all(item is not None for item in scored_vectors):
|
||||
import numpy as np
|
||||
|
||||
siglip_scores = embedder.match_probabilities(
|
||||
vector,
|
||||
np.stack(scored_vectors), # type: ignore[arg-type]
|
||||
)
|
||||
else:
|
||||
siglip_scores = [None] * len(hits)
|
||||
|
||||
print(f"\nQuery: {args.query!r}\n")
|
||||
for i, (hit, siglip_score) in enumerate(zip(hits, siglip_scores, strict=True), 1):
|
||||
meili_score = hit.get("_rankingScore")
|
||||
meili_text = (
|
||||
f" meili={meili_score:.4f}"
|
||||
if isinstance(meili_score, (float, int))
|
||||
else ""
|
||||
)
|
||||
siglip_text = (
|
||||
f" siglip={float(siglip_score):.4f}"
|
||||
if siglip_score is not None
|
||||
else ""
|
||||
)
|
||||
for i, scene in enumerate(scenes, 1):
|
||||
meili_score = scene.get("ranking_score")
|
||||
siglip_score = scene.get("siglip_score")
|
||||
meili_text = f" meili={meili_score:.4f}" if meili_score is not None else ""
|
||||
siglip_text = f" siglip={siglip_score:.4f}" if siglip_score is not None else ""
|
||||
print(
|
||||
f"{i:2d}. {hit.get('title')} @ "
|
||||
f"{format_ts(int(hit.get('timestamp_ms', 0)))}"
|
||||
f"{i:2d}. {scene.get('title')} @ "
|
||||
f"{format_ts(int(scene.get('timestamp_ms') or 0))}"
|
||||
f"{meili_text}{siglip_text}"
|
||||
)
|
||||
caption = hit.get("caption")
|
||||
if caption:
|
||||
print(f" {caption}")
|
||||
route = hit.get("caption_route")
|
||||
route_score = hit.get("route_score")
|
||||
if route:
|
||||
score_text = (
|
||||
f" score={float(route_score):.3f}"
|
||||
if isinstance(route_score, (float, int))
|
||||
else ""
|
||||
)
|
||||
print(f" route={route}{score_text}")
|
||||
jellyfin_item_id = hit.get("jellyfin_item_id")
|
||||
if jellyfin_item_id:
|
||||
tags = hit.get("jellyfin_tags") or []
|
||||
if scene.get("caption"):
|
||||
print(f" {scene['caption']}")
|
||||
if scene.get("caption_route"):
|
||||
route_score = scene.get("route_score")
|
||||
score_text = f" score={float(route_score):.3f}" if route_score is not None else ""
|
||||
print(f" route={scene['caption_route']}{score_text}")
|
||||
if scene.get("jellyfin_item_id"):
|
||||
tags = scene.get("tags") or []
|
||||
tags_text = f" tags={', '.join(tags)}" if tags else ""
|
||||
print(f" jellyfin={jellyfin_item_id}{tags_text}")
|
||||
print(f" {hit.get('thumb_path')}")
|
||||
print(f" jellyfin={scene['jellyfin_item_id']}{tags_text}")
|
||||
|
||||
|
||||
def cmd_serve(_: argparse.Namespace) -> None:
|
||||
from .api import main as api_main
|
||||
|
||||
api_main()
|
||||
|
||||
|
||||
def make_parser() -> argparse.ArgumentParser:
|
||||
@@ -799,6 +765,9 @@ def make_parser() -> argparse.ArgumentParser:
|
||||
status = sub.add_parser("status", help="Show canonical SQLite and Meilisearch index status")
|
||||
status.set_defaults(func=cmd_status)
|
||||
|
||||
serve = sub.add_parser("serve", help="Run the local HTTP search API")
|
||||
serve.set_defaults(func=cmd_serve)
|
||||
|
||||
search = sub.add_parser("search", help="Hybrid caption + visual scene search")
|
||||
search.add_argument("query")
|
||||
search.add_argument("--limit", type=int, default=10)
|
||||
|
||||
@@ -96,6 +96,10 @@ class Settings:
|
||||
"JELLYFIN_TAG_CONTEXT", "true"
|
||||
).strip().lower() not in {"0", "false", "no", "off"}
|
||||
|
||||
scene_api_host: str = os.getenv("SCENE_API_HOST", "127.0.0.1")
|
||||
scene_api_port: int = int(os.getenv("SCENE_API_PORT", "8097"))
|
||||
scene_api_key: str = os.getenv("SCENE_API_KEY", "")
|
||||
|
||||
data_dir: Path = Path(os.getenv("DATA_DIR", "./data")).resolve()
|
||||
|
||||
@property
|
||||
|
||||
@@ -87,6 +87,7 @@ class Meili:
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
semantic_ratio: float = 0.8,
|
||||
filter_expression: str | None = None,
|
||||
) -> dict:
|
||||
payload = {
|
||||
"q": query,
|
||||
@@ -98,6 +99,8 @@ class Meili:
|
||||
"limit": limit,
|
||||
"showRankingScore": True,
|
||||
}
|
||||
if filter_expression:
|
||||
payload["filter"] = filter_expression
|
||||
return self._check(self.session.post(
|
||||
f"{self.url}/indexes/{self.index}/search",
|
||||
json=payload,
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .database import connect, init_db, unpack_vector
|
||||
from .embedder import SiglipEmbedder
|
||||
from .meili import Meili
|
||||
|
||||
|
||||
class SearchService:
|
||||
"""Shared hybrid search implementation for CLI and HTTP consumers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
db_path: Path,
|
||||
meili_url: str,
|
||||
meili_master_key: str,
|
||||
meili_index: str,
|
||||
model_id: str,
|
||||
device: str,
|
||||
semantic_ratio: float,
|
||||
) -> None:
|
||||
self.db_path = db_path
|
||||
self.meili = Meili(meili_url, meili_master_key, meili_index)
|
||||
self.model_id = model_id
|
||||
self.device = device
|
||||
self.semantic_ratio = semantic_ratio
|
||||
|
||||
self._embedder: SiglipEmbedder | None = None
|
||||
self._model_lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def model_loaded(self) -> bool:
|
||||
return self._embedder is not None
|
||||
|
||||
def _get_embedder(self) -> SiglipEmbedder:
|
||||
if self._embedder is None:
|
||||
self._embedder = SiglipEmbedder(self.model_id, self.device)
|
||||
return self._embedder
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
init_db(self.db_path)
|
||||
with connect(self.db_path) as con:
|
||||
media_count = int(con.execute("SELECT COUNT(*) FROM media").fetchone()[0])
|
||||
frame_count = int(con.execute("SELECT COUNT(*) FROM frames").fetchone()[0])
|
||||
jellyfin_media_count = int(
|
||||
con.execute(
|
||||
"SELECT COUNT(*) FROM media WHERE source_kind = 'jellyfin'"
|
||||
).fetchone()[0]
|
||||
)
|
||||
|
||||
meili_error: str | None = None
|
||||
meili_documents: int | None = None
|
||||
try:
|
||||
stats = self.meili.stats()
|
||||
raw_count = stats.get("numberOfDocuments")
|
||||
if isinstance(raw_count, int):
|
||||
meili_documents = raw_count
|
||||
except Exception as exc:
|
||||
meili_error = str(exc)
|
||||
|
||||
return {
|
||||
"status": "ok" if meili_error is None else "degraded",
|
||||
"model_loaded": self.model_loaded,
|
||||
"media_count": media_count,
|
||||
"jellyfin_media_count": jellyfin_media_count,
|
||||
"frame_count": frame_count,
|
||||
"meili_documents": meili_documents,
|
||||
"meili_error": meili_error,
|
||||
}
|
||||
|
||||
def _vectors_for_hits(self, frame_ids: list[str]) -> dict[str, np.ndarray]:
|
||||
if not frame_ids:
|
||||
return {}
|
||||
|
||||
placeholders = ",".join("?" for _ in frame_ids)
|
||||
with connect(self.db_path) as con:
|
||||
rows = con.execute(
|
||||
f"""
|
||||
SELECT f.id, f.embedding_f16, m.embedding_dim
|
||||
FROM frames f
|
||||
JOIN media m ON m.id = f.media_id
|
||||
WHERE f.id IN ({placeholders})
|
||||
""",
|
||||
frame_ids,
|
||||
).fetchall()
|
||||
|
||||
return {
|
||||
str(row["id"]): unpack_vector(
|
||||
row["embedding_f16"],
|
||||
int(row["embedding_dim"]),
|
||||
)
|
||||
for row in rows
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _scene_from_hit(
|
||||
hit: dict[str, Any],
|
||||
siglip_score: float | None,
|
||||
) -> dict[str, Any]:
|
||||
ranking_score = hit.get("_rankingScore")
|
||||
return {
|
||||
"frame_id": hit.get("id"),
|
||||
"media_id": hit.get("media_id"),
|
||||
"jellyfin_item_id": hit.get("jellyfin_item_id"),
|
||||
"jellyfin_media_source_id": hit.get("jellyfin_media_source_id"),
|
||||
"jellyfin_library_id": hit.get("jellyfin_library_id"),
|
||||
"jellyfin_item_type": hit.get("jellyfin_item_type"),
|
||||
"title": hit.get("title"),
|
||||
"timestamp_ms": int(hit.get("timestamp_ms") or 0),
|
||||
"caption": hit.get("caption"),
|
||||
"caption_model": hit.get("caption_model"),
|
||||
"caption_route": hit.get("caption_route"),
|
||||
"route_score": hit.get("route_score"),
|
||||
"tags": hit.get("jellyfin_tags") or [],
|
||||
"ranking_score": (
|
||||
float(ranking_score)
|
||||
if isinstance(ranking_score, (float, int))
|
||||
else None
|
||||
),
|
||||
"siglip_score": siglip_score,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _group_items(
|
||||
scenes: list[dict[str, Any]],
|
||||
item_limit: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
grouped: dict[str, dict[str, Any]] = {}
|
||||
|
||||
for scene in scenes:
|
||||
item_id = scene.get("jellyfin_item_id")
|
||||
if not item_id:
|
||||
continue
|
||||
|
||||
existing = grouped.get(str(item_id))
|
||||
if existing is None:
|
||||
grouped[str(item_id)] = {
|
||||
"jellyfin_item_id": item_id,
|
||||
"jellyfin_library_id": scene.get("jellyfin_library_id"),
|
||||
"jellyfin_item_type": scene.get("jellyfin_item_type"),
|
||||
"title": scene.get("title"),
|
||||
"tags": scene.get("tags") or [],
|
||||
"best_timestamp_ms": scene.get("timestamp_ms"),
|
||||
"best_caption": scene.get("caption"),
|
||||
"ranking_score": scene.get("ranking_score"),
|
||||
"siglip_score": scene.get("siglip_score"),
|
||||
"scene_matches": 1,
|
||||
}
|
||||
else:
|
||||
existing["scene_matches"] = int(existing["scene_matches"]) + 1
|
||||
|
||||
return list(grouped.values())[:item_limit]
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
scene_limit: int = 50,
|
||||
item_limit: int = 20,
|
||||
jellyfin_only: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
query = query.strip()
|
||||
if not query:
|
||||
raise ValueError("Search query cannot be empty")
|
||||
|
||||
init_db(self.db_path)
|
||||
|
||||
# SigLIP / PyTorch model execution is kept behind one lock so concurrent
|
||||
# API requests do not fight over a single GPU model instance.
|
||||
with self._model_lock:
|
||||
embedder = self._get_embedder()
|
||||
text_vector = embedder.embed_text(query)
|
||||
|
||||
result = self.meili.search_vector(
|
||||
text_vector,
|
||||
query,
|
||||
scene_limit,
|
||||
self.semantic_ratio,
|
||||
filter_expression=(
|
||||
'source_kind = "jellyfin"'
|
||||
if jellyfin_only
|
||||
else None
|
||||
),
|
||||
)
|
||||
hits = list(result.get("hits") or [])
|
||||
|
||||
frame_ids = [
|
||||
str(hit["id"])
|
||||
for hit in hits
|
||||
if hit.get("id") is not None
|
||||
]
|
||||
vectors = self._vectors_for_hits(frame_ids)
|
||||
|
||||
scored_vectors = [vectors.get(str(hit.get("id"))) for hit in hits]
|
||||
if hits and all(vector is not None for vector in scored_vectors):
|
||||
siglip_scores = embedder.match_probabilities(
|
||||
text_vector,
|
||||
np.stack(scored_vectors), # type: ignore[arg-type]
|
||||
)
|
||||
scores: list[float | None] = [
|
||||
float(score) for score in siglip_scores
|
||||
]
|
||||
else:
|
||||
scores = [None] * len(hits)
|
||||
|
||||
scenes = [
|
||||
self._scene_from_hit(hit, score)
|
||||
for hit, score in zip(hits, scores, strict=True)
|
||||
]
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"semantic_ratio": self.semantic_ratio,
|
||||
"scene_count": len(scenes),
|
||||
"item_count": min(
|
||||
len({
|
||||
scene.get("jellyfin_item_id")
|
||||
for scene in scenes
|
||||
if scene.get("jellyfin_item_id")
|
||||
}),
|
||||
item_limit,
|
||||
),
|
||||
"items": self._group_items(scenes, item_limit),
|
||||
"scenes": scenes,
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from scene_search.api import require_api_key
|
||||
|
||||
|
||||
class ApiAuthTests(unittest.TestCase):
|
||||
def test_api_key_disabled_accepts_request(self) -> None:
|
||||
fake_settings = SimpleNamespace(scene_api_key="")
|
||||
with patch("scene_search.api.settings", fake_settings):
|
||||
require_api_key(authorization=None, x_scene_search_key=None)
|
||||
|
||||
def test_bearer_api_key_is_accepted(self) -> None:
|
||||
fake_settings = SimpleNamespace(scene_api_key="secret")
|
||||
with patch("scene_search.api.settings", fake_settings):
|
||||
require_api_key(
|
||||
authorization="Bearer secret",
|
||||
x_scene_search_key=None,
|
||||
)
|
||||
|
||||
def test_header_api_key_is_accepted(self) -> None:
|
||||
fake_settings = SimpleNamespace(scene_api_key="secret")
|
||||
with patch("scene_search.api.settings", fake_settings):
|
||||
require_api_key(
|
||||
authorization=None,
|
||||
x_scene_search_key="secret",
|
||||
)
|
||||
|
||||
def test_wrong_api_key_is_rejected(self) -> None:
|
||||
fake_settings = SimpleNamespace(scene_api_key="secret")
|
||||
with patch("scene_search.api.settings", fake_settings):
|
||||
with self.assertRaises(HTTPException) as caught:
|
||||
require_api_key(
|
||||
authorization="Bearer wrong",
|
||||
x_scene_search_key=None,
|
||||
)
|
||||
|
||||
self.assertEqual(caught.exception.status_code, 401)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -33,6 +33,22 @@ class MeiliVectorSerializationTests(unittest.TestCase):
|
||||
"http://localhost:7700/indexes/scenes/documents"
|
||||
)
|
||||
|
||||
def test_search_vector_adds_optional_filter(self) -> None:
|
||||
meili = Meili("http://localhost:7700", "key", "scenes")
|
||||
response = Mock()
|
||||
response.ok = True
|
||||
response.json.return_value = {"hits": []}
|
||||
meili.session.post = Mock(return_value=response)
|
||||
|
||||
meili.search_vector(
|
||||
[0.1, 0.2],
|
||||
"query",
|
||||
filter_expression='source_kind = "jellyfin"',
|
||||
)
|
||||
|
||||
payload = meili.session.post.call_args.kwargs["json"]
|
||||
self.assertEqual(payload["filter"], 'source_kind = "jellyfin"')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock
|
||||
|
||||
import numpy as np
|
||||
|
||||
from scene_search.database import connect, init_db, pack_vector
|
||||
from scene_search.search_service import SearchService
|
||||
|
||||
|
||||
class FakeEmbedder:
|
||||
def embed_text(self, text: str) -> np.ndarray:
|
||||
self.last_text = text
|
||||
return np.asarray([1.0, 0.0], dtype=np.float32)
|
||||
|
||||
def match_probabilities(
|
||||
self,
|
||||
text_vector: np.ndarray,
|
||||
image_vectors: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
return np.asarray([0.9, 0.7, 0.6], dtype=np.float32)
|
||||
|
||||
|
||||
class SearchServiceTests(unittest.TestCase):
|
||||
def test_search_returns_scenes_and_grouped_jellyfin_items(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = Path(tmp) / "scenes.sqlite3"
|
||||
init_db(db_path)
|
||||
|
||||
with connect(db_path) as con:
|
||||
con.execute(
|
||||
"""
|
||||
INSERT INTO media (
|
||||
id, source_path, title, content_sha256, size_bytes, mtime_ns,
|
||||
interval_ms, model_id, embedding_dim, source_kind,
|
||||
jellyfin_item_id, jellyfin_library_id, jellyfin_item_type,
|
||||
jellyfin_tags
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
"media-1",
|
||||
"/media/example.mkv",
|
||||
"Example Movie",
|
||||
"hash",
|
||||
1,
|
||||
0,
|
||||
10000,
|
||||
"model",
|
||||
2,
|
||||
"jellyfin",
|
||||
"jf-1",
|
||||
"lib-1",
|
||||
"Movie",
|
||||
'["tag one"]',
|
||||
),
|
||||
)
|
||||
for number, timestamp in enumerate((0, 10000, 20000)):
|
||||
con.execute(
|
||||
"""
|
||||
INSERT INTO frames (
|
||||
id, media_id, frame_number, timestamp_ms,
|
||||
thumb_path, embedding_f16, caption
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
f"frame-{number}",
|
||||
"media-1",
|
||||
number,
|
||||
timestamp,
|
||||
f"/tmp/frame-{number}.jpg",
|
||||
pack_vector(np.asarray([1.0, 0.0], dtype=np.float32)),
|
||||
f"caption {number}",
|
||||
),
|
||||
)
|
||||
|
||||
service = SearchService(
|
||||
db_path=db_path,
|
||||
meili_url="http://meili.test",
|
||||
meili_master_key="key",
|
||||
meili_index="scenes",
|
||||
model_id="model",
|
||||
device="cpu",
|
||||
semantic_ratio=0.8,
|
||||
)
|
||||
service._embedder = FakeEmbedder() # type: ignore[assignment]
|
||||
service.meili = Mock()
|
||||
service.meili.search_vector.return_value = {
|
||||
"hits": [
|
||||
{
|
||||
"id": "frame-0",
|
||||
"media_id": "media-1",
|
||||
"jellyfin_item_id": "jf-1",
|
||||
"jellyfin_library_id": "lib-1",
|
||||
"jellyfin_item_type": "Movie",
|
||||
"title": "Example Movie",
|
||||
"timestamp_ms": 0,
|
||||
"caption": "caption 0",
|
||||
"jellyfin_tags": ["tag one"],
|
||||
"_rankingScore": 0.95,
|
||||
},
|
||||
{
|
||||
"id": "frame-1",
|
||||
"media_id": "media-1",
|
||||
"jellyfin_item_id": "jf-1",
|
||||
"jellyfin_library_id": "lib-1",
|
||||
"jellyfin_item_type": "Movie",
|
||||
"title": "Example Movie",
|
||||
"timestamp_ms": 10000,
|
||||
"caption": "caption 1",
|
||||
"jellyfin_tags": ["tag one"],
|
||||
"_rankingScore": 0.85,
|
||||
},
|
||||
{
|
||||
"id": "frame-2",
|
||||
"media_id": "media-1",
|
||||
"jellyfin_item_id": "jf-1",
|
||||
"jellyfin_library_id": "lib-1",
|
||||
"jellyfin_item_type": "Movie",
|
||||
"title": "Example Movie",
|
||||
"timestamp_ms": 20000,
|
||||
"caption": "caption 2",
|
||||
"jellyfin_tags": ["tag one"],
|
||||
"_rankingScore": 0.75,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
result = service.search("person outside", scene_limit=10, item_limit=5)
|
||||
|
||||
self.assertEqual(result["scene_count"], 3)
|
||||
self.assertEqual(result["item_count"], 1)
|
||||
self.assertEqual(len(result["items"]), 1)
|
||||
self.assertEqual(result["items"][0]["jellyfin_item_id"], "jf-1")
|
||||
self.assertEqual(result["items"][0]["best_timestamp_ms"], 0)
|
||||
self.assertEqual(result["items"][0]["scene_matches"], 3)
|
||||
self.assertAlmostEqual(result["scenes"][0]["siglip_score"], 0.9, places=5)
|
||||
|
||||
call = service.meili.search_vector.call_args
|
||||
self.assertEqual(call.args[1], "person outside")
|
||||
self.assertEqual(
|
||||
call.kwargs["filter_expression"],
|
||||
'source_kind = "jellyfin"',
|
||||
)
|
||||
|
||||
def test_health_does_not_load_embedding_model(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = Path(tmp) / "scenes.sqlite3"
|
||||
service = SearchService(
|
||||
db_path=db_path,
|
||||
meili_url="http://meili.test",
|
||||
meili_master_key="key",
|
||||
meili_index="scenes",
|
||||
model_id="model",
|
||||
device="cpu",
|
||||
semantic_ratio=0.8,
|
||||
)
|
||||
service.meili = Mock()
|
||||
service.meili.stats.return_value = {"numberOfDocuments": 0}
|
||||
|
||||
result = service.health()
|
||||
|
||||
self.assertFalse(result["model_loaded"])
|
||||
self.assertEqual(result["status"], "ok")
|
||||
self.assertEqual(result["meili_documents"], 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user