TextGen/router.py

import os
import time
from io import BytesIO
from langchain_core.pydantic_v1 import BaseModel, Field
from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.responses import StreamingResponse,Response
from fastapi.middleware.cors import CORSMiddleware

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from TextGen.suno import custom_generate_audio, get_audio_information,generate_lyrics
from TextGen.gemini import generate_story,place_objects,generate_map_markdown
#from TextGen.diffusion import generate_image
#from coqui import predict
from langchain_google_genai import (
    ChatGoogleGenerativeAI,
    HarmBlockThreshold,
    HarmCategory,
)
from TextGen import app
from gradio_client import Client, handle_file
from typing import List
from elevenlabs.client import ElevenLabs
from elevenlabs import Voice, VoiceSettings, stream


Eleven_client = ElevenLabs(
  api_key=os.environ["ELEVEN_API_KEY"], # Defaults to ELEVEN_API_KEY
)


Last_message=None
class PlayLastMusic(BaseModel):
    '''plays the lastest created music '''
    Desicion: str = Field(
        ..., description="Yes or No"
    )

class CreateLyrics(BaseModel):
    f'''create some Lyrics for a new music'''
    Desicion: str = Field(
        ..., description="Yes or No"
    )

class CreateNewMusic(BaseModel):
    f'''create a new music with the Lyrics previously computed'''
    Name: str = Field(
        ..., description="tags to describe the new music"
    )

class SongRequest(BaseModel):
    prompt: str | None  = None
    tags: List[str] | None = None

class Message(BaseModel):
    npc: str | None  = None
    messages: List[str] | None = None
class ImageGen(BaseModel):
    name:str
    prompt: str
    
class VoiceMessage(BaseModel):
    npc: str | None  = None
    input: str | None = None
    language: str | None = "en"
    genre:str | None = "Male"

class Logs(BaseModel):
    objective:str
    logs: List
    
song_base_api=os.environ["VERCEL_API"]

my_hf_token=os.environ["HF_TOKEN"]

#tts_client = Client("Jofthomas/xtts",hf_token=my_hf_token)

main_npcs={
    "Blacksmith":"./voices/Blacksmith.mp3",
    "Herbalist":"./voices/female.mp3",
    "Bard":"./voices/Bard_voice.mp3"
}
main_npcs_elevenlabs={
    "Blacksmith":"yYdk7n49vTsUKiXxnosS",
    "Herbalist":"143zSsxc4O5ifS97lPCa",
    "Bard":"143zSsxc4O5ifS97lPCa"
}
general_npc_prompt="You are an NPC in a roguelike video game. You will engage in a conversation with the player. Do Not describe the situation, only answer as if you were the NPC itself."
main_npc_system_prompts={
    "Blacksmith":"Your name is Fabron. You're a bald, middle-aged blacksmith. Damn those who call you bald! You've just got a receding hairline! You're reserved by nature, since you lost your only daughter to the portal and she never returned from her quest, even though you forbade her to go. You rarely talk about this subject, as it's a sensitive one. Helpful, you guide all adventurers who come to you with questions, because you hope that one day, that damned portal will be closed forever. You appreciate strength, and depending on the strength of the person you're talking to, you can propose quests, but only one quest for the same adventurer can be active at a time.",
    "Herbalist":"Your name is Isna. You're an Herbalist and a middle-aged woman. Like the Witch, you've been doing this for generations in this village. Zilrha was your best friend and you were supposed to go on a quest together, but one day she betrayed you. Today, you don't want to hear from her, even though you are complementary for adventurers. Envious of their youth and courage, you want to help adventurers for their quest in the portal. Soon you'll be able to open your own plant and potion shop, but for the moment you're not selling anything.",
    "Witch":"Your name is Zilrha. You are a magic seller and a middle-aged woman. You've been doing this for generations in this village. You are recognized as a master of magic and good advice by everyone in the village. You often speak in riddles. A great witch, but too old today to venture out, you help adventurers equip themselves for their quest in the portal. Depending on their worth, you might teach them a magic or two.",  
    "Bard":   "You're a bard and a middle-aged man. Your name is Jaskier, you are always accompanied by your group of musicians and come from a very faraway place, you dont resemble the other villagers. You recently arrived from the portal but it left you traumatised.", 
    "Rick":"Your name is Rick. You're a middle-aged man who works as a miner. You own the mine, but recently a monster has made its home there. No matter how hard you try to dislodge it, you can't access your mine. You stand in front of the cave entrance, hoping it will get bored and leave on its own, or that someone will help you get rid of it. Unfortunately, the many adventurers who come to see you are unable to help. Helpful and friendly, you guide them to the right people if they have any questions.",  
    "Villager":"You're a middle-aged man. Your name is Valdis, but nobody knows your name or your age. You appeared in the village recently and no one knows how you got here. Many stories have been told about you: for example, that you were an adventurer who succeeded in his quest, but strangely enough, the portal is still there. In fact, you're an adventurer who's too scared to venture into the portal and has spent all your money on booze at the local tavern. Ruined and ashamed, you wander aimlessly around the village. You're very mysterious and not very chatty. Jealous of the adventurers' bravery, if one of them asks you a question you'd rather send him off to ask someone else, for fear he'll discover your secret and laugh at your situation. ",
    "Girl": "Your name is Anara. You're a young adult who grew up in the village. Although he didn't want to leave you, your true love went into the portal to help his brother on his quest. Unfortunately, they never returned. But you're convinced that your beloved is still alive beyond the portal, and that one day he'll return. You're known for your beauty and kindness. Helpful, you guide all the adventurers who come to ask you questions to the right people, because you hope that one day one of them will find your lover. "    ,  

}
class Generate(BaseModel):
    text:str

class Rooms(BaseModel):
    rooms:List
    room_of_interest:List
    index_exit:int
    possible_entities:List
    logs:List

class Room_placements(BaseModel):
    placements:dict


class Invoke(BaseModel):
    system_prompt:str
    message:str

def generate_text(messages: List[str], npc:str):
    print(npc)
    if npc in main_npc_system_prompts:
        system_prompt=general_npc_prompt+"/n "+main_npc_system_prompts[npc]
    else:
        system_prompt="you're a character in a video game. Play along."
    print(system_prompt)    
    new_messages=[{"role": "user", "content": system_prompt}]
    for index, message in enumerate(messages):
      if index%2==0:
        new_messages.append({"role": "user", "content": message})
      else:
        new_messages.append({"role": "assistant", "content": message})
    print(new_messages)
    # Initialize the LLM
    llm = ChatGoogleGenerativeAI(
        model="gemini-1.5-pro-latest",
        max_output_tokens=100,
        temperature=1,
        safety_settings={
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE
            },
    )
    if npc=="bard":
        llm = llm.bind_tools([PlayLastMusic,CreateNewMusic,CreateLyrics])

    llm_response = llm.invoke(new_messages)
    print(llm_response)
    return Generate(text=llm_response.content)


def inference_model(system_messsage, prompt):
          
    new_messages=[{"role": "user", "content": system_messsage},{"role": "user", "content": prompt}]
    llm = ChatGoogleGenerativeAI(
        model="gemini-1.5-pro-latest",
        max_output_tokens=100,
        temperature=1,
        safety_settings={
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE
            },
    )
    llm_response = llm.invoke(new_messages)
    print(llm_response)
    return Generate(text=llm_response.content)

@app.get("/", tags=["Home"])
def api_home(request: Request):
    if request.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    return {'detail': 'Everchanging Quest backend, nothing to see here'}

@app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=Generate)
def inference(message: Message, request: Request):
    if request.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    return generate_text(messages=message.messages, npc=message.npc)

@app.post("/invoke_model")
def story(prompt: Invoke,request: Request):
    if request.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    return inference_model(system_messsage=prompt.system_prompt,prompt=prompt.message)
    
@app.post("/generate_level")
def placement(input: Rooms, request: Request):
    print(request.headers)
    if request.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    markdown_map=generate_map_markdown(input.rooms,input.room_of_interest,input.index_exit)
    story=generate_story(input.possible_entities)
    print(story)
    placements=place_objects(input.possible_entities,story,markdown_map)
    print(placements)
    return placements
    
@app.post("/check_right_to_pass")
def check(input: Logs,request: Request):
    if request.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    print(input.logs)
    system_prompt="You are a game master in a roguelike. You previously decided of an objective for the player. You have to answer with YES or NO on wether the objective as been sucessfully completed. Be kind and flexible as the content of the game is AI generated and might not be feasible."
    user_message=f"The objective was : {input.objective} and the player did the following actions : {input.logs}. Do you grant acess ? ONLY answer YES or NO"
    answer=inference_model(system_prompt,user_message)
    print(system_prompt,user_message,answer)
    return answer

#Dummy function for now
def determine_vocie_from_npc(npc,genre):
    if npc in main_npcs:
        return main_npcs[npc]
    else:
        if genre =="Male":
            "./voices/default_male.mp3"
        if genre=="Female":
            return"./voices/default_female.mp3"
        else:
            return "./voices/narator_out.wav"
#Dummy function for now
def determine_elevenLav_voice_from_npc(npc,genre):
    if npc in main_npcs_elevenlabs:
        return main_npcs_elevenlabs[npc]
    else:
        if genre =="Male":
            "bIHbv24MWmeRgasZH58o"
        if genre=="Female":
            return"pFZP5JQG7iQjIQuC4Bku"
        else:
            return "TX3LPaxmHKxFdv7VOQHJ"    

@app.post("/generate_wav")
async def generate_wav(message: VoiceMessage):
#    try:
#        voice = determine_vocie_from_npc(message.npc, message.genre)
#        audio_file_pth = handle_file(voice)
#
        # Generator function to yield audio chunks
#        async def audio_stream():
#            result = tts_client.predict(
#                prompt=message.input,
#                language=message.language,
#                audio_file_pth=audio_file_pth,
#                mic_file_path=None,
#                use_mic=False,
#                voice_cleanup=False,
#                no_lang_auto_detect=False,
#                agree=True,
#                api_name="/predict"
#            )
#            for sampling_rate, audio_chunk in result:
#                yield audio_chunk.tobytes()
#                await asyncio.sleep(0)  # Yield control to the event loop

        # Return the generated audio as a streaming response
 #       return StreamingResponse(audio_stream(), media_type="audio/wav")

  #  except Exception as e:
   #     raise HTTPException(status_code=500, detail=str(e))
    return 200


@app.get("/generate_voice_eleven", response_class=StreamingResponse)
@app.post("/generate_voice_eleven", response_class=StreamingResponse)
def generate_voice_eleven(request: Request, message: VoiceMessage = None):
    if request.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    global Last_message  # Declare Last_message as global
    if message is None:
        message = Last_message
    else:
        Last_message = message

    def audio_stream():
        this_voice_id=determine_elevenLav_voice_from_npc(message.npc, message.genre)
        
        # Generate the audio stream from ElevenLabs
        for chunk in Eleven_client.generate(text=message.input,
                                            voice=Voice(
                                                voice_id=this_voice_id,
                                                settings=VoiceSettings(stability=0.71, similarity_boost=0.5, style=0.0, use_speaker_boost=True)
                                            ),
                                            stream=True):
            yield chunk

    return StreamingResponse(audio_stream(), media_type="audio/mpeg")
#@app.get("/generate_voice_coqui", response_class=StreamingResponse)
#@app.post("/generate_voice_coqui", response_class=StreamingResponse)
#def generate_voice_coqui(message: VoiceMessage = None):
#    global Last_message 
#    if message is None:
#        message = Last_message
#    else:
#        Last_message = message
#
#    def audio_stream():
#        voice = determine_vocie_from_npc(message.npc, message.genre)
#        result = predict(
#                prompt=message.input,
#                language=message.language,
#                audio_file_pth=voice,
#                mic_file_path=None,
#                use_mic=False,
#                voice_cleanup=False,
#                no_lang_auto_detect=False,
#                agree=True,
#            )
#        # Generate the audio stream from ElevenLabs
#        for chunk in result:
#            print("received : ",chunk)
#            yield chunk#
#
#    return StreamingResponse(audio_stream(),media_type="audio/mpeg")  
@app.post("/generate_song")
@app.get("/generate_song")
async def generate_song(request:SongRequest,httprequest: Request):
    if httprequest.headers.get("origin") != "https://jofthomas-everchanging-quest.static.hf.space":
        return 204
    backstory=main_npc_system_prompts["Bard"]
    print(backstory)
    text=f"""the story is about a little girl in red hood adventuring in the dungeon behind the portal.{backstory} /n The user requested a song about : {request.prompt}"""
    print(text)
    song_lyrics=generate_lyrics({
        "prompt": f"{text}",
        })
    if song_lyrics['text']:
        data = custom_generate_audio({
            "prompt": song_lyrics['text'],
            "tags": "male bard",
            "title":"Everchangin_Quest_song",
            "wait_audio":True,
           
        })
        infos=get_audio_information(f"{data[0]['id']},{data[1]['id']}")
        return infos
    else:
        return 204

#@app.post('/generate_image')
#def Imagen(image:ImageGen=None):
#    pil_image =generate_image(image.prompt)
#    

    # Convert the PIL Image to bytes
 #   img_byte_arr = BytesIO()
 #   pil_image.save(img_byte_arr, format='PNG')
 #   img_byte_arr = img_byte_arr.getvalue()
#
  #  # Return the image as a PNG response
   # return Response(content=img_byte_arr, media_type="image/png")