Update app.py

app.py

@@ -1,6 +1,5 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.responses import PlainTextResponse
-from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
@@ -13,20 +12,8 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 
 model.eval()
 
-class Message(BaseModel):
-    message: str
 
-@app.get("/")
-async def root():
-    return {"status": "TinyLlama работает"}
-
-from fastapi import Request
-
-@app.post("/")
-async def receive(request: Request):
-
-    data = await request.json()
-    message = data.get("message", "")
+def generate_ai(message: str):
 
     prompt = f"User: {message}\nAssistant:"
 
@@ -35,7 +22,7 @@ async def receive(request: Request):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=20,
+            max_new_tokens=25,
             do_sample=False
         )
 
@@ -44,14 +31,37 @@ async def receive(request: Request):
     if "Assistant:" in reply:
         reply = reply.split("Assistant:")[-1].strip()
 
-    return PlainTextResponse(reply)
+    return reply
 
-    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    if "Assistant:" in reply:
-        reply = reply.split("Assistant:")[-1].strip()
+@app.get("/")
+async def root():
+    return PlainTextResponse("AI server работает")
+
+
+@app.api_route("/", methods=["GET","POST","PUT","PATCH","DELETE","HEAD"])
+async def universal(request: Request):
+
+    try:
+
+        # пробуем получить JSON
+        try:
+            data = await request.json()
+            message = data.get("message", "")
+        except:
+            # если не JSON — читаем обычный текст
+            body = await request.body()
+            message = body.decode("utf-8")
+
+        if not message:
+            message = "Hello"
+
+        reply = generate_ai(message)
+
+        return PlainTextResponse(reply)
 
-    return PlainTextResponse(reply)
+    except Exception as e:
+        return PlainTextResponse(f"ERROR: {str(e)}")
 
 
 if __name__ == "__main__":