Niansuh commited on
Commit
3b4d658
·
verified ·
1 Parent(s): 443b651

Update api/utils.py

Browse files
Files changed (1) hide show
  1. api/utils.py +37 -28
api/utils.py CHANGED
@@ -6,15 +6,31 @@ import uuid
6
  import httpx
7
  from api.config import MODEL_MAPPING, AGENT_MODE, TRENDING_AGENT_MODE, headers
8
  from fastapi import HTTPException
9
- from api.auth import verify_app_secret # Correctly import verify_app_secret
10
-
11
- from api.config import APP_SECRET, BASE_URL
12
  from api.models import ChatRequest
13
-
14
  from api.logger import setup_logger
15
 
16
  logger = setup_logger(__name__)
17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  def message_to_dict(message):
19
  if isinstance(message.content, str):
20
  result = {"role": message.role, "content": message.content}
@@ -34,33 +50,16 @@ def message_to_dict(message):
34
  else:
35
  return {"role": message.role, "content": message.content}
36
 
37
- def create_chat_completion_data(
38
- content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
39
- ) -> Dict[str, Any]:
40
- return {
41
- "id": f"chatcmpl-{uuid.uuid4()}",
42
- "object": "chat.completion.chunk",
43
- "created": timestamp,
44
- "model": model,
45
- "choices": [
46
- {
47
- "index": 0,
48
- "delta": {"content": content, "role": "assistant"},
49
- "finish_reason": finish_reason,
50
- }
51
- ],
52
- "usage": None,
53
- }
54
-
55
-
56
  async def process_streaming_response(request: ChatRequest):
 
 
57
  json_data = {
58
  "messages": [message_to_dict(msg) for msg in request.messages],
59
  "previewToken": None,
60
  "userId": None,
61
  "codeModelMode": True,
62
- "agentMode": AGENT_MODE.get(request.model, {}),
63
- "trendingAgentMode": TRENDING_AGENT_MODE.get(request.model, {}),
64
  "isMicMode": False,
65
  "userSystemPrompt": None,
66
  "maxTokens": request.max_tokens,
@@ -89,7 +88,10 @@ async def process_streaming_response(request: ChatRequest):
89
  async for line in response.aiter_lines():
90
  timestamp = int(datetime.now().timestamp())
91
  if line:
92
- content = line + "\n"
 
 
 
93
  if content.startswith("$@$v=undefined-rv1$@$"):
94
  yield f"data: {json.dumps(create_chat_completion_data(content[21:], request.model, timestamp))}\n\n"
95
  else:
@@ -105,13 +107,15 @@ async def process_streaming_response(request: ChatRequest):
105
  raise HTTPException(status_code=500, detail="Request to external service failed")
106
 
107
  async def process_non_streaming_response(request: ChatRequest):
 
 
108
  json_data = {
109
  "messages": [message_to_dict(msg) for msg in request.messages],
110
  "previewToken": None,
111
  "userId": None,
112
  "codeModelMode": True,
113
- "agentMode": AGENT_MODE.get(request.model, {}),
114
- "trendingAgentMode": TRENDING_AGENT_MODE.get(request.model, {}),
115
  "isMicMode": False,
116
  "userSystemPrompt": None,
117
  "maxTokens": request.max_tokens,
@@ -143,6 +147,11 @@ async def process_non_streaming_response(request: ChatRequest):
143
  raise HTTPException(status_code=500, detail="Request to external service failed")
144
  if full_response.startswith("$@$v=undefined-rv1$@$"):
145
  full_response = full_response[21:]
 
 
 
 
 
146
  return {
147
  "id": f"chatcmpl-{uuid.uuid4()}",
148
  "object": "chat.completion",
 
6
  import httpx
7
  from api.config import MODEL_MAPPING, AGENT_MODE, TRENDING_AGENT_MODE, headers
8
  from fastapi import HTTPException
 
 
 
9
  from api.models import ChatRequest
 
10
  from api.logger import setup_logger
11
 
12
  logger = setup_logger(__name__)
13
 
14
+ def create_chat_completion_data(
15
+ content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
16
+ ) -> Dict[str, Any]:
17
+ # Clean the content to remove extra spaces
18
+ cleaned_content = '\n'.join([line.rstrip() for line in content.splitlines() if line.strip() != ''])
19
+ return {
20
+ "id": f"chatcmpl-{uuid.uuid4()}",
21
+ "object": "chat.completion.chunk",
22
+ "created": timestamp,
23
+ "model": model,
24
+ "choices": [
25
+ {
26
+ "index": 0,
27
+ "delta": {"content": cleaned_content, "role": "assistant"},
28
+ "finish_reason": finish_reason,
29
+ }
30
+ ],
31
+ "usage": None,
32
+ }
33
+
34
  def message_to_dict(message):
35
  if isinstance(message.content, str):
36
  result = {"role": message.role, "content": message.content}
 
50
  else:
51
  return {"role": message.role, "content": message.content}
52
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
53
  async def process_streaming_response(request: ChatRequest):
54
+ agent_mode = AGENT_MODE.get(request.model, {})
55
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
56
  json_data = {
57
  "messages": [message_to_dict(msg) for msg in request.messages],
58
  "previewToken": None,
59
  "userId": None,
60
  "codeModelMode": True,
61
+ "agentMode": agent_mode,
62
+ "trendingAgentMode": trending_agent_mode,
63
  "isMicMode": False,
64
  "userSystemPrompt": None,
65
  "maxTokens": request.max_tokens,
 
88
  async for line in response.aiter_lines():
89
  timestamp = int(datetime.now().timestamp())
90
  if line:
91
+ # Clean the content to remove extra spaces
92
+ content = line.strip()
93
+ content = '\n'.join([l.rstrip() for l in content.splitlines() if l.strip() != ''])
94
+ content += "\n"
95
  if content.startswith("$@$v=undefined-rv1$@$"):
96
  yield f"data: {json.dumps(create_chat_completion_data(content[21:], request.model, timestamp))}\n\n"
97
  else:
 
107
  raise HTTPException(status_code=500, detail="Request to external service failed")
108
 
109
  async def process_non_streaming_response(request: ChatRequest):
110
+ agent_mode = AGENT_MODE.get(request.model, {})
111
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
112
  json_data = {
113
  "messages": [message_to_dict(msg) for msg in request.messages],
114
  "previewToken": None,
115
  "userId": None,
116
  "codeModelMode": True,
117
+ "agentMode": agent_mode,
118
+ "trendingAgentMode": trending_agent_mode,
119
  "isMicMode": False,
120
  "userSystemPrompt": None,
121
  "maxTokens": request.max_tokens,
 
147
  raise HTTPException(status_code=500, detail="Request to external service failed")
148
  if full_response.startswith("$@$v=undefined-rv1$@$"):
149
  full_response = full_response[21:]
150
+
151
+ # Clean the content to remove extra spaces
152
+ full_response = full_response.strip()
153
+ full_response = '\n'.join([line.rstrip() for line in full_response.splitlines() if line.strip() != ''])
154
+
155
  return {
156
  "id": f"chatcmpl-{uuid.uuid4()}",
157
  "object": "chat.completion",