Update TextGen/router.py

TextGen/router.py

@@ -39,14 +39,25 @@ main_npcs={
     "Herbalist":"./voices/female.mp3",
     "Bard":"./voices/Bard_voice.mp3"
 }
+main_npc_system_prompts={
+    "Blacksmith":"You are a blacksmith in a video game",
+    "Herbalist":"You are an herbalist in a video game",
+    "Bard":"You are a bard in a video game"
+}
 class Generate(BaseModel):
     text:str
 
-def generate_text(messages: List[str]):
-    print(messages)
-    prompt=messages[-1]
-    prompt = PromptTemplate(template=prompt, input_variables=['Prompt'])
-
+def generate_text(messages: List[str], npc:str):
+    print(npc)
+    system_prompt=main_npc_system_prompts[npc]
+    print(system_prompt)    
+    new_messages=[{"role": "system", "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-pro",
@@ -54,17 +65,9 @@ def generate_text(messages: List[str]):
             HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
         },
     )
-
-    llmchain = LLMChain(
-        prompt=prompt,
-        llm=llm
-    )
-
-    llm_response = llmchain.run({"Prompt": prompt})
+    llm_response = llm.invoke(new_messages)
     return Generate(text=llm_response)
 
-        
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -79,7 +82,7 @@ def api_home():
 
 @app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=Generate)
 def inference(message: Message):
-    return generate_text(messages=message.messages)
+    return generate_text(messages=message.messages, npc=message.npc)
 
 #Dummy function for now
 def determine_vocie_from_npc(npc,genre):