Update api/utils.py
Browse files- 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 #
|
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":
|
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
|
125 |
-
|
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":
|
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 |
-
|
176 |
async with httpx.AsyncClient() as client:
|
177 |
try:
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
|
|
|
|
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))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|