48 lines
2.2 KiB
Python
48 lines
2.2 KiB
Python
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter
|
|
from fastapi.params import Depends
|
|
from starlette.requests import Request
|
|
|
|
from api import service
|
|
from api.dependency import get_generation_service
|
|
|
|
from api.models.GenerationRequest import GenerationResponse, GenerationRequest, PromptResponse, PromptRequest
|
|
from api.service.generation_service import GenerationService
|
|
from models.Generation import Generation
|
|
|
|
router = APIRouter(prefix='/api/generations', tags=["Generation"])
|
|
|
|
|
|
@router.post("/prompt-assistant", response_model=PromptResponse)
|
|
async def ask_prompt_assistant(prompt_request: PromptRequest, request: Request,
|
|
generation_service: GenerationService = Depends(
|
|
get_generation_service)) -> PromptResponse:
|
|
generated_prompt = await generation_service.ask_prompt_assistant(prompt_request.prompt, prompt_request.linked_assets)
|
|
return PromptResponse(prompt=generated_prompt)
|
|
|
|
|
|
@router.get("", response_model=List[GenerationResponse])
|
|
async def get_generations(character_id: Optional[str], limit: int = 10, offset: int = 0,
|
|
generation_service: GenerationService = Depends(get_generation_service)):
|
|
return await generation_service.get_generations(character_id, limit=limit, offset=offset)
|
|
|
|
|
|
@router.post("/_run", response_model=GenerationResponse)
|
|
async def post_generation(generation: GenerationRequest, request: Request,
|
|
generation_service: GenerationService = Depends(
|
|
get_generation_service)) -> GenerationResponse:
|
|
return await generation_service.create_generation_task(generation)
|
|
|
|
|
|
@router.get("/{generation_id}", response_model=GenerationResponse)
|
|
async def get_generation(generation_id: str,
|
|
generation_service: GenerationService = Depends(get_generation_service)) -> GenerationResponse:
|
|
return await generation_service.get_generation(generation_id)
|
|
|
|
|
|
@router.get("/running")
|
|
async def get_running_generations(request: Request,
|
|
generation_service: GenerationService = Depends(get_generation_service)):
|
|
return await generation_service.get_running_generations()
|