Open-Assistant/backend/app/api/v1/tasks.py

# -*- coding: utf-8 -*-
import random
from typing import Any
from uuid import UUID

from app.api import deps
from app.models.db_payload import TaskPayload
from app.prompt_repository import PromptRepository
from app.schemas import protocol as protocol_schema
from fastapi import APIRouter, Depends, HTTPException
from fastapi.security.api_key import APIKey
from loguru import logger
from sqlmodel import Session
from starlette.status import HTTP_400_BAD_REQUEST

router = APIRouter()


def generate_task(request: protocol_schema.TaskRequest) -> protocol_schema.Task:
    match (request.type):
        case protocol_schema.TaskRequestType.random:
            logger.info("Frontend requested a random task.")
            while request.type == protocol_schema.TaskRequestType.random:
                request.type = random.choice(list(protocol_schema.TaskRequestType)).value
            return generate_task(request)
        case protocol_schema.TaskRequestType.summarize_story:
            logger.info("Generating a SummarizeStoryTask.")
            task = protocol_schema.SummarizeStoryTask(
                story="This is a story. A very long story. So long, it needs to be summarized.",
            )
        case protocol_schema.TaskRequestType.rate_summary:
            logger.info("Generating a RateSummaryTask.")
            task = protocol_schema.RateSummaryTask(
                full_text="This is a story. A very long story. So long, it needs to be summarized.",
                summary="This is a summary.",
                scale=protocol_schema.RatingScale(min=1, max=5),
            )
        case protocol_schema.TaskRequestType.initial_prompt:
            logger.info("Generating an InitialPromptTask.")
            task = protocol_schema.InitialPromptTask(
                hint="Ask the assistant about a current event."  # this is optional
            )
        case protocol_schema.TaskRequestType.user_reply:
            logger.info("Generating a UserReplyTask.")
            task = protocol_schema.UserReplyTask(
                conversation=protocol_schema.Conversation(
                    messages=[
                        protocol_schema.ConversationMessage(
                            text="Hey, assistant, what's going on in the world?",
                            is_assistant=False,
                        ),
                        protocol_schema.ConversationMessage(
                            text="I'm not sure I understood correctly, could you rephrase that?",
                            is_assistant=True,
                        ),
                    ],
                )
            )
        case protocol_schema.TaskRequestType.assistant_reply:
            logger.info("Generating a AssistantReplyTask.")
            task = protocol_schema.AssistantReplyTask(
                conversation=protocol_schema.Conversation(
                    messages=[
                        protocol_schema.ConversationMessage(
                            text="Hey, assistant, write me an English essay about water.",
                            is_assistant=False,
                        ),
                    ],
                )
            )
        case protocol_schema.TaskRequestType.rank_initial_prompts:
            logger.info("Generating a RankInitialPromptsTask.")
            task = protocol_schema.RankInitialPromptsTask(
                prompts=[
                    "Please write a story about a time you were happy.",
                    "Please write a story about a time you were sad.",
                ]
            )
        case protocol_schema.TaskRequestType.rank_user_replies:
            logger.info("Generating a RankUserRepliesTask.")
            task = protocol_schema.RankUserRepliesTask(
                conversation=protocol_schema.Conversation(
                    messages=[
                        protocol_schema.ConversationMessage(
                            text="Hey, assistant, what's going on in the world?",
                            is_assistant=False,
                        ),
                        protocol_schema.ConversationMessage(
                            text="I'm not sure I understood correctly, could you rephrase that?",
                            is_assistant=True,
                        ),
                    ],
                ),
                replies=[
                    "Oh come oooooon!",
                    "What are the news?",
                ],
            )

        case protocol_schema.TaskRequestType.rank_assistant_replies:
            logger.info("Generating a RankAssistantRepliesTask.")
            task = protocol_schema.RankAssistantRepliesTask(
                conversation=protocol_schema.Conversation(
                    messages=[
                        protocol_schema.ConversationMessage(
                            text="Hey, assistant, what's going on in the world?",
                            is_assistant=False,
                        ),
                    ],
                ),
                replies=[
                    "I'm not sure I understood correctly, could you rephrase that?",
                    "The world is fine. All good.",
                    "Crap is hitting the fan. Start farming.",
                ],
            )
        case _:
            raise HTTPException(
                status_code=HTTP_400_BAD_REQUEST,
                detail="Invalid request type.",
            )

    logger.info(f"Generated {task=}.")

    return task


@router.post("/", response_model=protocol_schema.AnyTask)  # work with Union once more types are added
def request_task(
    *,
    db: Session = Depends(deps.get_db),
    api_key: APIKey = Depends(deps.get_api_key),
    request: protocol_schema.TaskRequest,
) -> Any:
    """
    Create new task.
    """
    api_client = deps.api_auth(api_key, db)

    try:
        task = generate_task(request)

        pr = PromptRepository(db, api_client, request.user)
        pr.store_task(task)

    except Exception:
        logger.exception("Failed to generate task.")
        raise HTTPException(
            status_code=HTTP_400_BAD_REQUEST,
        )
    return task


@router.post("/{task_id}/ack")
def acknowledge_task(
    *,
    db: Session = Depends(deps.get_db),
    api_key: APIKey = Depends(deps.get_api_key),
    task_id: UUID,
    ack_request: protocol_schema.TaskAck,
) -> Any:
    """
    The frontend acknowledges a task.
    """

    api_client = deps.api_auth(api_key, db)

    try:
        pr = PromptRepository(db, api_client, user=None)

        # here we store the post id in the database for the task
        pr.bind_frontend_post_id(task_id=task_id, post_id=ack_request.post_id)
        logger.info(f"Frontend acknowledges task {task_id=}, {ack_request=}.")

    except Exception:
        logger.exception("Failed to acknowledge task.")
        raise HTTPException(
            status_code=HTTP_400_BAD_REQUEST,
        )
    return {}


@router.post("/{task_id}/nack")
def acknowledge_task_failure(
    *,
    db: Session = Depends(deps.get_db),
    api_key: APIKey = Depends(deps.get_api_key),
    task_id: UUID,
    nack_request: protocol_schema.TaskNAck,
) -> Any:
    """
    The frontend reports failure to implement a task.
    """
    deps.api_auth(api_key, db)

    logger.info(f"Frontend reports failure to implement task {task_id=}, {nack_request=}.")
    # here we would store the post id in the database for the task
    return {}


@router.post("/interaction")
def post_interaction(
    *,
    db: Session = Depends(deps.get_db),
    api_key: APIKey = Depends(deps.get_api_key),
    interaction: protocol_schema.AnyInteraction,
) -> Any:
    """
    The frontend reports an interaction.
    """
    api_client = deps.api_auth(api_key, db)

    try:
        pr = PromptRepository(db, api_client, user=interaction.user)

        match type(interaction):
            case protocol_schema.TextReplyToPost:
                logger.info(
                    f"Frontend reports text reply to {interaction.post_id=} with {interaction.text=} by {interaction.user=}."
                )

                work_package = pr.fetch_workpackage_by_postid(interaction.post_id)
                work_payload: TaskPayload = work_package.payload.payload
                logger.info(f"found task work package in db: {work_payload}")

                # here we store the text reply in the database
                # ToDo: role user or agent?
                pr.store_text_reply(interaction, role="unknown")

                return protocol_schema.TaskDone()
            case protocol_schema.PostRating:
                logger.info(
                    f"Frontend reports rating of {interaction.post_id=} with {interaction.rating=} by {interaction.user=}."
                )

                # here we store the rating in the database
                pr.store_rating(interaction)

                return protocol_schema.TaskDone()
            case protocol_schema.PostRanking:
                logger.info(
                    f"Frontend reports ranking of {interaction.post_id=} with {interaction.ranking=} by {interaction.user=}."
                )

                # TODO: check if the ranking is valid
                pr.store_ranking(interaction)
                # here we would store the ranking in the database
                return protocol_schema.TaskDone()
            case _:
                raise HTTPException(
                    status_code=HTTP_400_BAD_REQUEST,
                    detail="Invalid response type.",
                )

    except Exception:
        logger.exception("Interaction request failed.")
        raise HTTPException(
            status_code=HTTP_400_BAD_REQUEST,
        )