Niansuh commited on
Commit
0eb0c17
·
verified ·
1 Parent(s): f9f806c

Update api/utils.py

Browse files
Files changed (1) hide show
  1. api/utils.py +9 -7
api/utils.py CHANGED
@@ -141,6 +141,7 @@ async def process_streaming_response(request: ChatRequest):
141
  "imageGenerationMode": False, # Added this line
142
  }
143
 
 
144
  async with httpx.AsyncClient() as client:
145
  try:
146
  async with client.stream(
@@ -152,10 +153,7 @@ async def process_streaming_response(request: ChatRequest):
152
  ) as response:
153
  response.raise_for_status()
154
 
155
- # Start processing the chunks and yield them one by one
156
  timestamp = int(datetime.now().timestamp())
157
- response_content = "" # Collect response content
158
-
159
  async for chunk in response.aiter_text():
160
  if chunk:
161
  content = chunk
@@ -173,14 +171,18 @@ async def process_streaming_response(request: ChatRequest):
173
  # Clean up the content
174
  cleaned_content = strip_model_prefix(content, model_prefix)
175
 
176
- # Yield each chunk as soon as it's ready
 
 
 
177
  yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
178
- response_content += cleaned_content # Collect the full response content
179
 
180
- # After all chunks are processed, add the advertisement text at the end only if it's not already present
181
  if ADVERTISEMENT_TEXT and ADVERTISEMENT_TEXT not in response_content:
182
  response_content += "\n\n" + ADVERTISEMENT_TEXT
183
- yield f"data: {json.dumps(create_chat_completion_data(response_content, request.model, timestamp, 'stop'))}\n\n"
 
 
184
 
185
  # Add the final "done" marker
186
  yield "data: [DONE]\n\n"
 
141
  "imageGenerationMode": False, # Added this line
142
  }
143
 
144
+ response_content = "" # Variable to hold the full response content
145
  async with httpx.AsyncClient() as client:
146
  try:
147
  async with client.stream(
 
153
  ) as response:
154
  response.raise_for_status()
155
 
 
156
  timestamp = int(datetime.now().timestamp())
 
 
157
  async for chunk in response.aiter_text():
158
  if chunk:
159
  content = chunk
 
171
  # Clean up the content
172
  cleaned_content = strip_model_prefix(content, model_prefix)
173
 
174
+ # Add the chunk to the full response content
175
+ response_content += cleaned_content
176
+
177
+ # Yield the cleaned chunk as part of the stream
178
  yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
179
 
180
+ # After all chunks are processed, append the advertisement text at the end if it's not already included
181
  if ADVERTISEMENT_TEXT and ADVERTISEMENT_TEXT not in response_content:
182
  response_content += "\n\n" + ADVERTISEMENT_TEXT
183
+
184
+ # Now yield the final chunk with the advertisement text appended at the end
185
+ yield f"data: {json.dumps(create_chat_completion_data(response_content, request.model, timestamp, 'stop'))}\n\n"
186
 
187
  # Add the final "done" marker
188
  yield "data: [DONE]\n\n"