Update endpoints.py

endpoints.py

@@ -45,19 +45,29 @@ model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 #                )
 # hf_llm = HuggingFacePipeline(pipeline=pipe)
 
-@app.post("/v1/chat/completions")
-async def chat_completions(messages: list, temperature: float = 1.0, max_tokens: int = 50, stream: bool = False):
+
+class ChatRequest(BaseModel):
+    messages: list
+    temperature: float = 1.0
+    max_tokens: int = 50
+    stream: bool = False
+
+class ChatResponse(BaseModel):
+    response: str
+
+@app.post("/v1/chat/completions", response_model=ChatResponse)
+async def chat_completions(request: ChatRequest):
     try:
         # Prepare input prompt
         input_prompt = ""
-        for message in messages:
+        for message in request.messages:
             role = message.get('role', 'user')
             content = message.get('content', '')
             input_prompt += f"{role}: {content}\n"
 
         # Tokenize and generate response
         input_ids = tokenizer.encode(input_prompt, return_tensors='pt')
-        output = model.generate(input_ids, max_length=1024, temperature=temperature, max_tokens=max_tokens)
+        output = model.generate(input_ids, max_length=1024, temperature=request.temperature, max_tokens=request.max_tokens)
 
         # Decode and send response
         response = tokenizer.decode(output[0], skip_special_tokens=True)