Update app.py

app.py

@@ -22,12 +22,14 @@ class MessageModel(BaseModel):
     role: str
     content: Any
 
-class ChatPayload(BaseModel):
+class RunPodInput(BaseModel):
     messages: List[MessageModel]
     max_tokens: Optional[int] = 16384
     temperature: Optional[float] = 1.00
     top_p: Optional[float] = 0.95
-    stream: Optional[bool] = False
+
+class RunPodPayload(BaseModel):
+    input: RunPodInput
 
 def verify_token(api_key: str = Security(api_key_header)):
     if not api_key or api_key.replace("Bearer ", "") != CLIENT_AUTH_TOKEN:
@@ -38,8 +40,9 @@ def verify_token(api_key: str = Security(api_key_header)):
 def home():
     return {"status": "online", "model": "Vedika AI (stepfun-ai/step-3.7-flash)"}
 
+@app.post("/run")
 @app.post("/v1/chat/completions")
-def process_chat(payload: ChatPayload, token: str = Security(verify_token)):
+def process_chat(payload: RunPodPayload, token: str = Security(verify_token)):
     if not NVIDIA_API_KEY:
         raise HTTPException(status_code=500, detail="Server Error: NVIDIA_API_KEY is missing")
 
@@ -48,14 +51,14 @@ def process_chat(payload: ChatPayload, token: str = Security(verify_token)):
         "Accept": "application/json"
     }
     
-    formatted_messages = [VEDIKA_SYSTEM_PROMPT] + [msg.model_dump() for msg in payload.messages]
+    formatted_messages = [VEDIKA_SYSTEM_PROMPT] + [msg.model_dump() for msg in payload.input.messages]
     
     nvidia_payload = {
         "model": MODEL_NAME,
         "messages": formatted_messages,
-        "max_tokens": payload.max_tokens,
-        "temperature": payload.temperature,
-        "top_p": payload.top_p,
+        "max_tokens": payload.input.max_tokens,
+        "temperature": payload.input.temperature,
+        "top_p": payload.input.top_p,
         "stream": False
     }