Update TextGen/router.py

TextGen/router.py

@@ -16,6 +16,12 @@ from langchain_google_genai import (
 from TextGen import app
 from gradio_client import Client
 
+
+class Message(BaseModel):
+    npc: str | None  = None
+    input: str | None = None
+
+
 song_base_api=os.environ["VERCEL_API"]
 
 my_hf_token=os.environ["HF_TOKEN"]
@@ -64,8 +70,8 @@ def api_home():
     return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}
 
 @app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=Generate)
-def inference(input_prompt: str):
-    return generate_text(prompt=input_prompt)
+def inference(message: Message):
+    return generate_text(prompt=message.input)
 
 
 @app.get("/generate_wav")