Update TextGen/router.py

TextGen/router.py

@@ -15,11 +15,11 @@ from langchain_google_genai import (
 )
 from TextGen import app
 from gradio_client import Client
-
+from typing import List
 
 class Message(BaseModel):
     npc: str | None  = None
-    input: str | None = None
+    messages: List[str] | None = None
     
 class VoiceMessage(BaseModel):
     npc: str | None  = None
@@ -42,27 +42,26 @@ main_npcs={
 class Generate(BaseModel):
     text:str
 
-def generate_text(prompt: str):
-    if prompt == "":
-        return {"detail": "Please provide a prompt."}
-    else:
-        prompt = PromptTemplate(template=prompt, input_variables=['Prompt'])
-
-        # Initialize the LLM
-        llm = ChatGoogleGenerativeAI(
-            model="gemini-pro",
-            safety_settings={
-                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-            },
-        )
-
-        llmchain = LLMChain(
-            prompt=prompt,
-            llm=llm
-        )
-
-        llm_response = llmchain.run({"Prompt": prompt})
-        return Generate(text=llm_response)
+def generate_text(messages: List[str]):
+    print(messages)
+    promptmessages[-1]
+    prompt = PromptTemplate(template=prompt, input_variables=['Prompt'])
+
+    # Initialize the LLM
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-pro",
+        safety_settings={
+            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+        },
+    )
+
+    llmchain = LLMChain(
+        prompt=prompt,
+        llm=llm
+    )
+
+    llm_response = llmchain.run({"Prompt": prompt})
+    return Generate(text=llm_response)