Niansuh commited on
Commit
771f426
·
verified ·
1 Parent(s): 33bafcd

Update api/utils.py

Browse files
Files changed (1) hide show
  1. api/utils.py +13 -2
api/utils.py CHANGED
@@ -63,6 +63,12 @@ def message_to_dict(message, model_prefix: Optional[str] = None):
63
  else:
64
  return {"role": message.role, "content": message.content}
65
 
 
 
 
 
 
 
66
  async def process_streaming_response(request: ChatRequest):
67
  agent_mode = AGENT_MODE.get(request.model, {})
68
  trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
@@ -112,7 +118,9 @@ async def process_streaming_response(request: ChatRequest):
112
  content = line
113
  if content.startswith("$@$v=undefined-rv1$@$"):
114
  content = content[21:]
115
- yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
 
 
116
 
117
  yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
118
  yield "data: [DONE]\n\n"
@@ -173,6 +181,9 @@ async def process_non_streaming_response(request: ChatRequest):
173
  if full_response.startswith("$@$v=undefined-rv1$@$"):
174
  full_response = full_response[21:]
175
 
 
 
 
176
  return {
177
  "id": f"chatcmpl-{uuid.uuid4()}",
178
  "object": "chat.completion",
@@ -181,7 +192,7 @@ async def process_non_streaming_response(request: ChatRequest):
181
  "choices": [
182
  {
183
  "index": 0,
184
- "message": {"role": "assistant", "content": full_response},
185
  "finish_reason": "stop",
186
  }
187
  ],
 
63
  else:
64
  return {"role": message.role, "content": message.content}
65
 
66
+ def strip_model_prefix(content: str, model_prefix: str) -> str:
67
+ """Remove the model prefix from the response content if present."""
68
+ if content.startswith(model_prefix):
69
+ return content[len(model_prefix):].strip()
70
+ return content
71
+
72
  async def process_streaming_response(request: ChatRequest):
73
  agent_mode = AGENT_MODE.get(request.model, {})
74
  trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
 
118
  content = line
119
  if content.startswith("$@$v=undefined-rv1$@$"):
120
  content = content[21:]
121
+ # Strip the model prefix from the response content
122
+ cleaned_content = strip_model_prefix(content, model_prefix)
123
+ yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
124
 
125
  yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
126
  yield "data: [DONE]\n\n"
 
181
  if full_response.startswith("$@$v=undefined-rv1$@$"):
182
  full_response = full_response[21:]
183
 
184
+ # Strip the model prefix from the full response
185
+ cleaned_full_response = strip_model_prefix(full_response, model_prefix)
186
+
187
  return {
188
  "id": f"chatcmpl-{uuid.uuid4()}",
189
  "object": "chat.completion",
 
192
  "choices": [
193
  {
194
  "index": 0,
195
+ "message": {"role": "assistant", "content": cleaned_full_response},
196
  "finish_reason": "stop",
197
  }
198
  ],