update: api formate

Author: LJ-Hao

src/flask_server_llm.py

@@ -16,6 +16,29 @@
 RKLLM_Handle_t = ctypes.c_void_p
 userdata = ctypes.c_void_p(None)
 
+LLMCallState = ctypes.c_int
+LLMCallState.RKLLM_RUN_NORMAL  = 0
+LLMCallState.RKLLM_RUN_WAITING  = 1
+LLMCallState.RKLLM_RUN_FINISH  = 2
+LLMCallState.RKLLM_RUN_ERROR   = 3
+
+RKLLMInputType = ctypes.c_int
+RKLLMInputType.RKLLM_INPUT_PROMPT      = 0
+
+RKLLMInferMode = ctypes.c_int
+RKLLMInferMode.RKLLM_INFER_GENERATE = 0
+
+class RKLLMExtendParam(ctypes.Structure):
+    _fields_ = [
+        ("base_domain_id", ctypes.c_int32),
+        ("embed_flash", ctypes.c_int8),
+        ("enabled_cpus_num", ctypes.c_int8),
+        ("enabled_cpus_mask", ctypes.c_uint32),
+        ("n_batch", ctypes.c_uint8),
+        ("use_cross_attn", ctypes.c_int8),
+        ("reserved", ctypes.c_uint8 * 104)
+    ]
+
 class RKLLMParam(ctypes.Structure):
     _fields_ = [
         ("model_path", ctypes.c_char_p),
@@ -36,30 +59,36 @@ class RKLLMParam(ctypes.Structure):
         ("img_start", ctypes.c_char_p),
         ("img_end", ctypes.c_char_p),
         ("img_content", ctypes.c_char_p),
+        ("extend_param", RKLLMExtendParam),
     ]
 
-class RKLLMInput(ctypes.Structure):
+class RKLLMInputUnion(ctypes.Union):
     _fields_ = [
-        ("role", ctypes.c_char_p),
-        ("enable_thinking", ctypes.c_bool),
-        ("input_type", ctypes.c_int),
-        ("input_data", ctypes.c_char_p)
+        ("prompt_input", ctypes.c_char_p),
     ]
 
-class RKLLMResult(ctypes.Structure):
+class RKLLMInput(ctypes.Structure):
     _fields_ = [
-        ("text", ctypes.c_char_p),
-        ("token_id", ctypes.c_int),
+        ("role", ctypes.c_char_p),
+        ("enable_thinking", ctypes.c_bool),
+        ("input_type", RKLLMInputType),
+        ("input_data", RKLLMInputUnion)
     ]
 
 class RKLLMInferParam(ctypes.Structure):
     _fields_ = [
-        ("mode", ctypes.c_int),
+        ("mode", RKLLMInferMode),
         ("lora_params", ctypes.c_void_p),
         ("prompt_cache_params", ctypes.c_void_p),
         ("keep_history", ctypes.c_int)
     ]
 
+class RKLLMResult(ctypes.Structure):
+    _fields_ = [
+        ("text", ctypes.c_char_p),
+        ("token_id", ctypes.c_int),
+    ]
+
 # 锁和状态变量
 lock = threading.Lock()
 is_blocking = False
@@ -71,11 +100,15 @@ class RKLLMInferParam(ctypes.Structure):
 # 回调函数
 def callback_impl(result, userdata, state):
     global global_text, global_state
-    if state == 2:  # FINISH
+    if state == LLMCallState.RKLLM_RUN_FINISH:
         global_state = state
-    elif state == 3:  # ERROR
+        print("\n")
+        sys.stdout.flush()
+    elif state == LLMCallState.RKLLM_RUN_ERROR:
         global_state = state
-    elif state == 0:  # NORMAL
+        print("run error")
+        sys.stdout.flush()
+    elif state == LLMCallState.RKLLM_RUN_NORMAL:
         global_state = state
         if result.contents.text:
             global_text.append(result.contents.text.decode('utf-8'))
@@ -87,6 +120,7 @@ def callback_impl(result, userdata, state):
 # RKLLM 类
 class RKLLM(object):
     def __init__(self, model_path, platform="rk3588"):
+        # 初始化 RKLLMParam
         rkllm_param = RKLLMParam()
         rkllm_param.model_path = bytes(model_path, 'utf-8')
         rkllm_param.max_context_len = 4096
@@ -96,24 +130,50 @@ def __init__(self, model_path, platform="rk3588"):
         rkllm_param.top_p = 0.9
         rkllm_param.temperature = 0.8
         rkllm_param.repeat_penalty = 1.1
+        rkllm_param.frequency_penalty = 0.0
+        rkllm_param.presence_penalty = 0.0
+        rkllm_param.mirostat = 0
+        rkllm_param.mirostat_tau = 5.0
+        rkllm_param.mirostat_eta = 0.1
         rkllm_param.skip_special_token = True
+        rkllm_param.is_async = False
+        rkllm_param.img_start = "".encode('utf-8')
+        rkllm_param.img_end = "".encode('utf-8')
+        rkllm_param.img_content = "".encode('utf-8')
+        
+        # 设置 extend_param
+        rkllm_param.extend_param.base_domain_id = 0
+        rkllm_param.extend_param.embed_flash = 1
+        rkllm_param.extend_param.n_batch = 1  # 关键修复：正确设置 n_batch
+        rkllm_param.extend_param.use_cross_attn = 0
+        rkllm_param.extend_param.enabled_cpus_num = 4
+        
+        # 根据平台设置 CPU 掩码
+        if platform.lower() in ["rk3576", "rk3588"]:
+            rkllm_param.extend_param.enabled_cpus_mask = (1 << 4) | (1 << 5) | (1 << 6) | (1 << 7)
+        else:
+            rkllm_param.extend_param.enabled_cpus_mask = (1 << 0) | (1 << 1) | (1 << 2) | (1 << 3)
 
         self.handle = RKLLM_Handle_t()
 
+        # 初始化函数
         self.rkllm_init = rkllm_lib.rkllm_init
         self.rkllm_init.argtypes = [ctypes.POINTER(RKLLM_Handle_t), ctypes.POINTER(RKLLMParam), callback_type]
         self.rkllm_init.restype = ctypes.c_int
+        
         ret = self.rkllm_init(ctypes.byref(self.handle), ctypes.byref(rkllm_param), callback)
         if ret != 0:
             print("rkllm init failed")
             sys.exit(1)
         else:
             print("rkllm init success!")
 
+        # 运行函数
         self.rkllm_run = rkllm_lib.rkllm_run
         self.rkllm_run.argtypes = [RKLLM_Handle_t, ctypes.POINTER(RKLLMInput), ctypes.POINTER(RKLLMInferParam), ctypes.c_void_p]
         self.rkllm_run.restype = ctypes.c_int
 
+        # 销毁函数
         self.rkllm_destroy = rkllm_lib.rkllm_destroy
         self.rkllm_destroy.argtypes = [RKLLM_Handle_t]
         self.rkllm_destroy.restype = ctypes.c_int
@@ -122,11 +182,11 @@ def run(self, prompt, role="user"):
         rkllm_input = RKLLMInput()
         rkllm_input.role = role.encode('utf-8')
         rkllm_input.enable_thinking = False
-        rkllm_input.input_type = 0  # RKLLM_INPUT_PROMPT
-        rkllm_input.input_data = ctypes.c_char_p(prompt.encode('utf-8'))
+        rkllm_input.input_type = RKLLMInputType.RKLLM_INPUT_PROMPT
+        rkllm_input.input_data.prompt_input = ctypes.c_char_p(prompt.encode('utf-8'))
 
         infer_param = RKLLMInferParam()
-        infer_param.mode = 0  # RKLLM_INFER_GENERATE
+        infer_param.mode = RKLLMInferMode.RKLLM_INFER_GENERATE
         infer_param.lora_params = None
         infer_param.prompt_cache_params = None
         infer_param.keep_history = 0
@@ -154,22 +214,26 @@ def chat_completions():
 
         messages = data['messages']
         stream = data.get('stream', False)
-        n_predict = data.get('n_predict', 512)
 
-        # 构建提示词
+        # 重置全局变量
+        global_text = []
+        global_state = -1
+        
+        # 构建提示词 - 简化的聊天格式
         prompt = ""
         for msg in messages:
             if msg['role'] == 'system':
-                prompt += f"System: {msg['content']}\n\n"
+                prompt += f"{msg['content']}\n\n"
             elif msg['role'] == 'user':
-                prompt += f"User: {msg['content']}\n\n"
+                prompt += f"User: {msg['content']}\n"
             elif msg['role'] == 'assistant':
-                prompt += f"Assistant: {msg['content']}\n\n"
-        prompt += "Assistant: "
+                prompt += f"Assistant: {msg['content']}\n"
 
-        # 重置全局变量
-        global_text = []
-        global_state = -1
+        # 添加最后的 Assistant: 提示
+        if prompt and not prompt.endswith("Assistant: "):
+            prompt += "Assistant: "
+        
+        print(f"Prompt: {prompt}")
 
         def generate_response():
             nonlocal prompt
@@ -200,7 +264,20 @@ def generate_response():
                     model_thread.join(timeout=0.01)
                     model_thread_finished = not model_thread.is_alive()
 
-                    if global_state == 2:  # FINISH
+                    if global_state == LLMCallState.RKLLM_RUN_FINISH:
+                        # 发送结束标记
+                        response_chunk = {
+                            "id": "chatcmpl-123",
+                            "object": "chat.completion.chunk",
+                            "created": int(time.time()),
+                            "model": "rkllm-model",
+                            "choices": [{
+                                "index": 0,
+                                "delta": {},
+                                "finish_reason": "stop"
+                            }]
+                        }
+                        yield f"data: {json.dumps(response_chunk, ensure_ascii=False)}\n\n"
                         break
 
                 # 发送结束标记
@@ -216,7 +293,7 @@ def generate_response():
                     model_thread.join(timeout=0.01)
                     model_thread_finished = not model_thread.is_alive()
 
-                    if global_state == 2:  # FINISH
+                    if global_state == LLMCallState.RKLLM_RUN_FINISH:
                         break
 
                 response = {
@@ -238,12 +315,13 @@ def generate_response():
                         "total_tokens": 0
                     }
                 }
-                yield json.dumps(response, ensure_ascii=False)
+                return json.dumps(response, ensure_ascii=False)
 
         if stream:
             return Response(generate_response(), content_type='text/event-stream')
         else:
-            return Response(generate_response(), content_type='application/json')
+            response_data = generate_response()
+            return Response(response_data, content_type='application/json')
 
     finally:
         lock.release()