import logging from typing import Optional, List, Tuple from fastapi import UploadFile, Request from bson import ObjectId import io from src.db.repositories.image_repository import image_repository from src.services.storage import StorageService from src.services.image_processor import ImageProcessor from src.services.embedding_service import EmbeddingService from src.services.pubsub_service import pubsub_service from src.models.image import ImageModel from src.models.user import UserModel from src.schemas.image import ImageResponse, ImageListResponse, ImageCreate, ImageUpdate logger = logging.getLogger(__name__) class ImageService: """Service class for handling image-related business logic""" def __init__(self): self.storage_service = StorageService() self.image_processor = ImageProcessor() self.embedding_service = EmbeddingService() def _generate_api_download_url(self, request: Request, image_id: str) -> str: """Generate API download URL for an image""" base_url = str(request.base_url).rstrip('/') return f"{base_url}/api/v1/images/{image_id}/download" async def upload_image( self, file: UploadFile, user: UserModel, request: Request, description: Optional[str] = None, collection_id: Optional[str] = None ) -> ImageResponse: """ Upload a new image Args: file: The uploaded file user: The user uploading the image request: The FastAPI request object for URL generation description: Optional description for the image collection_id: Optional collection ID to associate with the image Returns: ImageResponse: The created image metadata Raises: ValueError: If file validation fails RuntimeError: If upload fails """ # Validate file type if not file.content_type or not file.content_type.startswith('image/'): raise ValueError("File must be an image") # Validate file size (10MB limit) max_size = 10 * 1024 * 1024 # 10MB content = await file.read() if len(content) > max_size: raise ValueError("File size exceeds 10MB limit") # Reset file pointer await file.seek(0) try: # Upload to storage storage_path, content_type, file_size, metadata = await self.storage_service.upload_file( file, str(user.team_id) ) # Create image record image = ImageModel( filename=file.filename, original_filename=file.filename, file_size=file_size, content_type=content_type, storage_path=storage_path, public_url=None, # Will be set after we have the image ID team_id=user.team_id, uploader_id=user.id, description=description, metadata=metadata, collection_id=ObjectId(collection_id) if collection_id else None ) # Save to database created_image = await image_repository.create(image) # Generate API download URL now that we have the image ID api_download_url = self._generate_api_download_url(request, str(created_image.id)) # Update the image with the API download URL await image_repository.update(created_image.id, {"public_url": api_download_url}) created_image.public_url = api_download_url # Publish image processing task to Pub/Sub try: task_published = await pubsub_service.publish_image_processing_task( image_id=str(created_image.id), storage_path=storage_path, team_id=str(user.team_id) ) if not task_published: logger.warning(f"Failed to publish processing task for image {created_image.id}") except Exception as e: logger.warning(f"Failed to publish image processing task: {e}") # Convert to response return self._convert_to_response(created_image, request) except Exception as e: logger.error(f"Error uploading image: {e}") raise RuntimeError("Failed to upload image") async def list_images( self, user: UserModel, request: Request, skip: int = 0, limit: int = 50, collection_id: Optional[str] = None ) -> ImageListResponse: """ List images for the user's team or all images if user is admin Args: user: The requesting user request: The FastAPI request object for URL generation skip: Number of records to skip for pagination limit: Maximum number of records to return collection_id: Optional filter by collection ID Returns: ImageListResponse: List of images with pagination metadata """ # Check if user is admin - if so, get all images across all teams if user.is_admin: images = await image_repository.get_all_with_pagination( skip=skip, limit=limit, collection_id=ObjectId(collection_id) if collection_id else None, ) total = await image_repository.count_all( collection_id=ObjectId(collection_id) if collection_id else None, ) else: # Regular users only see images from their team images = await image_repository.get_by_team( user.team_id, skip=skip, limit=limit, collection_id=ObjectId(collection_id) if collection_id else None, ) total = await image_repository.count_by_team( user.team_id, collection_id=ObjectId(collection_id) if collection_id else None, ) # Convert to response response_images = [self._convert_to_response(image, request) for image in images] return ImageListResponse(images=response_images, total=total, skip=skip, limit=limit) async def get_image(self, image_id: str, user: UserModel, request: Request) -> ImageResponse: """ Get image metadata by ID Args: image_id: The image ID to retrieve user: The requesting user request: The FastAPI request object for URL generation Returns: ImageResponse: The image metadata Raises: ValueError: If image_id is invalid RuntimeError: If image not found PermissionError: If user not authorized to access the image """ try: obj_id = ObjectId(image_id) except Exception: raise ValueError("Invalid image ID") # Get image image = await image_repository.get_by_id(obj_id) if not image: raise RuntimeError("Image not found") # Check team access (admins can access any image) if not user.is_admin and image.team_id != user.team_id: raise PermissionError("Not authorized to access this image") return self._convert_to_response(image, request, include_last_accessed=True) async def download_image(self, image_id: str, user: UserModel) -> Tuple[bytes, str, str]: """ Download image file Args: image_id: The image ID to download user: The requesting user Returns: Tuple[bytes, str, str]: File content, content type, and filename Raises: ValueError: If image_id is invalid RuntimeError: If image not found or file not found in storage PermissionError: If user not authorized to access the image """ try: obj_id = ObjectId(image_id) except Exception: raise ValueError("Invalid image ID") # Get image image = await image_repository.get_by_id(obj_id) if not image: raise RuntimeError("Image not found") # Check team access (admins can access any image) if not user.is_admin and image.team_id != user.team_id: raise PermissionError("Not authorized to access this image") # Get file from storage file_content = self.storage_service.get_file(image.storage_path) if not file_content: raise RuntimeError("Image file not found in storage") # Update last accessed await image_repository.update_last_accessed(obj_id) return file_content, image.content_type, image.original_filename async def update_image( self, image_id: str, image_data: ImageUpdate, user: UserModel, request: Request ) -> ImageResponse: """ Update image metadata Args: image_id: The image ID to update image_data: The update data user: The requesting user request: The FastAPI request object for URL generation Returns: ImageResponse: The updated image metadata Raises: ValueError: If image_id is invalid RuntimeError: If image not found or update fails PermissionError: If user not authorized to update the image """ try: obj_id = ObjectId(image_id) except Exception: raise ValueError("Invalid image ID") # Get image image = await image_repository.get_by_id(obj_id) if not image: raise RuntimeError("Image not found") # Check team access (admins can update any image) if not user.is_admin and image.team_id != user.team_id: raise PermissionError("Not authorized to update this image") # Update image update_data = image_data.dict(exclude_unset=True) if not update_data: # No fields to update return self._convert_to_response(image, request) updated_image = await image_repository.update(obj_id, update_data) if not updated_image: raise RuntimeError("Failed to update image") return self._convert_to_response(updated_image, request) async def delete_image(self, image_id: str, user: UserModel) -> bool: """ Delete an image Args: image_id: The image ID to delete user: The requesting user Returns: bool: True if successfully deleted Raises: ValueError: If image_id is invalid RuntimeError: If image not found or deletion fails PermissionError: If user not authorized to delete the image """ try: obj_id = ObjectId(image_id) except Exception: raise ValueError("Invalid image ID") # Get image image = await image_repository.get_by_id(obj_id) if not image: raise RuntimeError("Image not found") # Check team access (admins can delete any image) if not user.is_admin and image.team_id != user.team_id: raise PermissionError("Not authorized to delete this image") # Delete from storage try: self.storage_service.delete_file(image.storage_path) except Exception as e: logger.warning(f"Failed to delete file from storage: {e}") # Delete from vector database if it has embeddings if image.has_embedding and image.embedding_id: try: await self.embedding_service.delete_embedding(image.embedding_id) except Exception as e: logger.warning(f"Failed to delete embedding: {e}") # Delete from database result = await image_repository.delete(obj_id) if not result: raise RuntimeError("Failed to delete image") return True def _convert_to_response( self, image: ImageModel, request: Request, include_last_accessed: bool = False ) -> ImageResponse: """Convert ImageModel to ImageResponse""" api_download_url = self._generate_api_download_url(request, str(image.id)) response_data = { "id": str(image.id), "filename": image.filename, "original_filename": image.original_filename, "file_size": image.file_size, "content_type": image.content_type, "storage_path": image.storage_path, "public_url": api_download_url, "team_id": str(image.team_id), "uploader_id": str(image.uploader_id), "upload_date": image.upload_date, "description": image.description, "metadata": image.metadata, "has_embedding": image.has_embedding, "collection_id": str(image.collection_id) if image.collection_id else None } if include_last_accessed: response_data["last_accessed"] = image.last_accessed return ImageResponse(**response_data)