Update app.py

app.py

@@ -367,28 +367,15 @@ def count_message_tokens(messages, model="gpt-3.5-turbo-0301"):
 
 def stream_notdiamond_response(response, model):
     """流式处理 notdiamond API 响应。"""
-    buffer = ""
-    content_pattern = re.compile(r'data: (.*?)\n\n', re.DOTALL)
-    full_content = ""
-    
+    last_sent = ""
     for chunk in response.iter_content(1024):
         if chunk:
-            buffer += chunk.decode('utf-8')
-            matches = content_pattern.findall(buffer)
-            
-            for match in matches:
-                try:
-                    data = json.loads(match)
-                    current_content = data['choices'][0]['delta'].get('content', '')
-                    full_content += current_content  # 拼接完整内容
-                    
-                    if current_content:
-                        yield create_openai_chunk(current_content, model)
-                except json.JSONDecodeError:
-                    logger.error(f"Failed to decode JSON: {match}")
-                    continue
+            new_content = chunk.decode('utf-8')
+            to_send = new_content[len(last_sent):]
+            if to_send:
+                yield create_openai_chunk(to_send, model)
+            last_sent = new_content
     
-    # 生成包含完整内容的最终 chunk
     yield create_openai_chunk('', model, 'stop')
 
 def handle_non_stream_response(response, model, prompt_tokens):
@@ -428,12 +415,9 @@ def handle_non_stream_response(response, model, prompt_tokens):
 def generate_stream_response(response, model, prompt_tokens):
     """生成流式 HTTP 响应。"""
     total_completion_tokens = 0
-    full_content = ""
     
     for chunk in stream_notdiamond_response(response, model):
         content = chunk['choices'][0]['delta'].get('content', '')
-        full_content += content  # 拼接完整内容
-        
         total_completion_tokens += count_tokens(content, model)
         
         chunk['usage'] = {
@@ -444,15 +428,6 @@ def generate_stream_response(response, model, prompt_tokens):
         
         yield f"data: {json.dumps(chunk)}\n\n"
     
-    # 生成包含完整内容的最终 chunk
-    final_chunk = create_openai_chunk('', model, 'stop')
-    final_chunk['usage'] = {
-        "prompt_tokens": prompt_tokens,
-        "completion_tokens": total_completion_tokens,
-        "total_tokens": prompt_tokens + total_completion_tokens
-    }
-    yield f"data: {json.dumps(final_chunk)}\n\n"
-    
     yield "data: [DONE]\n\n"
 
 def get_auth_credentials():
@@ -700,4 +675,4 @@ if __name__ == "__main__":
     health_check_thread.start()
     
     port = int(os.environ.get("PORT", 3000))
-    app.run(debug=False, host='0.0.0.0', port=port, threaded=True)
+    app.run(debug=False, host='0.0.0.0', port=port, threaded=True)