Spaces:
Sleeping
Sleeping
| import asyncio | |
| import time | |
| import traceback | |
| from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile, Depends | |
| from src.core.config import ( | |
| DEFAULT_PINECONE_KEY, IDX_FACES, IDX_OBJECTS, | |
| FACE_MATCH_THRESHOLD, REDIS_TTL_SEARCH, REDIS_TTL_INFERENCE, | |
| FACE_SPACE_URL, OBJECT_SPACE_URL, | |
| FACE_QUALITY_GATE_SEARCH, | |
| ) | |
| from src.core.security import get_verified_keys | |
| from src.services.db_client import ( | |
| intersect_face_results, merge_face_results, merge_object_results, | |
| pinecone_pool, search_faces, search_objects, | |
| ) | |
| from src.services.cache import cache_get, cache_set, make_face_search_key, make_obj_search_key, make_inference_key | |
| from src.core.logging import log | |
| from src.common.utils import face_ui_score, get_ip, is_default_key, to_list | |
| router = APIRouter() | |
| async def search_database( | |
| request: Request, | |
| file: UploadFile = File(...), | |
| detect_faces: bool = Form(True), | |
| user_id: str = Form(""), | |
| keys: dict = Depends(get_verified_keys) | |
| ): | |
| ip = get_ip(request) | |
| start = time.perf_counter() | |
| mode = "guest" if is_default_key(keys["pinecone_key"], DEFAULT_PINECONE_KEY) else "personal" | |
| log("INFO", "search.start", user_id=user_id or "anonymous", ip=ip, mode=mode, | |
| filename=file.filename, detect_faces=detect_faces) | |
| try: | |
| file_bytes = await file.read() | |
| ai_manager = request.app.state.ai | |
| sem = request.app.state.ai_semaphore | |
| if FACE_SPACE_URL and OBJECT_SPACE_URL: | |
| # Distributed mode: call remote Spaces for embedding, check Redis cache first | |
| print(f"[DEBUG] Calling face space at: {FACE_SPACE_URL}") # add this | |
| print(f"[DEBUG] Calling object space at: {OBJECT_SPACE_URL}") # add this | |
| from src.services.space_client import embed_face, embed_object | |
| inference_cache_key = make_inference_key(file_bytes, detect_faces, "search") | |
| vectors = await cache_get(inference_cache_key) | |
| if vectors is None: | |
| face_task = embed_face(file_bytes, quality_gate=FACE_QUALITY_GATE_SEARCH) if detect_faces else asyncio.sleep(0) | |
| obj_task = embed_object(file_bytes) | |
| face_vecs, obj_vecs = await asyncio.gather(face_task, obj_task, return_exceptions=True) | |
| if isinstance(face_vecs, Exception) or face_vecs is None: | |
| log("WARNING", "search.face_space_unavailable", user_id=user_id or "anonymous", ip=ip) | |
| face_vecs = [] | |
| if isinstance(obj_vecs, Exception) or obj_vecs is None: | |
| raise HTTPException(503, "Object AI service is waking up (free tier). Please retry in 60 seconds.") | |
| vectors = [*(face_vecs or []), *(obj_vecs or [])] | |
| await cache_set(inference_cache_key, vectors, ttl=REDIS_TTL_INFERENCE) | |
| else: | |
| # Monolith mode: local inference | |
| async with sem: | |
| vectors = await ai_manager.process_image_bytes_async(file_bytes, detect_faces=detect_faces, mode="search") | |
| inference_ms = round((time.perf_counter() - start) * 1000) | |
| face_vectors = [v for v in vectors if v["type"] == "face"] | |
| object_vectors = [v for v in vectors if v["type"] == "object"] | |
| lanes_used = list({v["type"] for v in vectors}) | |
| log("INFO", "search.inference_done", user_id=user_id or "anonymous", ip=ip, mode=mode, | |
| face_vecs=len(face_vectors), obj_vecs=len(object_vectors), inference_ms=inference_ms) | |
| pc = pinecone_pool.get(keys["pinecone_key"]) | |
| idx_obj = pc.Index(IDX_OBJECTS) | |
| idx_face = pc.Index(IDX_FACES) | |
| if detect_faces and face_vectors: | |
| return await _run_face_search(face_vectors, object_vectors, idx_face, idx_obj, start, user_id, ip, mode, lanes_used) | |
| else: | |
| return await _run_object_search(object_vectors, idx_obj, start, user_id, ip, mode, lanes_used) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| log("ERROR", "search.error", user_id=user_id or "anonymous", ip=ip, mode=mode, | |
| error=str(e), traceback=traceback.format_exc()[-800:]) | |
| raise HTTPException(500, str(e)) | |
| async def _run_face_search(face_vectors, object_vectors, idx_face, idx_obj, start, user_id, ip, mode, lanes_used) -> dict: | |
| async def _query_face(fv: dict) -> dict: | |
| vec = to_list(fv["vector"]) | |
| det_score = fv.get("det_score", 1.0) | |
| try: | |
| ck = make_face_search_key(vec, FACE_MATCH_THRESHOLD) | |
| image_map = await cache_get(ck) | |
| if image_map is None: | |
| image_map = await asyncio.to_thread(search_faces, idx_face, vec, det_score) | |
| await cache_set(ck, image_map, ttl=REDIS_TTL_SEARCH) | |
| except Exception as e: | |
| if "404" in str(e): | |
| raise HTTPException(404, "Pinecone index not found. Go to Settings → Verify & Save.") | |
| raise | |
| return { | |
| "query_face_idx": fv.get("face_idx", 0), | |
| "query_face_crop": fv.get("face_crop", ""), | |
| "query_bbox": fv.get("bbox", []), | |
| "det_score": det_score, | |
| "face_width_px": fv.get("face_width_px", 0), | |
| "matches": sorted( | |
| [ | |
| { | |
| "url": url, | |
| "score": face_ui_score(d["raw_score"]), | |
| "raw_score": round(d["raw_score"], 4), | |
| "face_crop": d["face_crop"], | |
| "folder": d["folder"], | |
| "caption": "👤 Verified Identity", | |
| } | |
| for url, d in image_map.items() | |
| ], | |
| key=lambda x: x["score"], reverse=True, | |
| )[:50], | |
| } | |
| async def _query_obj_single(ov: dict) -> list: | |
| vec = to_list(ov["vector"]) | |
| try: | |
| ck = make_obj_search_key(vec) | |
| result = await cache_get(ck) | |
| if result is None: | |
| result = await asyncio.to_thread(search_objects, idx_obj, vec) | |
| await cache_set(ck, result, ttl=REDIS_TTL_SEARCH) | |
| return result | |
| except Exception as e: | |
| if "404" in str(e): | |
| raise HTTPException(404, "Pinecone index not found.") | |
| raise | |
| face_tasks = [_query_face(fv) for fv in face_vectors] | |
| obj_tasks = [_query_obj_single(ov) for ov in object_vectors] | |
| all_results = await asyncio.gather(*face_tasks, *obj_tasks) | |
| raw_groups = list(all_results[:len(face_tasks)]) | |
| obj_nested = list(all_results[len(face_tasks):]) | |
| merged_face = merge_face_results(raw_groups) | |
| merged_objects = merge_object_results(obj_nested) | |
| face_groups = [g for g in raw_groups if g.get("matches")] | |
| # Photos where ALL searched faces appear together (only meaningful for multi-face queries) | |
| group_results = intersect_face_results(face_groups, min_faces_required=len(face_groups)) if len(face_groups) > 1 else [] | |
| duration_ms = round((time.perf_counter() - start) * 1000) | |
| log("INFO", "search.complete", user_id=user_id or "anonymous", ip=ip, mode=mode, | |
| lanes=["face", "object"], face_groups=len(face_groups), face_results=len(merged_face), | |
| group_results=len(group_results), object_results=len(merged_objects), duration_ms=duration_ms) | |
| return { | |
| "mode": "face", | |
| "face_groups": face_groups, | |
| "results": merged_face, | |
| "group_results": group_results, | |
| "object_results": merged_objects, | |
| } | |
| async def _run_object_search(object_vectors, idx_obj, start, user_id, ip, mode, lanes_used) -> dict: | |
| if not object_vectors: | |
| return {"mode": "object", "results": [], "face_groups": []} | |
| async def _query_obj(ov: dict) -> list: | |
| vec = to_list(ov["vector"]) | |
| try: | |
| ck = make_obj_search_key(vec) | |
| result = await cache_get(ck) | |
| if result is None: | |
| result = await asyncio.to_thread(search_objects, idx_obj, vec) | |
| await cache_set(ck, result, ttl=REDIS_TTL_SEARCH) | |
| return result | |
| except Exception as e: | |
| if "404" in str(e): | |
| raise HTTPException(404, "Pinecone index not found.") | |
| raise | |
| nested = await asyncio.gather(*[_query_obj(ov) for ov in object_vectors]) | |
| final = merge_object_results(nested) | |
| duration_ms = round((time.perf_counter() - start) * 1000) | |
| log("INFO", "search.complete", user_id=user_id or "anonymous", ip=ip, mode=mode, | |
| lanes=lanes_used, results=len(final), duration_ms=duration_ms) | |
| return {"mode": "object", "results": final, "face_groups": []} |