Update TextGen/router.py

TextGen/router.py

@@ -88,6 +88,11 @@ class Generate(BaseModel):
     text:str
 
 class Invoke(BaseModel):
+    rooms:List
+    room_of_interest:List
+    index_exit:int
+
+class Rooms(BaseModel):
     system_prompt:str
     message:str
 
@@ -158,9 +163,17 @@ def api_home():
 def inference(message: Message):
     return generate_text(messages=message.messages, npc=message.npc)
 
-@app.post("/invoke_model", response_model=Generate)
+@app.post("/invoke_model")
 def story(prompt: Invoke):
     return inference_model(system_messsage=prompt.system_prompt,prompt=prompt.message)
+    
+@app.post("/generate_level")
+def placement(input: Room):
+    print(input)
+    answer={
+        "key":"value"
+    }
+    return answer
 
 #Dummy function for now
 def determine_vocie_from_npc(npc,genre):