Search an approved image library with an image query

image-search-pilot/
  app.py
  cli.py
  search_core.py
  manifest.template.json
  requirements.txt
  test_app.py
  test_search_core.py
  manifest.json              # your private inventory
  images/                    # your approved image files
  private-search.sqlite      # created locally at indexing

"""Small, single-operator image search pilot. Only Embedder calls a provider."""
from __future__ import annotations

from datetime import datetime, timezone
from hashlib import sha256
from io import BytesIO
import hmac
import json
import math
import os
from pathlib import Path
import re
import sqlite3
import warnings
from PIL import Image

MODEL = "gemini-embedding-2"
DIM = 768
MAX_IMAGE = 5 * 1024 * 1024
MAX_ITEMS = 50
MAX_CALLS_PER_UTC_DAY = 100
ID = re.compile(r"[a-z0-9][a-z0-9_-]{0,63}\Z")


def now():
    return datetime.now(timezone.utc).isoformat()


def image_type(data: bytes) -> str:
    if not data or len(data) > MAX_IMAGE:
        raise ValueError("image size")
    mime = "image/png" if data.startswith(b"\x89PNG\r\n\x1a\n") else (
        "image/jpeg" if data.startswith(b"\xff\xd8\xff") and data.endswith(b"\xff\xd9") else None)
    if not mime:
        raise ValueError("image type")
    Image.MAX_IMAGE_PIXELS = 20_000_000
    try:
        with warnings.catch_warnings():
            warnings.simplefilter("error", Image.DecompressionBombWarning)
            with Image.open(BytesIO(data), formats=["PNG", "JPEG"]) as image:
                if image.format != mime[6:].upper() or image.width * image.height > 20_000_000:
                    raise ValueError("format/pixels")
                image.load()
    except (OSError, Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
        raise ValueError("decode") from exc
    return mime


def unit_vector(values):
    if len(values) != DIM or not all(isinstance(x, (int, float)) and math.isfinite(x) for x in values):
        raise ValueError("embedding dimension or finite-value check failed")
    length = math.sqrt(sum(x * x for x in values))
    if length == 0:
        raise ValueError("zero embedding")
    return [float(x) / length for x in values]


def load_manifest(path: Path, library: Path):
    rows = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(rows, list) or not 1 <= len(rows) <= MAX_ITEMS:
        raise ValueError("manifest needs 1 to 50 images")
    library = library.resolve(strict=True)
    found = []
    seen = set()
    for row in rows:
        if not isinstance(row, dict) or set(row) != {"id", "file", "title", "tags", "permissionRef", "revision"}:
            raise ValueError("manifest row fields invalid")
        ident, filename = row["id"], row["file"]
        if not isinstance(ident, str) or not ID.fullmatch(ident) or ident in seen:
            raise ValueError("invalid or duplicate id")
        seen.add(ident)
        if not isinstance(filename, str) or not filename or Path(filename).name != filename:
            raise ValueError("file must be a direct filename")
        filepath = library / filename
        if filepath.is_symlink() or not filepath.is_file() or filepath.resolve().parent != library:
            raise ValueError("image must be a regular file directly in library")
        title = row["title"]
        tags = row["tags"]
        if not isinstance(title, str) or not 1 <= len(title) <= 120:
            raise ValueError("title invalid")
        if not isinstance(tags, list) or len(tags) > 10 or len(set(map(str, tags))) != len(tags) or any(not isinstance(t, str) or not ID.fullmatch(t) for t in tags):
            raise ValueError("tags invalid")
        if not isinstance(row["permissionRef"], str) or not 1 <= len(row["permissionRef"]) <= 120:
            raise ValueError("permissionRef required")
        if not isinstance(row["revision"], int) or row["revision"] < 1:
            raise ValueError("revision must be positive integer")
        data = filepath.read_bytes()
        mime = image_type(data)
        found.append({**row, "mime": mime, "bytes": data, "sha256": sha256(data).hexdigest()})
    return found


class SearchStore:
    def __init__(self, db_path: Path, embedder, *, daily_cap=MAX_CALLS_PER_UTC_DAY):
        self.db = sqlite3.connect(db_path, timeout=30)
        self.db.row_factory = sqlite3.Row
        self.embedder = embedder
        self.daily_cap = daily_cap
        self.db.executescript("""
        CREATE TABLE IF NOT EXISTS assets (
          id TEXT PRIMARY KEY, file TEXT NOT NULL, title TEXT NOT NULL,
          tags TEXT NOT NULL, permission_ref TEXT NOT NULL, revision INTEGER NOT NULL,
          sha256 TEXT NOT NULL, mime TEXT NOT NULL, model TEXT NOT NULL,
          dimension INTEGER NOT NULL, vector TEXT NOT NULL, active INTEGER NOT NULL);
        CREATE TABLE IF NOT EXISTS operations (
          op_id TEXT PRIMARY KEY, kind TEXT NOT NULL, binding TEXT NOT NULL,
          state TEXT NOT NULL CHECK(state IN ('started','ready','done','abandoned')),
          vector TEXT, result TEXT, started_at TEXT NOT NULL, run_id TEXT);
        CREATE TABLE IF NOT EXISTS runs (
          run_id TEXT PRIMARY KEY, binding TEXT NOT NULL, base_fingerprint TEXT NOT NULL,
          state TEXT NOT NULL CHECK(state IN ('preparing','done','aborted')),
          started_at TEXT NOT NULL, activated_fingerprint TEXT);
        """)

    def close(self):
        self.db.close()

    def _begin(self, op_id, kind, binding, run_id=None, expected_fingerprint=None):
        # Reserve atomically even if a CLI and the local API touch the same DB.
        self.db.execute("BEGIN IMMEDIATE")
        try:
            if kind == "query" and self._pending_run():
                raise RuntimeError("index run incomplete; query is paused")
            # A completed replay is read under this lock only after validating its corpus.
            if kind == "query" and self.corpus_fingerprint() != expected_fingerprint:
                raise RuntimeError("corpus changed before query; use a new request id")
            if kind == "index" and not self.db.execute(
                "SELECT 1 FROM runs WHERE run_id=? AND state='preparing'", (run_id,)).fetchone():
                raise RuntimeError("index run is no longer active")
            existing = self.db.execute("SELECT * FROM operations WHERE op_id=?", (op_id,)).fetchone()
            if existing:
                if existing["kind"] != kind or existing["binding"] != binding or existing["run_id"] != run_id:
                    raise ValueError("operation id reused with different input or corpus")
                if existing["state"] in ("started", "abandoned"):
                    raise RuntimeError("uncertain call: reconcile provider usage; abandon explicitly before a new id")
                self.db.commit()
                return existing
            if self.db.execute("SELECT 1 FROM operations WHERE state='started' LIMIT 1").fetchone():
                raise RuntimeError("another provider call is uncertain; reconcile it first")
            day = now()[:10]
            used = self.db.execute("SELECT count(*) FROM operations WHERE substr(started_at,1,10)=?", (day,)).fetchone()[0]
            if used >= self.daily_cap:
                raise RuntimeError("daily provider-call ceiling reached")
            self.db.execute("INSERT INTO operations VALUES (?,?,?,?,?,?,?,?)", (op_id, kind, binding, "started", None, None, now(), run_id))
            self.db.commit()
            return None
        except Exception:
            self.db.rollback()
            raise

    def abandon(self, op_id):
        row = self.db.execute("SELECT state FROM operations WHERE op_id=?", (op_id,)).fetchone()
        if not row or row["state"] != "started":
            raise ValueError("only a started call can be abandoned")
        with self.db:
            self.db.execute("UPDATE operations SET state='abandoned' WHERE op_id=?", (op_id,))

    def abort_run(self, run_id):
        self.db.execute("BEGIN IMMEDIATE")
        try:
            run = self.db.execute("SELECT state FROM runs WHERE run_id=?", (run_id,)).fetchone()
            if not run or run["state"] != "preparing":
                raise ValueError("only an incomplete run can be aborted")
            if self.db.execute("SELECT 1 FROM operations WHERE run_id=? AND state='started'", (run_id,)).fetchone():
                raise RuntimeError("reconcile and abandon uncertain provider calls before aborting run")
            self.db.execute("UPDATE operations SET state='abandoned' WHERE run_id=? AND state='ready'", (run_id,))
            self.db.execute("UPDATE runs SET state='aborted' WHERE run_id=?", (run_id,))
            self.db.commit()
        except Exception:
            self.db.rollback()
            raise

    def _pending_run(self):
        return self.db.execute("SELECT run_id FROM runs WHERE state='preparing' LIMIT 1").fetchone()

    def _vector(self, op_id, existing, data, mime):
        if existing:
            return json.loads(existing["vector"])
        # The started receipt is committed before this potentially chargeable call.
        vector = unit_vector(self.embedder.embed(data, mime))
        with self.db:
            self.db.execute("UPDATE operations SET state='ready',vector=? WHERE op_id=?", (json.dumps(vector), op_id))
        return vector

    def _prevalidate(self, rows):
        seen = set()
        if not 1 <= len(rows) <= MAX_ITEMS:
            raise ValueError("manifest needs 1 to 50 images")
        for row in rows:
            if row["id"] in seen:
                raise ValueError("duplicate id")
            seen.add(row["id"])
            old = self.db.execute("SELECT * FROM assets WHERE id=?", (row["id"],)).fetchone()
            if old and row["revision"] < old["revision"]:
                raise ValueError("revision went backwards")
            if old and row["revision"] == old["revision"] and (
                old["sha256"] != row["sha256"] or old["file"] != row["file"] or
                old["title"] != row["title"] or json.loads(old["tags"]) != row["tags"] or
                old["permission_ref"] != row["permissionRef"]):
                raise ValueError("changed bytes or metadata require a new revision")

    def _run_binding(self, rows):
        fields = ("id", "file", "title", "tags", "permissionRef", "revision", "sha256", "mime")
        return json.dumps([[r[f] for f in fields] for r in sorted(rows, key=lambda r: r["id"])] + [MODEL, DIM], sort_keys=True)

    def _start_run(self, rows, run_id, binding):
        self.db.execute("BEGIN IMMEDIATE")
        try:
            prior = self.db.execute("SELECT * FROM runs WHERE run_id=?", (run_id,)).fetchone()
            if prior:
                if prior["binding"] != binding or prior["state"] == "aborted":
                    raise ValueError("run id reused with changed manifest or aborted run")
                if prior["state"] == "preparing" and prior["base_fingerprint"] != self.corpus_fingerprint():
                    raise RuntimeError("base corpus changed during incomplete run")
                if prior["state"] == "done" and prior["activated_fingerprint"] != self.corpus_fingerprint():
                    raise ValueError("completed run does not match current corpus; choose a new run id")
                self.db.commit()
                return prior["state"] == "done"
            if self._pending_run():
                raise RuntimeError("another index run is incomplete; resume or abort it")
            self._prevalidate(rows)  # Reject the whole manifest before any row mutation.
            self.db.execute("INSERT INTO runs VALUES (?,?,?,?,?,?)",
                (run_id, binding, self.corpus_fingerprint(), "preparing", now(), None))
            self.db.commit()
            return False
        except Exception:
            self.db.rollback()
            raise

    def index(self, rows, run_id):
        if not ID.fullmatch(run_id):
            raise ValueError("run id invalid")
        binding = self._run_binding(rows)
        if self._start_run(rows, run_id, binding):
            return {"active": len(rows), "model": MODEL, "dimension": DIM, "runState": "done"}
        ids = {r["id"] for r in rows}
        vectors = {}
        for row in rows:
            old = self.db.execute("SELECT * FROM assets WHERE id=?", (row["id"],)).fetchone()
            same_image = old and old["sha256"] == row["sha256"] and old["model"] == MODEL and old["dimension"] == DIM
            if same_image:
                vector = json.loads(old["vector"])
            else:
                op_id = f"{run_id}-{row['id']}-{row['revision']}"
                binding = json.dumps([row["id"], row["revision"], row["sha256"], MODEL, DIM])
                prior = self._begin(op_id, "index", binding, run_id)
                vector = self._vector(op_id, prior, row["bytes"], row["mime"])
            vectors[row["id"]] = vector
        # All public corpus changes, including removals and run completion, commit together.
        self.db.execute("BEGIN IMMEDIATE")
        try:
            run = self.db.execute("SELECT * FROM runs WHERE run_id=?", (run_id,)).fetchone()
            if not run or run["state"] != "preparing" or run["base_fingerprint"] != self.corpus_fingerprint():
                raise RuntimeError("base corpus changed before activation")
            self._prevalidate(rows)
            for row in rows:
                self.db.execute("""INSERT INTO assets VALUES (?,?,?,?,?,?,?,?,?,?,?,?)
                  ON CONFLICT(id) DO UPDATE SET file=excluded.file,title=excluded.title,tags=excluded.tags,
                  permission_ref=excluded.permission_ref,revision=excluded.revision,sha256=excluded.sha256,
                  mime=excluded.mime,model=excluded.model,dimension=excluded.dimension,
                  vector=excluded.vector,active=1""",
                  (row["id"], row["file"], row["title"], json.dumps(row["tags"]), row["permissionRef"], row["revision"], row["sha256"], row["mime"], MODEL, DIM, json.dumps(vectors[row["id"]]), 1))
            for asset in self.db.execute("SELECT id FROM assets WHERE active=1").fetchall():
                if asset["id"] not in ids:
                    self.db.execute("UPDATE assets SET active=0 WHERE id=?", (asset["id"],))
            self.db.execute("UPDATE operations SET state='done' WHERE run_id=? AND state='ready'", (run_id,))
            self.db.execute("UPDATE runs SET state='done',activated_fingerprint=? WHERE run_id=?",
                (self.corpus_fingerprint(), run_id))
            self.db.commit()
        except Exception:
            self.db.rollback()
            raise
        return {"active": len(ids), "model": MODEL, "dimension": DIM, "runState": "done"}

    def corpus_fingerprint(self):
        rows = self.db.execute("SELECT id,file,title,sha256,revision,tags,permission_ref,model,dimension FROM assets WHERE active=1 ORDER BY id").fetchall()
        return sha256(json.dumps([tuple(r) for r in rows]).encode()).hexdigest()

    def search(self, data, *, request_id, token, expected_token, allowed_permissions, tag=None, limit=5):
        if not expected_token or not hmac.compare_digest(token or "", expected_token):
            raise PermissionError("invalid bearer token")
        mime = image_type(data)
        if not ID.fullmatch(request_id) or not isinstance(limit, int) or not 1 <= limit <= 10:
            raise ValueError("request id or limit invalid")
        if tag is not None and not ID.fullmatch(tag):
            raise ValueError("tag invalid")
        if not allowed_permissions or any(not isinstance(p, str) for p in allowed_permissions):
            raise ValueError("permission scope missing")
        scope = sorted(set(allowed_permissions))
        if self._pending_run():
            raise RuntimeError("index run incomplete; query is paused")
        fingerprint = self.corpus_fingerprint()
        binding = json.dumps([sha256(data).hexdigest(), mime, fingerprint, scope, tag, limit, MODEL, DIM])
        prior = self._begin(request_id, "query", binding, expected_fingerprint=fingerprint)
        if prior and prior["state"] == "done":
            return json.loads(prior["result"])
        vector = self._vector(request_id, prior, data, mime)
        self.db.execute("BEGIN IMMEDIATE")
        try:
            if self._pending_run() or self.corpus_fingerprint() != fingerprint:
                self.db.execute("UPDATE operations SET state='abandoned',result=? WHERE op_id=?",
                    (json.dumps({"reason": "corpus changed or index incomplete after embedding"}), request_id))
                self.db.commit()
                raise RuntimeError("corpus changed during query; use a new request id")
            ranked = []
            for row in self.db.execute("SELECT * FROM assets WHERE active=1"):
                if row["permission_ref"] not in scope or (tag is not None and tag not in json.loads(row["tags"])):
                    continue
                if row["model"] != MODEL or row["dimension"] != DIM:
                    raise RuntimeError("mixed embedding spaces")
                score = sum(a*b for a,b in zip(vector, json.loads(row["vector"])))
                ranked.append({"id": row["id"], "title": row["title"], "tags": json.loads(row["tags"]), "revision": row["revision"], "score": round(score, 6)})
            ranked.sort(key=lambda r: (-r["score"], r["id"]))
            result = {"requestId": request_id, "corpusFingerprint": fingerprint, "results": ranked[:limit]}
            self.db.execute("UPDATE operations SET state='done',result=? WHERE op_id=?", (json.dumps(result), request_id))
            self.db.commit()
            return result
        except Exception:
            self.db.rollback()
            raise

    def status(self):
        return {"active": self.db.execute("SELECT count(*) FROM assets WHERE active=1").fetchone()[0],
                "corpusFingerprint": self.corpus_fingerprint(),
                "pendingRuns": [dict(r) for r in self.db.execute("SELECT run_id,state,started_at FROM runs WHERE state='preparing'")],
                "operations": [dict(r) for r in self.db.execute("SELECT op_id,kind,state,started_at,run_id FROM operations ORDER BY started_at")]}


class GeminiEmbedder:
    def __init__(self, api_key):
        from google.genai import Client, types
        self.client = Client(api_key=api_key, vertexai=False, enterprise=False,
            http_options=types.HttpOptions(timeout=30000, retry_options=types.HttpRetryOptions(attempts=1)))

    def embed(self, data, mime):
        from google.genai import types
        response = self.client.models.embed_content(model=MODEL,
            contents=[types.Part.from_bytes(data=data, mime_type=mime)],
            config=types.EmbedContentConfig(output_dimensionality=DIM))
        if len(response.embeddings or []) != 1:
            raise ValueError("embedding count")
        return response.embeddings[0].values

    def close(self):
        self.client.close()

"""Operator CLI. Index is chargeable; check and status are offline."""
import argparse
import json
import os
from pathlib import Path
from search_core import GeminiEmbedder, SearchStore, load_manifest


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=("check", "index", "status", "abandon", "abort-run"))
    parser.add_argument("--manifest", type=Path)
    parser.add_argument("--library", type=Path)
    parser.add_argument("--db", type=Path)
    parser.add_argument("--run-id")
    parser.add_argument("--op-id")
    args = parser.parse_args()
    if args.command in ("check", "index"):
        if not args.manifest or not args.library:
            parser.error("check/index require --manifest and --library")
        rows = load_manifest(args.manifest, args.library)
        print(json.dumps({"validatedRows": len(rows), "ids": [r["id"] for r in rows]}))
        if args.command == "check":
            return
    if not args.db:
        parser.error("index/status/abandon require --db")
    if args.command == "index" and not args.run_id:
        parser.error("index requires --run-id")
    if args.command == "abandon" and not args.op_id:
        parser.error("abandon requires --op-id")
    if args.command == "abort-run" and not args.run_id:
        parser.error("abort-run requires --run-id")
    embedder = GeminiEmbedder(os.environ["GEMINI_API_KEY"]) if args.command == "index" else None
    store = SearchStore(args.db, embedder)
    try:
        if args.command == "index":
            print(json.dumps(store.index(rows, args.run_id)))
        elif args.command == "status":
            print(json.dumps(store.status()))
        elif args.command == "abort-run":
            store.abort_run(args.run_id)
            print(json.dumps({"abortedRun": args.run_id}))
        else:
            store.abandon(args.op_id)
            print(json.dumps({"abandoned": args.op_id}))
    finally:
        store.close()
        if embedder:
            embedder.close()


if __name__ == "__main__":
    main()

"""Local-only API. Start one worker and bind 127.0.0.1."""
import os
import hmac
from pathlib import Path
import fastapi
from search_core import GeminiEmbedder, MAX_IMAGE, SearchStore

app = fastapi.FastAPI(openapi_url=None)


@app.post("/search")
async def search(request: fastapi.Request, request_id: str, tag: str | None = None, limit: int = 5):
    token = request.headers.get("authorization", "")
    expected = os.environ.get("SEARCH_API_TOKEN")
    if not expected or not hmac.compare_digest(token, f"Bearer {expected}"):
        raise fastapi.HTTPException(503 if not expected else 401)
    if request.headers.get("content-type") not in ("image/png", "image/jpeg"):
        raise fastapi.HTTPException(415)
    body = bytearray()
    async for chunk in request.stream():
        if len(body) + len(chunk) > MAX_IMAGE:
            raise fastapi.HTTPException(413)
        body += chunk
    embedder = None
    store = None
    try:
        key = os.environ["GEMINI_API_KEY"]
        db = Path(os.environ["SEARCH_DB"])
        scopes = [v for v in os.environ["SEARCH_ALLOWED_PERMISSION_REFS"].split(",") if v]
        embedder = GeminiEmbedder(key)
        store = SearchStore(db, embedder)
        return store.search(bytes(body), request_id=request_id,
            token=token[7:], expected_token=expected,
            allowed_permissions=scopes, tag=tag, limit=limit)
    except (PermissionError, ValueError, RuntimeError, KeyError) as exc:
        raise fastapi.HTTPException({PermissionError: 401, ValueError: 400}.get(type(exc), 409))
    finally:
        if store:
            store.close()
        if embedder:
            embedder.close()

[
  {
    "id": "your-owned-image-id",
    "file": "your-owned-image.png",
    "title": "Your descriptive title",
    "tags": ["your-approved-tag"],
    "permissionRef": "your-internal-access-scope",
    "revision": 1
  }
]

google-genai==2.33.0
fastapi==0.162.0
uvicorn==0.54.0
pillow==13.0.0

import json
import os
from base64 import b64decode
from pathlib import Path
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from search_core import DIM, GeminiEmbedder, SearchStore, image_type, load_manifest

PNG_A = b64decode("iVBORw0KGgoAAAANSUhEUgAAAAMAAAADCAIAAADZSiLoAAAAEUlEQVR4nGP8zwAFTDAGEgsAHVsBBeg3qMUAAAAASUVORK5CYII=")
PNG_B = b64decode("iVBORw0KGgoAAAANSUhEUgAAAAMAAAADCAIAAADZSiLoAAAAE0lEQVR4nGNkYPjPAAZMEAqFBQAbXQEF2cdRlAAAAABJRU5ErkJggg==")


class FakeEmbedder:
    def __init__(self):
        self.calls = []

    def embed(self, data, mime):
        self.calls.append(data)
        vector = [0.0] * DIM
        vector[0 if data == PNG_A else 1] = 1.0
        return vector


class SearchCoreTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.root = Path(self.tmp.name)
        self.lib = self.root / "images"
        self.lib.mkdir()
        (self.lib / "a.png").write_bytes(PNG_A)
        (self.lib / "b.png").write_bytes(PNG_B)
        self.manifest = self.root / "manifest.json"
        self.rows = [
            {"id": "a", "file": "a.png", "title": "A", "tags": ["red"], "permissionRef": "team-a", "revision": 1},
            {"id": "b", "file": "b.png", "title": "B", "tags": ["blue"], "permissionRef": "team-b", "revision": 1},
        ]
        self.manifest.write_text(json.dumps(self.rows))
        self.fake = FakeEmbedder()
        self.store = SearchStore(self.root / "search.sqlite", self.fake)
        self.addCleanup(self.store.close)

    def load(self):
        return load_manifest(self.manifest, self.lib)

    def query(self, **kw):
        return self.store.search(PNG_A, request_id="q1", token="secret",
            expected_token="secret", allowed_permissions=["team-a"], **kw)

    def test_image_validation_and_manifest_path(self):
        self.assertEqual(image_type(PNG_A), "image/png")
        with self.assertRaises(ValueError):
            image_type(b"not an image")
        self.rows[0]["file"] = "../a.png"
        self.manifest.write_text(json.dumps(self.rows))
        with self.assertRaises(ValueError):
            self.load()

    def test_index_filter_update_and_removal(self):
        self.store.index(self.load(), "run1")
        self.assertEqual(len(self.fake.calls), 2)
        result = self.query(tag="red")
        self.assertEqual([r["id"] for r in result["results"]], ["a"])
        self.assertEqual(len(self.fake.calls), 3)
        self.assertEqual(self.query(tag="red"), result)
        self.assertEqual(len(self.fake.calls), 3)
        with self.assertRaises(ValueError):
            self.query(tag="blue")  # Same request id, changed binding.
        self.rows[0]["title"] = "Updated A"
        self.rows[0]["revision"] = 2
        self.rows.pop()
        self.manifest.write_text(json.dumps(self.rows))
        self.store.index(self.load(), "run2")
        self.assertEqual(len(self.fake.calls), 3)  # Metadata-only edit saves a call.
        self.assertEqual(self.store.status()["active"], 1)
        with self.assertRaises(ValueError):
            self.query(tag="red")  # Corpus changed under request id.

    def test_auth_before_embedding_and_permission_filter(self):
        self.store.index(self.load(), "run1")
        with self.assertRaises(PermissionError):
            self.store.search(PNG_A, request_id="bad", token="wrong",
                expected_token="secret", allowed_permissions=["team-a"])
        self.assertEqual(len(self.fake.calls), 2)
        result = self.store.search(PNG_A, request_id="q2", token="secret",
            expected_token="secret", allowed_permissions=["team-b"])
        self.assertEqual([r["id"] for r in result["results"]], ["b"])

    def test_uncertain_call_never_auto_retries(self):
        self.fake.embed = lambda data, mime: (_ for _ in ()).throw(TimeoutError("unknown provider outcome"))
        with self.assertRaises(TimeoutError):
            self.store.index(self.load(), "run1")
        self.assertEqual(self.store.status()["operations"][0]["state"], "started")
        with self.assertRaises(RuntimeError):
            self.store.index(self.load(), "run1")
        self.store.abandon("run1-a-1")
        with self.assertRaises(RuntimeError):
            self.store.index(self.load(), "run1")

    def test_daily_call_ceiling(self):
        limited = SearchStore(self.root / "limited.sqlite", FakeEmbedder(), daily_cap=1)
        self.addCleanup(limited.close)
        with self.assertRaises(RuntimeError):
            limited.index(self.load(), "run1")
        self.assertEqual(len(limited.status()["operations"]), 1)

    def test_google_sdk_contract_and_route_without_transport(self):
        with patch.dict(os.environ, {"GOOGLE_GENAI_USE_VERTEXAI": "true", "GOOGLE_GENAI_USE_ENTERPRISE": "true"}):
            adapter = GeminiEmbedder("offline-key")
            self.assertEqual(adapter.client._api_client._http_options.base_url, "https://generativelanguage.googleapis.com/")
            with patch.object(adapter.client.models, "embed_content", return_value=SimpleNamespace(embeddings=[SimpleNamespace(values=[1.0] * DIM)])) as call:
                adapter.embed(PNG_A, "image/png")
                args = call.call_args.kwargs
                self.assertEqual(args["contents"][0].inline_data.mime_type, "image/png")
                self.assertEqual(args["config"].output_dimensionality, DIM)
            adapter.close()

    def test_query_rejects_concurrent_metadata_permission_update(self):
        self.store.index(self.load(), "seed")
        original = self.store.corpus_fingerprint()
        def update_during_call(data, mime):
            writer = SearchStore(self.root / "search.sqlite", FakeEmbedder())
            try:
                changed = self.load()
                changed[0]["title"] = "New A"
                changed[0]["permissionRef"] = "team-b"
                changed[0]["revision"] = 2
                writer.index(changed, "metadata")
            finally:
                writer.close()
            return [1.0] + [0.0] * (DIM - 1)
        self.fake.embed = update_during_call
        with self.assertRaisesRegex(RuntimeError, "corpus changed"):
            self.query()
        self.assertNotEqual(self.store.corpus_fingerprint(), original)
        op = self.store.db.execute("SELECT state,result FROM operations WHERE op_id='q1'").fetchone()
        self.assertEqual(op["state"], "abandoned")
        self.assertNotIn("results", op["result"])
        self.fake.embed = FakeEmbedder().embed
        result = self.store.search(PNG_A, request_id="q2", token="secret", expected_token="secret",
            allowed_permissions=["team-a"])
        self.assertEqual(result["results"], [])

    def test_query_rejects_concurrent_removal(self):
        self.store.index(self.load(), "seed")
        def remove_during_call(data, mime):
            writer = SearchStore(self.root / "search.sqlite", FakeEmbedder())
            try:
                writer.index([self.load()[1]], "remove-a")
            finally:
                writer.close()
            return [1.0] + [0.0] * (DIM - 1)
        self.fake.embed = remove_during_call
        with self.assertRaisesRegex(RuntimeError, "corpus changed"):
            self.query()
        self.assertEqual(self.store.status()["active"], 1)
        self.assertEqual(self.store.db.execute("SELECT state FROM operations WHERE op_id='q1'").fetchone()[0], "abandoned")

    def _cached_replay_race(self, change):
        self.store.index(self.load(), "seed")
        cached = self.query(tag="red")
        before = len(self.fake.calls)
        original_begin = self.store._begin
        def change_before_begin(*args, **kwargs):
            writer = SearchStore(self.root / "search.sqlite", FakeEmbedder())
            try:
                rows = self.load()
                writer.index(change(rows), "concurrent")
            finally:
                writer.close()
            return original_begin(*args, **kwargs)
        with patch.object(self.store, "_begin", side_effect=change_before_begin):
            with self.assertRaisesRegex(RuntimeError, "corpus changed before query"):
                self.query(tag="red")
        self.assertEqual(len(self.fake.calls), before)
        self.assertNotEqual(self.store.corpus_fingerprint(), cached["corpusFingerprint"])

    def test_cached_replay_rejects_permission_revocation(self):
        def change(rows):
            rows[0]["permissionRef"] = "team-b"
            rows[0]["revision"] = 2
            return rows
        self._cached_replay_race(change)
        self.assertEqual(self.store.search(PNG_A, request_id="q2", token="secret",
            expected_token="secret", allowed_permissions=["team-a"])["results"], [])

    def test_cached_replay_rejects_removal(self):
        self._cached_replay_race(lambda rows: rows[1:])
        self.assertEqual(self.store.status()["active"], 1)

    def test_cached_replay_rejects_tag_change(self):
        def change(rows):
            rows[0]["tags"] = ["blue"]
            rows[0]["revision"] = 2
            return rows
        self._cached_replay_race(change)

    def test_cached_replay_rejects_revision_change(self):
        def change(rows):
            rows[0]["title"] = "Revised A"
            rows[0]["revision"] = 2
            return rows
        self._cached_replay_race(change)

    def test_unchanged_cached_replay_skips_embedding(self):
        self.store.index(self.load(), "seed")
        cached = self.query(tag="red")
        before = len(self.fake.calls)
        self.assertEqual(self.query(tag="red"), cached)
        self.assertEqual(len(self.fake.calls), before)

    def test_prevalidates_entire_manifest_before_metadata_changes(self):
        self.rows[1]["revision"] = 2
        self.manifest.write_text(json.dumps(self.rows))
        self.store.index(self.load(), "seed")
        changed = self.load()
        changed[0]["revision"] = 2
        changed[0]["title"] = "Partial A"
        changed[1]["revision"] = 1
        with self.assertRaisesRegex(ValueError, "revision went backwards"):
            self.store.index(changed, "invalid")
        self.assertEqual(self.store.db.execute("SELECT title FROM assets WHERE id='a'").fetchone()[0], "A")
        self.assertEqual(self.store.status()["pendingRuns"], [])

    def test_interrupted_run_keeps_old_corpus_and_can_abort(self):
        self.store.index(self.load(), "seed")
        old_fingerprint = self.store.corpus_fingerprint()
        changed = self.load()[:1]
        changed[0]["revision"] = 2
        changed[0]["bytes"] = PNG_B
        changed[0]["sha256"] = __import__("hashlib").sha256(PNG_B).hexdigest()
        self.fake.embed = lambda data, mime: (_ for _ in ()).throw(TimeoutError("uncertain"))
        with self.assertRaises(TimeoutError):
            self.store.index(changed, "replace-and-remove")
        self.assertEqual(self.store.corpus_fingerprint(), old_fingerprint)
        self.assertEqual(self.store.status()["pendingRuns"][0]["run_id"], "replace-and-remove")
        with self.assertRaisesRegex(RuntimeError, "index run incomplete"):
            self.query()
        with self.assertRaisesRegex(RuntimeError, "reconcile"):
            self.store.abort_run("replace-and-remove")
        self.store.abandon("replace-and-remove-a-2")
        self.store.abort_run("replace-and-remove")
        self.assertEqual(self.store.status()["pendingRuns"], [])
        self.assertEqual(self.store.corpus_fingerprint(), old_fingerprint)
        self.fake.embed = FakeEmbedder().embed
        self.store.index(changed, "new-run")
        self.assertEqual(self.store.status()["active"], 1)

    def test_later_row_failure_keeps_metadata_and_removals_unapplied(self):
        self.store.index(self.load(), "seed")
        changed = self.load()
        changed[0]["title"] = "Not yet visible"
        changed[0]["revision"] = 2
        changed[1]["revision"] = 2
        changed[1]["bytes"] = PNG_A
        changed[1]["sha256"] = __import__("hashlib").sha256(PNG_A).hexdigest()
        self.fake.embed = lambda data, mime: (_ for _ in ()).throw(TimeoutError("uncertain"))
        with self.assertRaises(TimeoutError):
            self.store.index(changed, "late-failure")
        self.assertEqual(self.store.db.execute("SELECT title FROM assets WHERE id='a'").fetchone()[0], "A")
        self.assertEqual(self.store.status()["active"], 2)
        self.assertEqual(self.store.status()["pendingRuns"][0]["run_id"], "late-failure")

    def test_ready_vector_replay_activates_without_new_provider_call(self):
        rows = self.load()[:1]
        original = self.store._prevalidate
        calls = 0
        def interrupt_after_embedding(items):
            nonlocal calls
            calls += 1
            if calls == 2:
                raise RuntimeError("local interruption before activation")
            return original(items)
        with patch.object(self.store, "_prevalidate", side_effect=interrupt_after_embedding):
            with self.assertRaisesRegex(RuntimeError, "before activation"):
                self.store.index(rows, "ready-replay")
        self.assertEqual(self.store.status()["active"], 0)
        self.assertEqual(self.store.status()["pendingRuns"][0]["run_id"], "ready-replay")
        self.assertEqual(len(self.fake.calls), 1)
        self.store.index(rows, "ready-replay")
        self.assertEqual(len(self.fake.calls), 1)
        self.assertEqual(self.store.status()["active"], 1)


if __name__ == "__main__":
    unittest.main()

import os
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch
from fastapi.testclient import TestClient
import app as module
from search_core import SearchStore, DIM, MAX_IMAGE
from test_search_core import PNG_A as PNG


class FakeEmbedder:
    def __init__(self, key):
        self.calls = 0
    def embed(self, data, mime):
        self.calls += 1
        return [1.0] + [0.0] * (DIM - 1)
    def close(self):
        pass


class ApiTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.db = Path(self.tmp.name) / "db.sqlite"
        fake = FakeEmbedder("fake")
        store = SearchStore(self.db, fake)
        store.index([{"id":"a","file":"a.png","title":"A","tags":["red"],
            "permissionRef":"team-a","revision":1,"mime":"image/png",
            "bytes":PNG,"sha256":__import__("hashlib").sha256(PNG).hexdigest()}], "seed")
        store.close()
        self.env = patch.dict(os.environ, {"SEARCH_API_TOKEN":"secret", "GEMINI_API_KEY":"fake",
            "SEARCH_DB":str(self.db), "SEARCH_ALLOWED_PERMISSION_REFS":"team-a"})
        self.env.start()
        self.addCleanup(self.env.stop)
        self.mock = patch.object(module, "GeminiEmbedder", FakeEmbedder)
        self.mock.start()
        self.addCleanup(self.mock.stop)
        self.client = TestClient(module.app)

    def post(self, rid="req1", token="Bearer secret", tag="red", data=PNG, mime="image/png", status=200):
        result = self.client.post("/search", params={"request_id": rid, "tag": tag},
            content=data, headers={"content-type": mime, **({"authorization": token} if token else {})})
        self.assertEqual(result.status_code, status)
        return result

    def test_boundary_and_filter(self):
        for token in (None, "Bearer wrong"):
            self.post(token=token, status=401)
        good = self.post().json()
        self.assertEqual((good["results"][0]["id"], self.post().json()), ("a", good))
        self.post(tag="blue", status=400)

    def test_request_size_and_type(self):
        self.post(data=b"a" * (MAX_IMAGE + 1), status=413)
        self.post(data=b"bad", status=400)
        self.post(mime="text/plain", status=415)


if __name__ == "__main__":
    unittest.main()

Set-Location 'C:\path\to\image-search-pilot'
py -3.12 -m venv .venv
$Python = (Resolve-Path '.\.venv\Scripts\python.exe').Path
& $Python --version
& $Python -m pip install -r .\requirements.txt
& $Python -B -m unittest -v test_search_core.py test_app.py
& $Python .\cli.py check --manifest .\manifest.json --library .\images

$env:GEMINI_API_KEY = [System.Net.NetworkCredential]::new('', (Read-Host 'Gemini API key from your secret manager' -AsSecureString)).Password
$env:SEARCH_API_TOKEN = [System.Net.NetworkCredential]::new('', (Read-Host 'Local API token from your secret manager' -AsSecureString)).Password

$env:SEARCH_DB = (Join-Path (Get-Location) 'private-search.sqlite')
$env:SEARCH_ALLOWED_PERMISSION_REFS = 'catalog-team'
& $Python .\cli.py index --manifest .\manifest.json --library .\images --db .\private-search.sqlite --run-id first-owned-library
& $Python .\cli.py status --db .\private-search.sqlite

& $Python .\cli.py status --db .\private-search.sqlite
& $Python .\cli.py abandon --db .\private-search.sqlite --op-id 'OP_ID_FROM_STATUS'
& $Python .\cli.py abort-run --db .\private-search.sqlite --run-id 'RUN_ID_FROM_STATUS'

& '.\.venv\Scripts\uvicorn.exe' app:app --host 127.0.0.1 --port 8765 --workers 1 --no-access-log

Set-Location 'C:\path\to\image-search-pilot'
$QueryImage = (Resolve-Path '.\images\your-approved-query.png').Path
$RequestId = 'owned-query-001'
$Tag = 'your-approved-tag'
$ImageBytes = [System.IO.File]::ReadAllBytes($QueryImage)
$Uri = 'http://127.0.0.1:8765/search?request_id=' + [uri]::EscapeDataString($RequestId) + '&tag=' + [uri]::EscapeDataString($Tag) + '&limit=5'
$Response = Invoke-RestMethod -Method Post -Uri $Uri -Headers @{ Authorization = "Bearer $env:SEARCH_API_TOKEN" } -ContentType 'image/png' -Body $ImageBytes
$Response | ConvertTo-Json -Depth 6