feat(fastapi): adds image generation service

servers/fastapi/api/lifespan.py

@@ -4,67 +4,8 @@ import os
 from fastapi import FastAPI
 from sqlmodel import SQLModel
 
-from constants.supported_ollama_models import SUPPORTED_OLLAMA_MODELS
-from enums.llm_provider import LLMProvider
 from services import SQL_ENGINE
-from utils.custom_llm_provider import list_available_custom_models
-from utils.llm_provider import (
-    get_llm_provider,
-    is_custom_llm_selected,
-    is_ollama_selected,
-)
-from utils.ollama import pull_ollama_model
-
-
-can_change_keys = os.getenv("CAN_CHANGE_KEYS") != "false"
-
-
-async def check_llm_model_availability():
-    if not can_change_keys:
-        if get_llm_provider() == LLMProvider.OPENAI:
-            openai_api_key = os.getenv("OPENAI_API_KEY")
-            if not openai_api_key:
-                raise Exception("OPENAI_API_KEY must be provided")
-
-        elif get_llm_provider() == LLMProvider.GOOGLE:
-            google_api_key = os.getenv("GOOGLE_API_KEY")
-            if not google_api_key:
-                raise Exception("GOOGLE_API_KEY must be provided")
-
-        elif is_ollama_selected():
-            ollama_model = os.getenv("OLLAMA_MODEL")
-            if not ollama_model:
-                raise Exception("OLLAMA_MODEL must be provided")
-
-            if ollama_model not in SUPPORTED_OLLAMA_MODELS:
-                raise Exception(f"Model {ollama_model} is not supported")
-
-            print("-" * 50)
-            print("Pulling model: ", ollama_model)
-            async for event in pull_ollama_model(ollama_model):
-                print(event)
-            print("Pulled model: ", ollama_model)
-            print("-" * 50)
-
-        elif is_custom_llm_selected():
-            custom_model = os.getenv("CUSTOM_MODEL")
-            custom_llm_url = os.getenv("CUSTOM_LLM_URL")
-            custom_llm_api_key = os.getenv("CUSTOM_LLM_API_KEY")
-            if not custom_model:
-                raise Exception("CUSTOM_MODEL must be provided")
-            if not custom_llm_url:
-                raise Exception("CUSTOM_LLM_URL must be provided")
-            if not custom_llm_api_key:
-                raise Exception("CUSTOM_LLM_API_KEY must be provided")
-            print("-" * 50)
-            print("Selecting model: ", custom_model)
-            models = await list_available_custom_models(
-                custom_llm_url, custom_llm_api_key
-            )
-            print("Available models: ", models)
-            print("-" * 50)
-            if custom_model not in models:
-                raise Exception(f"Model {custom_model} is not available")
+from utils.check_llm_model_availability import check_llm_model_availability
 
 
 @asynccontextmanager

servers/fastapi/api/main.py

@@ -1,8 +1,11 @@
 from fastapi import FastAPI
+from fastapi.staticfiles import StaticFiles
 from api.v1.ppt.router import API_V1_PPT_ROUTER
 from api.lifespan import app_lifespan
 
 
 APP = FastAPI(lifespan=app_lifespan)
 
+APP.mount("/static", StaticFiles(directory="static"), name="static")
+
 APP.include_router(API_V1_PPT_ROUTER)

servers/fastapi/api/v1/ppt/endpoints/files.py

@@ -10,7 +10,7 @@ from services import TEMP_FILE_SERVICE
 from services.documents_loader import DocumentsLoader
 from utils.validators import validate_files
 
-FILES_ROUTER = APIRouter(prefix="/files")
+FILES_ROUTER = APIRouter(prefix="/files", tags=["Files"])
 
 
 @FILES_ROUTER.post("/upload", response_model=List[str])

servers/fastapi/api/v1/ppt/endpoints/images.py

@@ -0,0 +1,9 @@
+from typing import Annotated, List
+from fastapi import APIRouter, File, UploadFile
+
+IMAGES_ROUTER = APIRouter(prefix="/images", tags=["Images"])
+
+
+@IMAGES_ROUTER.post("/upload")
+async def upload_images(images: Annotated[List[UploadFile], File()]):
+    pass

servers/fastapi/api/v1/ppt/endpoints/outlines.py

@@ -8,7 +8,7 @@ from models.sse_response import SSECompleteResponse, SSEResponse, SSEStatusRespo
 from services.database import get_sql_session
 from utils.llm_calls.generate_presentation_outlines import generate_ppt_outline
 
-OUTLINES_ROUTER = APIRouter(prefix="/outlines")
+OUTLINES_ROUTER = APIRouter(prefix="/outlines", tags=["Outlines"])
 
 
 @OUTLINES_ROUTER.get("/stream")

servers/fastapi/api/v1/ppt/endpoints/presentation.py

@@ -19,9 +19,9 @@ from utils.llm_calls.generate_document_summary import generate_document_summary
 from utils.llm_calls.generate_presentation_structure import (
     generate_presentation_structure,
 )
-from utils.llm_calls.generate_slide import get_slide_content_from_type_and_outline
+from utils.llm_calls.generate_slide_content import get_slide_content_from_type_and_outline
 
-PRESENTATION_ROUTER = APIRouter(prefix="/presentation")
+PRESENTATION_ROUTER = APIRouter(prefix="/presentation", tags=["Presentation"])
 
 
 @PRESENTATION_ROUTER.post("/create", response_model=PresentationModel)

servers/fastapi/api/v1/ppt/router.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter
 
 from api.v1.ppt.endpoints.files import FILES_ROUTER
+from api.v1.ppt.endpoints.images import IMAGES_ROUTER
 from api.v1.ppt.endpoints.outlines import OUTLINES_ROUTER
 from api.v1.ppt.endpoints.presentation import PRESENTATION_ROUTER
 
@@ -10,3 +11,4 @@ API_V1_PPT_ROUTER = APIRouter(prefix="/api/v1/ppt")
 API_V1_PPT_ROUTER.include_router(FILES_ROUTER)
 API_V1_PPT_ROUTER.include_router(OUTLINES_ROUTER)
 API_V1_PPT_ROUTER.include_router(PRESENTATION_ROUTER)
+API_V1_PPT_ROUTER.include_router(IMAGES_ROUTER)

servers/fastapi/models/image_prompt.py

@@ -0,0 +1,10 @@
+from typing import Optional
+from pydantic import BaseModel
+
+
+class ImagePrompt(BaseModel):
+    prompt: str
+    theme_prompt: Optional[str] = None
+
+    def get_image_prompt(self, with_theme: bool = False) -> str:
+        return f"{self.prompt}, {self.theme_prompt}" if with_theme else self.prompt

servers/fastapi/services/image_generation_service.py

@@ -0,0 +1,100 @@
+import asyncio
+import os
+import uuid
+import aiohttp
+from google import genai
+from google.genai.types import GenerateContentConfig
+from models.image_prompt import ImagePrompt
+from utils.get_env import get_pexels_api_key_env
+from utils.llm_provider import (
+    get_llm_client,
+    is_google_selected,
+    is_openai_selected,
+)
+from utils.randomizers import get_random_uuid
+
+
+class ImageGenerationService:
+
+    def __init__(self, output_directory: str):
+        self.output_directory = output_directory
+        os.makedirs(output_directory, exist_ok=True)
+
+        self.use_pexels = get_pexels_api_key_env() is not None
+        self.image_gen_func = self.get_image_gen_func()
+
+    def get_image_gen_func(self):
+        if self.use_pexels:
+            return self.get_image_from_pexels
+        elif is_google_selected():
+            return self.generate_image_google
+        elif is_openai_selected():
+            return self.generate_image_openai
+        return None
+
+    async def generate_image(self, prompt: ImagePrompt) -> str:
+        if not self.image_gen_func:
+            print("No image generation function found. Using placeholder image.")
+            return "static/images/placeholder.jpg"
+
+        image_prompt = prompt.get_image_prompt(not self.use_pexels)
+        print(f"Request - Generating Image for {image_prompt}")
+
+        try:
+            image_path = await self.image_gen_func(image_prompt, self.output_directory)
+            if image_path and os.path.exists(image_path):
+                return image_path
+            raise Exception(f"Image not found at {image_path}")
+
+        except Exception as e:
+            print(f"Error generating image: {e}")
+            return "static/images/placeholder.jpg"
+
+    async def generate_image_openai(prompt: str, output_directory: str) -> str:
+        client = get_llm_client()
+        result = await client.images.generate(
+            model="dall-e-3",
+            prompt=prompt,
+            n=1,
+            quality="standard",
+            size="1024x1024",
+        )
+        image_url = result.data[0].url
+        async with aiohttp.ClientSession() as session:
+            async with session.get(image_url) as response:
+                image_bytes = await response.read()
+                image_path = os.path.join(output_directory, f"{get_random_uuid()}.jpg")
+                with open(image_path, "wb") as f:
+                    f.write(image_bytes)
+                return image_path
+
+    async def generate_image_google(prompt: str, output_directory: str) -> str:
+        client = genai.Client()
+        response = await asyncio.to_thread(
+            client.models.generate_content,
+            model="gemini-2.0-flash-preview-image-generation",
+            contents=[prompt],
+            config=GenerateContentConfig(response_modalities=["TEXT", "IMAGE"]),
+        )
+
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                print(part.text)
+            elif part.inline_data is not None:
+                image_path = os.path.join(output_directory, f"{str(uuid.uuid4())}.jpg")
+                with open(image_path, "wb") as f:
+                    f.write(part.inline_data.data)
+
+        return image_path
+
+    async def get_image_from_pexels(prompt: str, output_directory: str) -> str:
+        async with aiohttp.ClientSession() as session:
+            response = await session.get(
+                f"https://api.pexels.com/v1/search?query={prompt}&per_page=1",
+                headers={"Authorization": f'{os.getenv("PEXELS_API_KEY")}'},
+            )
+            data = await response.json()
+            image_url = data["photos"][0]["src"]["large"]
+            image_path = os.path.join(output_directory, f"{str(uuid.uuid4())}.jpg")
+            await download_file(image_url, image_path)
+            return image_path