Update TextGen/router.py

TextGen/router.py

@@ -9,7 +9,7 @@ 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 TextGen.diffusion import generate_image
 #from coqui import predict
 from langchain_google_genai import (
     ChatGoogleGenerativeAI,
@@ -73,6 +73,11 @@ main_npcs={
     "Herbalist":"./voices/female.mp3",
     "Bard":"./voices/Bard_voice.mp3"
 }
+main_npcs_elevenlabs={
+    "Blacksmith":"",
+    "Herbalist":"",
+    "Bard":""
+}
 main_npc_system_prompts={
     "Blacksmith":"You are a blacksmith in a video game",
     "Herbalist":"You are an herbalist in a video game",
@@ -82,6 +87,10 @@ main_npc_system_prompts={
 class Generate(BaseModel):
     text:str
 
+class Invoke(BaseModel):
+    system_prompt:str
+    message:str
+
 def generate_text(messages: List[str], npc:str):
     print(npc)
     if npc in main_npcs:
@@ -123,6 +132,24 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+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():
     return {'detail': 'Everchanging Quest backend, nothing to see here'}
@@ -131,6 +158,10 @@ def api_home():
 def inference(message: Message):
     return generate_text(messages=message.messages, npc=message.npc)
 
+@app.post("/invoke_model", response_model=Generate)
+def story(prompt: Invoke):
+    return inference_model(system_messsage=prompt.system_prompt,prompt=prompt.message)
+
 #Dummy function for now
 def determine_vocie_from_npc(npc,genre):
     if npc in main_npcs:
@@ -142,7 +173,17 @@ def determine_vocie_from_npc(npc,genre):
             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:
+        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"    
 
 @app.post("/generate_wav")
 async def generate_wav(message: VoiceMessage):
@@ -234,15 +275,15 @@ async def generate_song():
     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)
-    
-
-    # 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()
-
+#@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")
+#    return Response(content=img_byte_arr, media_type="image/png")