Update TextGen/router.py

TextGen/router.py

@@ -3,14 +3,14 @@ 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.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.diffusion import generate_image
-#from coqui import predict
+from TextGen.suno import custom_generate_audio, get_audio_information, generate_lyrics
+# from TextGen.diffusion import generate_image
+# from coqui import predict
 from langchain_google_genai import (
     ChatGoogleGenerativeAI,
     HarmBlockThreshold,
@@ -22,13 +22,12 @@ 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
+    api_key=os.environ["ELEVEN_API_KEY"],  # Defaults to ELEVEN_API_KEY
 )
 
+Last_message = None
 
-Last_message=None
 class PlayLastMusic(BaseModel):
     '''plays the lastest created music '''
     Desicion: str = Field(
@@ -36,85 +35,88 @@ class PlayLastMusic(BaseModel):
     )
 
 class CreateLyrics(BaseModel):
-    f'''create some Lyrics for a new music'''
+    '''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'''
+    '''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
+    prompt: str | None = None
     tags: List[str] | None = None
 
 class Message(BaseModel):
-    npc: str | None  = None
+    npc: str | None = None
     messages: List[str] | None = None
+
 class ImageGen(BaseModel):
-    prompt: str | None  = None  
+    prompt: str | None = None
+
 class VoiceMessage(BaseModel):
-    npc: str | None  = None
+    npc: str | None = None
     input: str | None = None
     language: str | None = "en"
-    genre:str | None = "Male"
-    
-song_base_api=os.environ["VERCEL_API"]
+    genre: str | None = "Male"
 
-my_hf_token=os.environ["HF_TOKEN"]
+song_base_api = os.environ["VERCEL_API"]
+my_hf_token = os.environ["HF_TOKEN"]
 
-#tts_client = Client("Jofthomas/xtts",hf_token=my_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 = {
+    "Blacksmith": "./voices/Blacksmith.mp3",
+    "Herbalist": "./voices/female.mp3",
+    "Bard": "./voices/Bard_voice.mp3"
 }
-main_npcs_elevenlabs={
-    "Blacksmith":"yYdk7n49vTsUKiXxnosS",
-    "Herbalist":"143zSsxc4O5ifS97lPCa",
-    "Bard":"143zSsxc4O5ifS97lPCa"
+
+main_npcs_elevenlabs = {
+    "Blacksmith": "yYdk7n49vTsUKiXxnosS",
+    "Herbalist": "143zSsxc4O5ifS97lPCa",
+    "Bard": "143zSsxc4O5ifS97lPCa"
 }
-main_npc_system_prompts={
-    "Blacksmith":"You are a blacksmith in a video game",
-    "Herbalist":"You are an herbalist in a video game",
-    "Witch":"You are a witch in a video game. You are disguised as a potion seller in a small city where adventurers come to challenge the portal. You are selling some magic spells in a UI that the player only sees. Don't event too much lore and just follow the standard role of a merchant.",
-    "Bard":"You are a bard in a video game"
+
+main_npc_system_prompts = {
+    "Blacksmith": "You are a blacksmith in a video game",
+    "Herbalist": "You are an herbalist in a video game",
+    "Witch": "You are a witch in a video game. You are disguised as a potion seller in a small city where adventurers come to challenge the portal. You are selling some magic spells in a UI that the player only sees. Don't event too much lore and just follow the standard role of a merchant.",
+    "Bard": "You are a bard in a video game"
 }
+
 class Generate(BaseModel):
-    text:str
+    text: str
 
 class Rooms(BaseModel):
-    rooms:List
-    room_of_interest:List
-    index_exit:int
-    possible_entities:List
-    logs:List
+    rooms: List
+    room_of_interest: List
+    index_exit: int
+    possible_entities: List
+    logs: List
 
 class Room_placements(BaseModel):
-    placements:dict
-
+    placements: dict
 
 class Invoke(BaseModel):
-    system_prompt:str
-    message:str
+    system_prompt: str
+    message: str
 
-def generate_text(messages: List[str], npc:str):
+def generate_text(messages: List[str], npc: str):
     print(npc)
     if npc in main_npcs:
-        system_prompt=main_npc_system_prompts[npc]
+        system_prompt = 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}]
+        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})
+        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(
@@ -122,14 +124,14 @@ def generate_text(messages: List[str], npc:str):
         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
-            },
+            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])
+    if npc == "bard":
+        llm = llm.bind_tools([PlayLastMusic, CreateNewMusic, CreateLyrics])
 
     llm_response = llm.invoke(new_messages)
     print(llm_response)
@@ -144,18 +146,17 @@ app.add_middleware(
 )
 
 def inference_model(system_messsage, prompt):
-          
-    new_messages=[{"role": "user", "content": system_messsage},{"role": "user", "content": 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
-            },
+            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)
@@ -171,71 +172,70 @@ def inference(message: Message):
 
 @app.post("/invoke_model")
 def story(prompt: Invoke):
-    return inference_model(system_messsage=prompt.system_prompt,prompt=prompt.message)
-    
+    return inference_model(system_messsage=prompt.system_prompt, prompt=prompt.message)
+
 @app.post("/generate_level")
 def placement(input: Rooms):
     print(input)
-    markdown_map=generate_map_markdown(input)
+    markdown_map = generate_map_markdown(input)
     print(markdown_map)
-    answer={
-        "key":"value"
+    answer = {
+        "key": "value"
     }
     return answer
 
-#Dummy function for now
-def determine_vocie_from_npc(npc,genre):
+# 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"
+        if genre == "Male":
+            return "./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):
+
+# 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"
+        if genre == "Male":
+            return "bIHbv24MWmeRgasZH58o"
+        if genre == "Female":
+            return "pFZP5JQG7iQjIQuC4Bku"
         else:
-            return "TX3LPaxmHKxFdv7VOQHJ"    
+            return "TX3LPaxmHKxFdv7VOQHJ"
 
-@app.post("/generate_wav")
+@app.post("/generate_wav", response_class=StreamingResponse)
 async def generate_wav(message: VoiceMessage):
-#    try:
-#        voice = determine_vocie_from_npc(message.npc, message.genre)
-#        audio_file_pth = handle_file(voice)
-#
+    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
+        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
+        return StreamingResponse(audio_stream(), media_type="audio/wav")
 
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/generate_voice_eleven", response_class=StreamingResponse)
 @app.post("/generate_voice_eleven", response_class=StreamingResponse)
@@ -247,8 +247,8 @@ def generate_voice_eleven(message: VoiceMessage = None):
         Last_message = message
 
     def audio_stream():
-        this_voice_id=determine_elevenLav_voice_from_npc(message.npc, message.genre)
-        
+        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(
@@ -259,66 +259,66 @@ def generate_voice_eleven(message: VoiceMessage = None):
             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
+
+# @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#
+#     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")   
+#     return StreamingResponse(audio_stream(), media_type="audio/mpeg")
+
 @app.get("/generate_song")
 async def generate_song():
-    text="""You are a bard in a video game singing the tales of a little girl in red hood."""
+    text = """You are a bard in a video game singing the tales of a little girl in red hood."""
 
-    song_lyrics=generate_lyrics({
+    song_lyrics = generate_lyrics({
         "prompt": f"{text}",
-        })
+    })
     data = custom_generate_audio({
         "prompt": song_lyrics['text'],
         "tags": "male bard",
-        "title":"Everchangin_Quest_song",
-        "wait_audio":True,
-       
+        "title": "Everchangin_Quest_song",
+        "wait_audio": True,
     })
-    infos=get_audio_information(f"{data[0]['id']},{data[1]['id']}")
+    infos = get_audio_information(f"{data[0]['id']},{data[1]['id']}")
     return infos
 
-#@app.post('/generate_image')
-#def Imagen(image:ImageGen=None):
-#    pil_image =generate_image(image.prompt)
-#    
+# @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()
+#     # 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")
+#     # Return the image as a PNG response
+#     return Response(content=img_byte_arr, media_type="image/png")
 
 def generate_map_markdown(data):
     import numpy as np
-    
+
     # Define the room structure with walls and markers
     def create_room(room_char):
         return [
@@ -326,22 +326,22 @@ def generate_map_markdown(data):
             f"║ {room_char} ║",
             f"╚═══╝"
         ]
-    
+
     # Extract rooms and rooms of interest
     rooms = [eval(room) for room in data["rooms"]]
     rooms_of_interest = [eval(room) for room in data["room_of_interest"]]
-    
+
     # Determine grid size
     min_x = min(room[0] for room in rooms)
     max_x = max(room[0] for room in rooms)
     min_y = min(room[1] for room in rooms)
     max_y = max(room[1] for room in rooms)
-    
+
     # Create grid with empty spaces represented by a room-like structure
     map_height = (max_y - min_y + 1) * 3
     map_width = (max_x - min_x + 1) * 5
     grid = np.full((map_height, map_width), " ")
-    
+
     # Populate grid with rooms and their characteristics
     for i, room in enumerate(rooms):
         x, y = room
@@ -356,9 +356,9 @@ def generate_map_markdown(data):
         room_structure = create_room(room_char)
         for j, row in enumerate(room_structure):
             grid[y_offset + j, x_offset:x_offset + 5] = list(row)
-    
+
     # Convert grid to a string format suitable for display
     markdown_map = "\n".join("".join(row) for row in grid)
-    
+
     # Return the map wrapped in triple backticks for proper display in markdown
     return f"```\n{markdown_map}\n```"