⚠️[FIX] Fix /chat/completion with max_completion_tokens and stop parameters

Author: fairyshine

src/fastmindapi/model/llama_cpp/LLM.py

@@ -69,9 +69,10 @@ def generate(self,
 
         return generation_output
 
-    def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, logprobs: bool = False, top_logprobs: int = 10):
+    def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, logprobs: bool = False, top_logprobs: int = 10, stop: list[str] = None):
         response = self.model.create_chat_completion(messages, 
                                                      max_tokens=max_completion_tokens, 
                                                      logprobs=logprobs, 
-                                                     top_logprobs=top_logprobs if logprobs else None)
+                                                     top_logprobs=top_logprobs if logprobs else None,
+                                                     stop=stop)
         return response

src/fastmindapi/model/openai/ChatModel.py

@@ -61,14 +61,15 @@ def generate(self,
                              "logits": logits_list}
         return generation_output
 
-    def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, logprobs: bool = False, top_logprobs: int =10):
+    def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, logprobs: bool = False, top_logprobs: int =10, stop: list[str] = None):
         try:
             completion = self.client.chat.completions.create(
             model= self.model_name,
             messages=messages,
-            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_completion_tokens,
             logprobs=logprobs,
             top_logprobs=top_logprobs if logprobs else None,
+            stop=stop
             )
             return completion.model_dump()
         except Exception as e:

src/fastmindapi/model/transformers/CausalLM.py

@@ -105,7 +105,7 @@ def generate(self,
 
         return generation_output
 
-    def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, logprobs: bool = False, top_logprobs: int = 10):
+    def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, logprobs: bool = False, top_logprobs: int = 10, stop: list[str] = None):
         import torch
         import time
 
@@ -118,9 +118,10 @@ def chat(self, messages: list[ChatMessage], max_completion_tokens: int = None, l
         input_text += "assistant: "
 
         inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
-        
+
         generate_kwargs = {
             "max_new_tokens": max_completion_tokens,
+            "stop_strings": stop
         }
 
         with torch.no_grad():

src/fastmindapi/server/router/openai.py

@@ -7,10 +7,14 @@ class ChatMessage(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
     model: str
     messages: list[ChatMessage]
     max_completion_tokens: int = None
+    logprobs: bool = False
+    top_logprobs: int = 10
+    stop: list[str] = None
 
     model_config=ConfigDict(protected_namespaces=())
 
@@ -24,7 +28,10 @@ def chat_completions(request: Request, item: ChatRequest):
 
     outputs = server.module["model"].loaded_models[item.model].chat(
         messages=item.messages, 
-        max_completion_tokens=item.max_completion_tokens
+        max_completion_tokens=item.max_completion_tokens,
+        logprobs=item.logprobs,
+        top_logprobs=item.top_logprobs,
+        stop=item.stop
     )
     return outputs