Update app.py

app.py

@@ -368,8 +368,8 @@ def count_message_tokens(messages, model="gpt-3.5-turbo-0301"):
 def stream_notdiamond_response(response, model):
     """流式处理 notdiamond API 响应。"""
     buffer = ""
-    previous_content = ""
     content_pattern = re.compile(r'data: (.*?)\n\n', re.DOTALL)
+    full_content = ""
     
     for chunk in response.iter_content(1024):
         if chunk:
@@ -380,15 +380,15 @@ def stream_notdiamond_response(response, model):
                 try:
                     data = json.loads(match)
                     current_content = data['choices'][0]['delta'].get('content', '')
-                    new_content = current_content[len(previous_content):]  # 提取新内容
-                    previous_content = current_content  # 更新已处理的内容
+                    full_content += current_content  # 拼接完整内容
                     
-                    if new_content:
-                        yield create_openai_chunk(new_content, model)
+                    if current_content:
+                        yield create_openai_chunk(current_content, model)
                 except json.JSONDecodeError:
                     logger.error(f"Failed to decode JSON: {match}")
                     continue
     
+    # 生成包含完整内容的最终 chunk
     yield create_openai_chunk('', model, 'stop')
 
 def handle_non_stream_response(response, model, prompt_tokens):
@@ -428,14 +428,13 @@ def handle_non_stream_response(response, model, prompt_tokens):
 def generate_stream_response(response, model, prompt_tokens):
     """生成流式 HTTP 响应。"""
     total_completion_tokens = 0
-    previous_content = ""
+    full_content = ""
     
     for chunk in stream_notdiamond_response(response, model):
         content = chunk['choices'][0]['delta'].get('content', '')
-        new_content = content[len(previous_content):]  # 提取新内容
-        previous_content = content  # 更新已处理的内容
+        full_content += content  # 拼接完整内容
         
-        total_completion_tokens += count_tokens(new_content, model)
+        total_completion_tokens += count_tokens(content, model)
         
         chunk['usage'] = {
             "prompt_tokens": prompt_tokens,
@@ -445,6 +444,15 @@ 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():