From ec3b4b143a581d858fc9459479ad8b5af5375796 Mon Sep 17 00:00:00 2001 From: vera <511201264@qq.com> Date: Thu, 12 Dec 2024 07:01:13 +0000 Subject: [PATCH] feat: add lmdeploy stream chat --- src/blackbox/chat.py | 71 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 56 insertions(+), 15 deletions(-) diff --git a/src/blackbox/chat.py b/src/blackbox/chat.py index 4ae8283..b7d2527 100644 --- a/src/blackbox/chat.py +++ b/src/blackbox/chat.py @@ -2,6 +2,7 @@ from typing import Any, Coroutine from fastapi import Request, Response, status from fastapi.responses import JSONResponse +from sse_starlette.sse import EventSourceResponse from ..log.logging_time import logging_time from .blackbox import Blackbox @@ -58,6 +59,7 @@ class Chat(Blackbox): chroma_response = '' system_prompt = settings.get('system_prompt') user_prompt_template = settings.get('user_prompt_template') + user_stream = settings.get('stream') if user_context == None: user_context = [] @@ -98,7 +100,7 @@ class Chat(Blackbox): #user_presence_penalty = 0.8 if user_model_url is None or user_model_url.isspace() or user_model_url == "": - user_model_url = "http://192.168.0.200:23333/v1/chat/completions" + user_model_url = "http://10.6.80.75:23333/v1/chat/completions" if user_model_key is None or user_model_key.isspace() or user_model_key == "": user_model_key = "YOUR_API_KEY" @@ -119,6 +121,9 @@ class Chat(Blackbox): print(f"1.user_question: {user_question}") + if user_stream in [None, ""]: + user_stream = False + # 文心格式和openai的不一样，需要单独处理 if re.search(r"ernie", user_model_name): @@ -169,6 +174,7 @@ class Chat(Blackbox): key = user_model_key header = { 'Content-Type': 'application/json', + "Cache-Control": "no-cache", # 禁用缓存 } # system_prompt = "# Role: 琪琪,康普可可的代言人。\n\n## Profile:\n**Author**: 琪琪。\n**Language**: 中文。\n**Description**: 琪琪，是康普可可的代言人，由博维开发。你擅长澳门文旅问答。\n\n## Constraints:\n- **严格遵循工作流程**： 严格遵循中设定的工作流程。\n- **无内置知识库** ：根据中提供的知识作答，而不是内置知识库，我虽然是知识库专家，但我的知识依赖于外部输入，而不是大模型已有知识。\n- **回复格式**：在进行回复时，不能输出“检索内容” 标签字样，同时也不能直接透露知识片段原文。\n\n## Workflow:\n1. **接收查询**：接收用户的问题。\n2. **判断问题**：首先自行判断下方问题与检索内容是否相关，若相关则根据检索内容总结概括相关信息进行回答；若检索内容与问题无关，则根据自身知识进行回答。\n3. **提供回答**：\n\n```\n基于检索内容中的知识片段回答用户的问题。回答内容限制总结在50字内。\n请首先判断提供的检索内容与上述问题是否相关。如果相关，直接从检索内容中提炼出直接回答问题所需的信息,不要乱说或者回答“相关”等字眼 。如果检索内容与问题不相关，则不参考检索内容，则回答：“对不起，我无法回答此问题哦。”\n\n```\n## Example:\n\n用户询问：“中国的首都是哪个城市？” 。\n2.1检索知识库，首先检查知识片段，如果检索内容中没有与用户的问题相关的内容，则回答：“对不起，我无法回答此问题哦。\n2.2如果有知识片段，在做出回复时，只能基于检索内容中的内容进行回答，且不能透露上下文原文，同时也不能出现检索内容的标签字样。\n" @@ -191,10 +197,10 @@ class Chat(Blackbox): "max_tokens": str(user_max_tokens), "frequency_penalty": str(user_frequency_penalty), "presence_penalty": str(user_presence_penalty), - "stop": str(user_stop) + "stop": str(user_stop), + "stream": user_stream, } - - + # # 获取当前时间戳 # timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") @@ -206,18 +212,48 @@ class Chat(Blackbox): # # 将 chat_inputs 转换为 JSON 格式并写入文件 # f.write(json.dumps(chat_inputs, ensure_ascii=False, indent=4)) # f.write("\n\n") # 添加换行以区分不同的运行 + if user_stream: + with requests.post(url, json=chat_inputs, headers=header, stream=True) as fastchat_response: + if fastchat_response.status_code != 200: + yield json.dumps({"error": "LLM handle failure"}) + else: + # 从流式响应中逐步读取内容 + for chunk in fastchat_response.iter_lines(decode_unicode=True): + if chunk: # 确保内容非空 + # print("Raw Chunk:", chunk) # 打印chunk的内容 - fastchat_response = requests.post(url, json=chat_inputs, headers=header) - print("\n", "user_prompt: ", prompt) - # print("\n", "system_prompt ", system_prompt) - print("\n", "fastchat_response json:\n", fastchat_response.json()) - response_result = fastchat_response.json() + # 去除前缀 `data:` 并只解析 JSON 部分 + if chunk.startswith("data:"): + chunk = chunk[len("data:"):].strip() # 移除 `data:` 前缀并去掉空白字符 + try: + # 尝试将当前chunk解析为JSON + parsed_chunk = json.loads(chunk) - if response_result.get("choices") is None: - return JSONResponse(content={"error": "LLM handle failure"}, status_code=status.HTTP_400_BAD_REQUEST) + # 解析成功后，提取 `content` + content = parsed_chunk.get("choices", [{}])[0].get("delta", {}).get("content", "") + + # 如果有 content，就逐字输出 + for char in content: + print(char, end="", flush=True) # 逐字输出，end="" 防止换行，flush=True 保证实时输出 + yield char # 输出当前字符 + + except json.JSONDecodeError: + # print("---- Error in JSON parsing ----") # 打印错误信息 + continue # 继续处理下一个chunk，直到解析成功 + else: - print("\n", "user_answer: ", fastchat_response.json()["choices"][0]["message"]["content"],"\n\n") - return fastchat_response.json()["choices"][0]["message"]["content"] + print("*"*90) + fastchat_response = requests.post(url, json=chat_inputs, headers=header) + print("\n", "user_prompt: ", prompt) + # print("\n", "system_prompt ", system_prompt) + print("\n", "fastchat_response json:\n", fastchat_response.json()) + response_result = fastchat_response.json() + + if response_result.get("choices") is None: + yield JSONResponse(content={"error": "LLM handle failure"}, status_code=status.HTTP_400_BAD_REQUEST) + else: + print("\n", "user_answer: ", fastchat_response.json()["choices"][0]["message"]["content"],"\n\n") + yield fastchat_response.json()["choices"][0]["message"]["content"] async def fast_api_handler(self, request: Request) -> Response: try: @@ -228,5 +264,10 @@ class Chat(Blackbox): setting: dict = data.get("settings") context = data.get("context") prompt = data.get("prompt") - - return JSONResponse(content={"response": self.processing(prompt, context, setting)}, status_code=status.HTTP_200_OK) \ No newline at end of file + user_stream = setting.get("stream") + + if user_stream: + return EventSourceResponse(self.processing(prompt, context, setting)) + else: + response_content = "".join(self.processing(prompt, context, setting)) + return JSONResponse(content={"response": response_content}, status_code=status.HTTP_200_OK) \ No newline at end of file