Bladeren bron

修复groq流式输出错误

GXX 1 maand geleden
bovenliggende
commit
36f7b1d73f
1 gewijzigde bestanden met toevoegingen van 19 en 3 verwijderingen
  1. 19 3
      chat_model.py

+ 19 - 3
chat_model.py

@@ -80,13 +80,29 @@ def groq_chat(messages, model = "deepseek-r1-distill-llama-70b"):
 
     completion = client.chat.completions.create(
         model=model,
-        stream=False,
+        stream=True,
         messages=messages
     )
+    ans = ""
+    for chunk in completion:
+        # print(chunk)
+        if chunk.choices[0].delta.content:
+            ans += chunk.choices[0].delta.content
 
-    ans = completion.choices[0].message.content
+    start_tag = "<think>"
+    end_tag = "</think>"
 
-    return ans
+    text = ans
+
+    while start_tag in text and end_tag in text:
+        start_idx = text.find(start_tag)
+        end_idx = text.find(end_tag) + len(end_tag)
+        text = text[:start_idx] + text[end_idx:]
+
+    # 清理多余的空白行
+    cleaned_text = "\n".join([line for line in text.split("\n") if line.strip()])
+
+    return cleaned_text.strip()
 
 def siliflow_chat(messages, model = "THUDM/glm-4-9b-chat"):
     """