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)