Niansuh commited on
Commit
7b631fd
·
verified ·
1 Parent(s): dd52fd4

Update api/utils.py

Browse files
Files changed (1) hide show
  1. api/utils.py +14 -101
api/utils.py CHANGED
@@ -1,5 +1,3 @@
1
- # api/utils.py
2
-
3
  from datetime import datetime
4
  import json
5
  from typing import Any, Dict, Optional
@@ -19,69 +17,13 @@ from api.models import ChatRequest
19
 
20
  from api.logger import setup_logger
21
 
22
- import uuid # Added import for uuid
23
 
24
  logger = setup_logger(__name__)
25
 
26
- def create_chat_completion_data(
27
- content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
28
- ) -> Dict[str, Any]:
29
- return {
30
- "id": f"chatcmpl-{uuid.uuid4()}",
31
- "object": "chat.completion.chunk",
32
- "created": timestamp,
33
- "model": model,
34
- "choices": [
35
- {
36
- "index": 0,
37
- "delta": {"content": content, "role": "assistant"},
38
- "finish_reason": finish_reason,
39
- }
40
- ],
41
- "usage": None,
42
- }
43
-
44
- def message_to_dict(message, model_prefix: Optional[str] = None):
45
- if isinstance(message.content, str):
46
- content = message.content
47
- if model_prefix:
48
- content = f"{model_prefix} {content}"
49
- return {"role": message.role, "content": content}
50
- elif isinstance(message.content, list) and len(message.content) == 2:
51
- content = message.content[0]["text"]
52
- if model_prefix:
53
- content = f"{model_prefix} {content}"
54
- return {
55
- "role": message.role,
56
- "content": content,
57
- "data": {
58
- "imageBase64": message.content[1]["image_url"]["url"],
59
- "fileText": "",
60
- "title": "snapshot",
61
- },
62
- }
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 and language specifier from the response content if present."""
68
- # Remove the model prefix
69
- if content.startswith(model_prefix):
70
- content = content[len(model_prefix):].strip()
71
-
72
- # Remove language specifier from code blocks (e.g., ```html -> ```)
73
- if content.startswith("```"):
74
- first_newline = content.find("\n")
75
- if first_newline != -1:
76
- # Retain the triple backticks and the content after the language specifier
77
- content = "```\n" + content[first_newline + 1:]
78
-
79
- return content
80
-
81
  async def process_streaming_response(request: ChatRequest):
82
  agent_mode = AGENT_MODE.get(request.model, {})
83
  trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
84
- model_prefix = MODEL_PREFIXES.get(request.model, "")
85
  referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
86
  referer_url = f"{BASE_URL}{referer_path}"
87
 
@@ -90,7 +32,7 @@ async def process_streaming_response(request: ChatRequest):
90
  dynamic_headers['Referer'] = referer_url
91
 
92
  json_data = {
93
- "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
94
  "previewToken": None,
95
  "userId": None,
96
  "codeModelMode": True,
@@ -121,18 +63,8 @@ async def process_streaming_response(request: ChatRequest):
121
  timeout=100,
122
  ) as response:
123
  response.raise_for_status()
124
- async for line in response.aiter_lines():
125
- timestamp = int(datetime.now().timestamp())
126
- if line:
127
- content = line
128
- if content.startswith("$@$v=undefined-rv1$@$"):
129
- content = content[21:]
130
- # Strip the model prefix and handle code blocks
131
- cleaned_content = strip_model_prefix(content, model_prefix)
132
- yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
133
-
134
- yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
135
- yield "data: [DONE]\n\n"
136
  except httpx.HTTPStatusError as e:
137
  logger.error(f"HTTP error occurred: {e}")
138
  raise HTTPException(status_code=e.response.status_code, detail=str(e))
@@ -143,7 +75,6 @@ async def process_streaming_response(request: ChatRequest):
143
  async def process_non_streaming_response(request: ChatRequest):
144
  agent_mode = AGENT_MODE.get(request.model, {})
145
  trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
146
- model_prefix = MODEL_PREFIXES.get(request.model, "")
147
  referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
148
  referer_url = f"{BASE_URL}{referer_path}"
149
 
@@ -152,7 +83,7 @@ async def process_non_streaming_response(request: ChatRequest):
152
  dynamic_headers['Referer'] = referer_url
153
 
154
  json_data = {
155
- "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
156
  "previewToken": None,
157
  "userId": None,
158
  "codeModelMode": True,
@@ -172,38 +103,20 @@ async def process_non_streaming_response(request: ChatRequest):
172
  "mobileClient": False,
173
  "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
174
  }
175
- full_response = ""
176
  async with httpx.AsyncClient() as client:
177
  try:
178
- async with client.stream(
179
- method="POST", url=f"{BASE_URL}/api/chat", headers=dynamic_headers, json=json_data
180
- ) as response:
181
- response.raise_for_status()
182
- async for chunk in response.aiter_text():
183
- full_response += chunk
 
 
184
  except httpx.HTTPStatusError as e:
185
  logger.error(f"HTTP error occurred: {e}")
186
  raise HTTPException(status_code=e.response.status_code, detail=str(e))
187
  except httpx.RequestError as e:
188
  logger.error(f"Error occurred during request: {e}")
189
  raise HTTPException(status_code=500, detail=str(e))
190
- if full_response.startswith("$@$v=undefined-rv1$@$"):
191
- full_response = full_response[21:]
192
-
193
- # Strip the model prefix and handle code blocks
194
- cleaned_full_response = strip_model_prefix(full_response, model_prefix)
195
-
196
- return {
197
- "id": f"chatcmpl-{uuid.uuid4()}",
198
- "object": "chat.completion",
199
- "created": int(datetime.now().timestamp()),
200
- "model": request.model,
201
- "choices": [
202
- {
203
- "index": 0,
204
- "message": {"role": "assistant", "content": cleaned_full_response},
205
- "finish_reason": "stop",
206
- }
207
- ],
208
- "usage": None,
209
- }
 
 
 
1
  from datetime import datetime
2
  import json
3
  from typing import Any, Dict, Optional
 
17
 
18
  from api.logger import setup_logger
19
 
20
+ import uuid # Ensure this is imported if needed elsewhere
21
 
22
  logger = setup_logger(__name__)
23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  async def process_streaming_response(request: ChatRequest):
25
  agent_mode = AGENT_MODE.get(request.model, {})
26
  trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
 
27
  referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
28
  referer_url = f"{BASE_URL}{referer_path}"
29
 
 
32
  dynamic_headers['Referer'] = referer_url
33
 
34
  json_data = {
35
+ "messages": request.messages, # Directly pass messages without processing
36
  "previewToken": None,
37
  "userId": None,
38
  "codeModelMode": True,
 
63
  timeout=100,
64
  ) as response:
65
  response.raise_for_status()
66
+ async for chunk in response.aiter_text():
67
+ yield chunk # Directly yield the chunk without modification
 
 
 
 
 
 
 
 
 
 
68
  except httpx.HTTPStatusError as e:
69
  logger.error(f"HTTP error occurred: {e}")
70
  raise HTTPException(status_code=e.response.status_code, detail=str(e))
 
75
  async def process_non_streaming_response(request: ChatRequest):
76
  agent_mode = AGENT_MODE.get(request.model, {})
77
  trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
 
78
  referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
79
  referer_url = f"{BASE_URL}{referer_path}"
80
 
 
83
  dynamic_headers['Referer'] = referer_url
84
 
85
  json_data = {
86
+ "messages": request.messages, # Directly pass messages without processing
87
  "previewToken": None,
88
  "userId": None,
89
  "codeModelMode": True,
 
103
  "mobileClient": False,
104
  "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
105
  }
106
+
107
  async with httpx.AsyncClient() as client:
108
  try:
109
+ response = await client.post(
110
+ f"{BASE_URL}/api/chat",
111
+ headers=dynamic_headers,
112
+ json=json_data,
113
+ timeout=100,
114
+ )
115
+ response.raise_for_status()
116
+ return response.json() # Directly return the JSON response
117
  except httpx.HTTPStatusError as e:
118
  logger.error(f"HTTP error occurred: {e}")
119
  raise HTTPException(status_code=e.response.status_code, detail=str(e))
120
  except httpx.RequestError as e:
121
  logger.error(f"Error occurred during request: {e}")
122
  raise HTTPException(status_code=500, detail=str(e))