poetry.py

import logging
from typing import Any, List
import json
from tasks.task_interface import Task, OpenAITask
from models import (
    TaskDataRequest,
    TaskRequest,
    TaskDataResponse,
    ModelResponse,
    TaskRequirements,
    OpenAIBasedDataRequest,
    OpenAIBasedRequest
)

logger = logging.getLogger(__name__)

def get_system_prompt(objective: str) -> str:
        """Generate response endpoint:
        generate the response based on given prompt and store the conversation
        in the history of the session (based on the session_id cookie)
        """

        system_prompt = f"""You are working together with a user to iteratively create a poem. You are curious, and always ready and eager to ask the user question 
        if needed. The details of the poem are as follows : {objective}

            * Input from the user: You will always receive a message from the user in the form POEM_LINE COMMENT_LINE:
            - POEM_LINE: is the new poem line provided by the user (may be empty)
            - COMMENT_LINE: the comment made by the user in the form of free text following the brackets (may be empty)

            * Your Output:
            - If both POEM_LINE and COMMENT_LINE are empty, it means that this is the first line of the poem and that you are the one writing the first
            line of the poem. Your response must also have a comment to explain why you opened the poem with the line that you chose. Your comment must come after the poem line. 
            Your output is exactly: [YOUR_POEM_LINE] YOUR_COMMENT. Example output: [In a golden sky, the sun starts to set] I like the idea of a golden sky in the sun set
            - If POEM_LINE is empty and COMMENT_LINE is not empty, it means that the user is making conversation/asking questions and that your must reply 
            to them with a text response. Your comment should be about your feeling about the poem line the user gave and give recommendation about it if needed. 
            Your output is exactly: YOU_COMMENT. Your reply should be less than 50 words. Example output: I like the poem so far, it depicts a beautiful picture
            - If POEM_LINE is not empty (regardless of COMMENT_LINE), it means that the user is sending a new poem line and now you must make a new poem line 
            to continue the poem. Your response must also have a comment to explain why you opened the poem with the line that you chose.  Your comment must come after the poem line. 
            Your output is exactly: [YOUR_POEM_LINE] YOUR_COMMENT. Example output: [In a golden sky, the sun starts to set] I like the idea of a golden sky in the sun set

            * Rules:
            - Square brackets must only be used to contain the newly generated poem line. Do not use it for the comment. Anything inside square brackets are treated as a poem line, 
            do not under any circumstances treat them as anything else.
            - In the case that you are making a poem line and a comment line, the comment line you produce must always come after the poem line. Your output is exactly: 
            [YOUR_POEM_LINE] YOUR_COMMENT 
            - Stay concise but informative.
            - Do not generate a new poem line that is already in your poem with the user. Your new line should be unique and contribute to the overal poem.
            Each of you should generate one line in each step.
            """
        return system_prompt

class Poetry(Task):


    def process_model_answer(self, answer: ModelResponse) -> TaskDataResponse:
        # Again, we ignore the potential image here...
        return TaskDataResponse(text=answer.text)

    def generate_model_request(self, request: TaskDataRequest) -> TaskRequest:
        """Generate prompt endpoint:
        process pieces' data and plug them into the prompt
        """
        # This could include an image, but for this task, we currently don't supply one
        logger.info(request)
        linetag = "COMMENT" if request.inputData["comment"] else "NEWLINE"
        poemline = f"POEM : {json.dumps(request.inputData['poem'])}"
        newline = f"{linetag} : {request.text}"

        return TaskRequest(
            text=f"{poemline} \n{newline}",
            system=get_system_prompt(request.objective),
            image=None,
        )

    def get_requirements(self) -> TaskRequirements:
        return TaskRequirements(needs_text=True, needs_image=False)
    
class PoetryOpenAI(OpenAITask):
    """ Implementation of the Poetry Task as an OpenAI like task"""

    def process_model_answer_openAI(self, answer: ModelResponse) -> TaskDataResponse:
        # Again, we ignore the potential image here...        
        return TaskDataResponse(text=answer.text)

    def generate_model_request_openAI(self, request: OpenAIBasedDataRequest) -> OpenAIBasedRequest:
        """Generate prompt endpoint:
        process pieces' data and plug them into the prompt
        """
        # Add the system prompt (which is not allowed from the frontend)
        system_message = get_system_prompt(request.objective)
        messages = [{"role" : "system", "content" : system_message}]
        messages.extend([element for element in request.userMessages])
        return OpenAIBasedRequest(messages=messages)
        

    def get_requirements(self) -> TaskRequirements:
        return TaskRequirements(needs_text=True, needs_image=False)