cloudzy/routes/search.py

"""Semantic search endpoint using FAISS"""
from fastapi import APIRouter, Query, Depends, HTTPException
from sqlmodel import Session, select
import numpy as np

from cloudzy.database import get_session
from cloudzy.models import Photo
from cloudzy.schemas import SearchResponse, SearchResult
from cloudzy.search_engine import SearchEngine
# from cloudzy.ai_utils import generate_filename_embedding
from cloudzy.ai_utils import  ImageEmbeddingGenerator
import os

router = APIRouter(tags=["search"])


@router.get("/search", response_model=SearchResponse)
async def search_photos(
    q: str = Query(..., min_length=1, max_length=200, description="Search query"),
    top_k: int = Query(5, ge=1, le=50, description="Number of results"),
    session: Session = Depends(get_session),
):
    """
    Semantic search for similar photos using FAISS.
    
    Converts query to embedding and finds most similar images.
    
    Args:
        q: Search query (used to generate embedding)
        top_k: Number of results to return (max 50)
    
    Returns: List of similar photos with distance scores
    """

    generator = ImageEmbeddingGenerator()
    query_embedding = generator._embed_text(q)


    
    # Search in FAISS
    search_engine = SearchEngine()
    search_results = search_engine.search(query_embedding, top_k=top_k)
    
    
    if not search_results:
        return SearchResponse(
            query=q,
            results=[],
            total_results=0,
        )
    
    APP_DOMAIN = os.getenv("APP_DOMAIN")

    
    
    # Fetch photo details from database
    result_objects = []
    for photo_id, distance in search_results:
        statement = select(Photo).where(Photo.id == photo_id)
        photo = session.exec(statement).first()
        
        if photo:  # Only include if photo exists in DB
            result_objects.append(
                SearchResult(
                    photo_id=photo.id,
                    filename=photo.filename,
                    image_url = f"{APP_DOMAIN}uploads/{photo.filename}",
                    tags=photo.get_tags(),
                    caption=photo.caption,
                    distance=distance,
                )
            )
    
    return SearchResponse(
        query=q,
        results=result_objects,
        total_results=len(result_objects),
    )


@router.post("/search/image-to-image")
async def image_to_image_search(
    reference_photo_id: int = Query(..., description="Reference photo ID"),
    top_k: int = Query(5, ge=1, le=50),
    session: Session = Depends(get_session),
):
    """
    Find similar images to a reference photo (image-to-image search).
    
    Args:
        reference_photo_id: ID of the reference photo
        top_k: Number of similar results
    
    Returns: Similar photos
    """
    # Get reference photo
    statement = select(Photo).where(Photo.id == reference_photo_id)
    reference_photo = session.exec(statement).first()
    
    if not reference_photo:
        raise HTTPException(status_code=404, detail=f"Photo {reference_photo_id} not found")
    
    # Get reference embedding
    reference_embedding = reference_photo.get_embedding()
    if not reference_embedding:
        raise HTTPException(status_code=400, detail="Photo has no embedding")
    
    # Search in FAISS
    search_engine = SearchEngine()
    search_results = search_engine.search(
        np.array(reference_embedding, dtype=np.float32),
        top_k=top_k + 1  # +1 to skip the reference photo itself
    )
    
    # Build results (skip first result which is the reference photo itself)
    result_objects = []
    for photo_id, distance in search_results[1:]:  # Skip first result
        statement = select(Photo).where(Photo.id == photo_id)
        photo = session.exec(statement).first()
        
        if photo:
            result_objects.append(
                SearchResult(
                    photo_id=photo.id,
                    filename=photo.filename,
                    tags=photo.get_tags(),
                    caption=photo.caption,
                    distance=distance,
                )
            )
    
    return SearchResponse(
        query=f"Similar to photo {reference_photo_id}",
        results=result_objects[:top_k],
        total_results=len(result_objects),
    )