api/routes1.py

"""
Routes for the Image Similarity Search API
Contains all endpoints for the application
"""

from fastapi import APIRouter, FastAPI, File, UploadFile, Form, Query, Path 
from typing import List, Optional
from pydantic import BaseModel

from services.embedding_service import ImageEmbeddingModel
from services.vector_db_service import VectorDatabaseClient


class SearchResponse(BaseModel):
    """Response model for search results"""
    image_id: str
    similarity: float
    metadata: Optional[dict] = None


def register_routes(
    app: FastAPI,
    embedding_model: ImageEmbeddingModel,
    vector_db: VectorDatabaseClient,
    # Remove security_service parameter
):
    """Register all routes with the FastAPI app"""
    router = APIRouter()

    @router.post("/upload", response_model=dict)
    async def upload_image(
        file: UploadFile = File(...),
        metadata: Optional[str] = Form(None),
        # Remove security dependency: api_key: str = Depends(security_service.verify_api_key)
    ):
        """Upload an image and store its embedding"""
        # Process the image and generate embedding
        image_data = await file.read()
        embedding = embedding_model.generate_embedding(image_data)
        
        # Store in vector database with optional metadata
        image_id = vector_db.add_embedding(embedding, file.filename, metadata)
        
        return {"image_id": image_id, "message": "Image uploaded successfully"}
    
    @router.get("/search/by-id/{image_id}", response_model=List[SearchResponse])
    async def search_by_id(
        image_id: str = Path(..., description="ID of the uploaded image to use as query"),
        limit: int = Query(5, description="Maximum number of results to return"),
        # Remove security dependency: api_key: str = Depends(security_service.verify_api_key)
    ):
        """Search for similar images using an existing image ID as the query"""
        results = vector_db.search_by_id(image_id, limit)
        return [
            SearchResponse(
                image_id=result.id,
                similarity=result.score,
                metadata=result.metadata
            )
            for result in results
        ]
    
    @router.post("/search/by-image", response_model=List[SearchResponse])
    async def search_by_image(
        file: UploadFile = File(...),
        limit: int = Query(5, description="Maximum number of results to return"),
        # Remove security dependency: api_key: str = Depends(security_service.verify_api_key)
    ):
        """Search for similar images by uploading a new image"""
        # Process the image and generate embedding
        image_data = await file.read()
        embedding = embedding_model.generate_embedding(image_data)
        
        # Search using the embedding
        results = vector_db.search_by_embedding(embedding, limit)
        return [
            SearchResponse(
                image_id=result.id,
                similarity=result.score,
                metadata=result.metadata
            )
            for result in results
        ]
    
    @router.delete("/images/{image_id}")
    async def delete_image(
        image_id: str = Path(..., description="ID of the image to delete"),
        # Remove security dependency: api_key: str = Depends(security_service.verify_api_key)
    ):
        """Delete an image from the database"""
        success = vector_db.delete_embedding(image_id)
        if success:
            return {"message": f"Image {image_id} deleted successfully"}
        return {"message": f"Image {image_id} not found"}
    
    
    # Add the router to the app
    app.include_router(router, prefix="/api/v1")