Update app/main.py

app/main.py

@@ -1,30 +1,15 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
-from app import model as ai_model  # safe import
+import app.model as ai_model
 
 app = FastAPI()
 
-# Load model
 model, tokenizer = ai_model.load_model()
 
-class QueryRequest(BaseModel):
+class PromptRequest(BaseModel):
     prompt: str
-    max_new_tokens: int = 512
-    temperature: float = 0.7
-    top_p: float = 0.95
 
 @app.post("/generate")
-async def generate_text(req: QueryRequest):
-    result = ai_model.generate_response(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=req.prompt,
-        max_new_tokens=req.max_new_tokens,
-        temperature=req.temperature,
-        top_p=req.top_p
-    )
-    return {"response": result}
-
-@app.get("/")
-def health_check():
-    return {"status": "ok", "message": "AyurEze Agent Running"}
+async def generate(prompt_request: PromptRequest):
+    response = ai_model.generate_response(model, tokenizer, prompt_request.prompt)
+    return {"response": response}