Niansuh commited on
Commit
0e9cdf4
·
verified ·
1 Parent(s): 370673c

Upload 20 files

Browse files
Files changed (7) hide show
  1. Dockerfile +0 -3
  2. api/config.py +115 -172
  3. api/logger.py +20 -20
  4. api/models.py +14 -14
  5. api/routes.py +60 -60
  6. api/utils.py +480 -241
  7. requirements.txt +6 -7
Dockerfile CHANGED
@@ -9,9 +9,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
9
  # Set the working directory in the container
10
  WORKDIR /app
11
 
12
- # Create the error.log file and set permissions
13
- RUN touch /app/error.log && chmod 666 /app/error.log
14
-
15
  # Copy the requirements file first
16
  COPY requirements.txt /app/
17
 
 
9
  # Set the working directory in the container
10
  WORKDIR /app
11
 
 
 
 
12
  # Copy the requirements file first
13
  COPY requirements.txt /app/
14
 
api/config.py CHANGED
@@ -1,172 +1,115 @@
1
- import os
2
- from dotenv import load_dotenv
3
-
4
- load_dotenv()
5
-
6
- BASE_URL = "https://www.blackbox.ai"
7
- headers = {
8
- 'accept': '*/*',
9
- 'accept-language': 'en-US,en;q=0.9',
10
- 'origin': 'https://www.blackbox.ai',
11
- 'priority': 'u=1, i',
12
- 'sec-ch-ua': '"Google Chrome";v="129", "Not=A?Brand";v="8", "Chromium";v="129"',
13
- 'sec-ch-ua-mobile': '?0',
14
- 'sec-ch-ua-platform': '"Windows"',
15
- 'sec-fetch-dest': 'empty',
16
- 'sec-fetch-mode': 'cors',
17
- 'sec-fetch-site': 'same-origin',
18
- 'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
19
- 'AppleWebKit/537.36 (KHTML, like Gecko) '
20
- 'Chrome/129.0.0.0 Safari/537.36',
21
- }
22
- APP_SECRET = os.getenv("APP_SECRET")
23
-
24
- ALLOWED_MODELS = [
25
- {"id": "blackboxai", "name": "blackboxai"},
26
- {"id": "blackboxai-pro", "name": "blackboxai-pro"},
27
- {"id": "flux", "name": "flux"},
28
- {"id": "llama-3.1-8b", "name": "llama-3.1-8b"},
29
- {"id": "llama-3.1-70b", "name": "llama-3.1-70b"},
30
- {"id": "llama-3.1-405b", "name": "llama-3.1-405b"},
31
- {"id": "gpt-4o", "name": "gpt-4o"},
32
- {"id": "gemini-pro", "name": "gemini-pro"},
33
- {"id": "gemini-1.5-flash", "name": "gemini-1.5-flash"},
34
- {"id": "claude-sonnet-3.5", "name": "claude-sonnet-3.5"},
35
- {"id": "PythonAgent", "name": "PythonAgent"},
36
- {"id": "JavaAgent", "name": "JavaAgent"},
37
- {"id": "JavaScriptAgent", "name": "JavaScriptAgent"},
38
- {"id": "HTMLAgent", "name": "HTMLAgent"},
39
- {"id": "GoogleCloudAgent", "name": "GoogleCloudAgent"},
40
- {"id": "AndroidDeveloper", "name": "AndroidDeveloper"},
41
- {"id": "SwiftDeveloper", "name": "SwiftDeveloper"},
42
- {"id": "Next.jsAgent", "name": "Next.jsAgent"},
43
- {"id": "MongoDBAgent", "name": "MongoDBAgent"},
44
- {"id": "PyTorchAgent", "name": "PyTorchAgent"},
45
- {"id": "ReactAgent", "name": "ReactAgent"},
46
- {"id": "XcodeAgent", "name": "XcodeAgent"},
47
- {"id": "AngularJSAgent", "name": "AngularJSAgent"},
48
- {"id": "RepoMap", "name": "RepoMap"},
49
- {"id": "gemini-1.5-pro-latest", "name": "gemini-pro"},
50
- {"id": "gemini-1.5-pro", "name": "gemini-1.5-pro"},
51
- {"id": "claude-3-5-sonnet-20240620", "name": "claude-sonnet-3.5"},
52
- {"id": "claude-3-5-sonnet", "name": "claude-sonnet-3.5"},
53
- ]
54
-
55
- MODEL_MAPPING = {
56
- "blackboxai": "blackboxai",
57
- "blackboxai-pro": "blackboxai-pro",
58
- "flux": "flux",
59
- "ImageGeneration": "flux",
60
- "llama-3.1-8b": "llama-3.1-8b",
61
- "llama-3.1-70b": "llama-3.1-70b",
62
- "llama-3.1-405b": "llama-3.1-405b",
63
- "gpt-4o": "gpt-4o",
64
- "gemini-pro": "gemini-pro",
65
- "gemini-1.5-flash": "gemini-1.5-flash",
66
- "claude-sonnet-3.5": "claude-sonnet-3.5",
67
- "PythonAgent": "PythonAgent",
68
- "JavaAgent": "JavaAgent",
69
- "JavaScriptAgent": "JavaScriptAgent",
70
- "HTMLAgent": "HTMLAgent",
71
- "GoogleCloudAgent": "GoogleCloudAgent",
72
- "AndroidDeveloper": "AndroidDeveloper",
73
- "SwiftDeveloper": "SwiftDeveloper",
74
- "Next.jsAgent": "Next.jsAgent",
75
- "MongoDBAgent": "MongoDBAgent",
76
- "PyTorchAgent": "PyTorchAgent",
77
- "ReactAgent": "ReactAgent",
78
- "XcodeAgent": "XcodeAgent",
79
- "AngularJSAgent": "AngularJSAgent",
80
- "RepoMap": "RepoMap",
81
- # Additional mappings
82
- "gemini-flash": "gemini-1.5-flash",
83
- "claude-3.5-sonnet": "claude-sonnet-3.5",
84
- "flux": "flux",
85
- "gemini-1.5-pro-latest": "gemini-pro",
86
- "gemini-1.5-pro": "gemini-1.5-pro",
87
- "claude-3-5-sonnet-20240620": "claude-sonnet-3.5",
88
- "claude-3-5-sonnet": "claude-sonnet-3.5",
89
- }
90
-
91
- USER_SELECTED_MODEL = {
92
- "gpt-4o": "gpt-4o",
93
- "gemini-pro": "gemini-pro",
94
- 'claude-sonnet-3.5': "claude-sonnet-3.5",
95
- "PythonAgent": "PythonAgent",
96
- "JavaAgent": "JavaAgent",
97
- "JavaScriptAgent": "JavaScriptAgent",
98
- "HTMLAgent": "HTMLAgent",
99
- "GoogleCloudAgent": "GoogleCloudAgent",
100
- "AndroidDeveloper": "AndroidDeveloper",
101
- "SwiftDeveloper": "SwiftDeveloper",
102
- "Next.jsAgent": "Next.jsAgent",
103
- "MongoDBAgent": "MongoDBAgent",
104
- "PyTorchAgent": "PyTorchAgent",
105
- "ReactAgent": "ReactAgent",
106
- "XcodeAgent": "XcodeAgent",
107
- "AngularJSAgent": "AngularJSAgent",
108
- "RepoMap": "RepoMap",
109
- }
110
-
111
- MODEL_PREFIXES = {
112
- 'gpt-4o': '@GPT-4o',
113
- 'gemini-pro': '@Gemini-PRO',
114
- 'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
115
- 'PythonAgent': '@Python Agent',
116
- 'JavaAgent': '@Java Agent',
117
- 'JavaScriptAgent': '@JavaScript Agent',
118
- 'HTMLAgent': '@HTML Agent',
119
- 'GoogleCloudAgent': '@Google Cloud Agent',
120
- 'AndroidDeveloper': '@Android Developer',
121
- 'SwiftDeveloper': '@Swift Developer',
122
- 'Next.jsAgent': '@Next.js Agent',
123
- 'MongoDBAgent': '@MongoDB Agent',
124
- 'PyTorchAgent': '@PyTorch Agent',
125
- 'ReactAgent': '@React Agent',
126
- 'XcodeAgent': '@Xcode Agent',
127
- 'AngularJSAgent': '@AngularJS Agent',
128
- 'blackboxai-pro': '@BLACKBOXAI-PRO',
129
- 'ImageGeneration': '@Image Generation',
130
- }
131
-
132
- MODEL_REFERERS = {
133
- "blackboxai": "/?model=blackboxai",
134
- "gpt-4o": "/?model=gpt-4o",
135
- "gemini-pro": "/?model=gemini-pro",
136
- "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
137
- # Add more referers if needed
138
- }
139
-
140
- MODEL_ALIASES = {
141
- "gemini-flash": "gemini-1.5-flash",
142
- "claude-3.5-sonnet": "claude-sonnet-3.5",
143
- "flux": "ImageGeneration",
144
- }
145
-
146
- # Agent modes
147
- AGENT_MODE = {
148
- 'flux': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "flux"},
149
- }
150
-
151
- TRENDING_AGENT_MODE = {
152
- "blackboxai": {},
153
- "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
154
- "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
155
- 'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
156
- 'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
157
- 'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
158
- 'PythonAgent': {'mode': True, 'id': "Python Agent"},
159
- 'JavaAgent': {'mode': True, 'id': "Java Agent"},
160
- 'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
161
- 'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
162
- 'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
163
- 'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
164
- 'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
165
- 'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
166
- 'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
167
- 'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
168
- 'ReactAgent': {'mode': True, 'id': "React Agent"},
169
- 'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
170
- 'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
171
- 'RepoMap': {'mode': True, 'id': "repomap"},
172
- }
 
1
+ import os
2
+ from dotenv import load_dotenv
3
+
4
+ load_dotenv()
5
+
6
+ BASE_URL = "https://www.blackbox.ai"
7
+ headers = {
8
+ 'accept': '*/*',
9
+ 'accept-language': 'en-US,en;q=0.9',
10
+ 'origin': 'https://www.blackbox.ai',
11
+ 'priority': 'u=1, i',
12
+ 'sec-ch-ua': '"Google Chrome";v="129", "Not=A?Brand";v="8", "Chromium";v="129"',
13
+ 'sec-ch-ua-mobile': '?0',
14
+ 'sec-ch-ua-platform': '"Windows"',
15
+ 'sec-fetch-dest': 'empty',
16
+ 'sec-fetch-mode': 'cors',
17
+ 'sec-fetch-site': 'same-origin',
18
+ 'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
19
+ 'AppleWebKit/537.36 (KHTML, like Gecko) '
20
+ 'Chrome/129.0.0.0 Safari/537.36',
21
+ }
22
+ APP_SECRET = os.getenv("APP_SECRET")
23
+
24
+ ALLOWED_MODELS = [
25
+ {"id": "blackboxai", "name": "blackboxai"},
26
+ {"id": "blackboxai-pro", "name": "blackboxai-pro"},
27
+ {"id": "flux", "name": "flux"},
28
+ {"id": "llama-3.1-8b", "name": "llama-3.1-8b"},
29
+ {"id": "llama-3.1-70b", "name": "llama-3.1-70b"},
30
+ {"id": "llama-3.1-405b", "name": "llama-3.1-405b"},
31
+ {"id": "gpt-4o", "name": "gpt-4o"},
32
+ {"id": "gemini-pro", "name": "gemini-pro"},
33
+ {"id": "gemini-1.5-flash", "name": "gemini-1.5-flash"},
34
+ {"id": "claude-sonnet-3.5", "name": "claude-sonnet-3.5"},
35
+ {"id": "PythonAgent", "name": "PythonAgent"},
36
+ {"id": "JavaAgent", "name": "JavaAgent"},
37
+ {"id": "JavaScriptAgent", "name": "JavaScriptAgent"},
38
+ {"id": "HTMLAgent", "name": "HTMLAgent"},
39
+ {"id": "GoogleCloudAgent", "name": "GoogleCloudAgent"},
40
+ {"id": "AndroidDeveloper", "name": "AndroidDeveloper"},
41
+ {"id": "SwiftDeveloper", "name": "SwiftDeveloper"},
42
+ {"id": "Next.jsAgent", "name": "Next.jsAgent"},
43
+ {"id": "MongoDBAgent", "name": "MongoDBAgent"},
44
+ {"id": "PyTorchAgent", "name": "PyTorchAgent"},
45
+ {"id": "ReactAgent", "name": "ReactAgent"},
46
+ {"id": "XcodeAgent", "name": "XcodeAgent"},
47
+ {"id": "AngularJSAgent", "name": "AngularJSAgent"},
48
+ {"id": "RepoMap", "name": "RepoMap"},
49
+ {"id": "gemini-1.5-pro-latest", "name": "gemini-pro"},
50
+ {"id": "gemini-1.5-pro", "name": "gemini-1.5-pro"},
51
+ {"id": "claude-3-5-sonnet-20240620", "name": "claude-sonnet-3.5"},
52
+ {"id": "claude-3-5-sonnet", "name": "claude-sonnet-3.5"},
53
+ ]
54
+
55
+ MODEL_MAPPING = {
56
+ "blackboxai": "blackboxai",
57
+ "blackboxai-pro": "blackboxai-pro",
58
+ "ImageGeneration": "flux",
59
+ "llama-3.1-8b": "llama-3.1-8b",
60
+ "llama-3.1-70b": "llama-3.1-70b",
61
+ "llama-3.1-405b": "llama-3.1-405b",
62
+ "gpt-4o": "gpt-4o",
63
+ "gemini-pro": "gemini-pro",
64
+ "gemini-1.5-flash": "gemini-1.5-flash",
65
+ "claude-sonnet-3.5": "claude-sonnet-3.5",
66
+ "PythonAgent": "PythonAgent",
67
+ "JavaAgent": "JavaAgent",
68
+ "JavaScriptAgent": "JavaScriptAgent",
69
+ "HTMLAgent": "HTMLAgent",
70
+ "GoogleCloudAgent": "GoogleCloudAgent",
71
+ "AndroidDeveloper": "AndroidDeveloper",
72
+ "SwiftDeveloper": "SwiftDeveloper",
73
+ "Next.jsAgent": "Next.jsAgent",
74
+ "MongoDBAgent": "MongoDBAgent",
75
+ "PyTorchAgent": "PyTorchAgent",
76
+ "ReactAgent": "ReactAgent",
77
+ "XcodeAgent": "XcodeAgent",
78
+ "AngularJSAgent": "AngularJSAgent",
79
+ "RepoMap": "RepoMap",
80
+ # Additional mappings
81
+ "gemini-flash": "gemini-1.5-flash",
82
+ "claude-3.5-sonnet": "claude-sonnet-3.5",
83
+ "gemini-1.5-pro-latest": "gemini-pro",
84
+ "gemini-1.5-pro": "gemini-1.5-pro",
85
+ "claude-3-5-sonnet-20240620": "claude-sonnet-3.5",
86
+ "claude-3-5-sonnet": "claude-sonnet-3.5",
87
+ }
88
+
89
+ # Agent modes
90
+ AGENT_MODE = {
91
+ 'flux': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "flux"},
92
+ }
93
+
94
+ TRENDING_AGENT_MODE = {
95
+ "blackboxai": {},
96
+ "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
97
+ "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
98
+ 'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
99
+ 'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
100
+ 'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
101
+ 'PythonAgent': {'mode': True, 'id': "Python Agent"},
102
+ 'JavaAgent': {'mode': True, 'id': "Java Agent"},
103
+ 'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
104
+ 'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
105
+ 'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
106
+ 'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
107
+ 'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
108
+ 'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
109
+ 'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
110
+ 'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
111
+ 'ReactAgent': {'mode': True, 'id': "React Agent"},
112
+ 'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
113
+ 'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
114
+ 'RepoMap': {'mode': True, 'id': "repomap"},
115
+ }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
api/logger.py CHANGED
@@ -1,20 +1,20 @@
1
- import logging
2
-
3
- def setup_logger(name):
4
- logger = logging.getLogger(name)
5
- if not logger.handlers:
6
- logger.setLevel(logging.INFO)
7
- formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
8
-
9
- # Console handler
10
- console_handler = logging.StreamHandler()
11
- console_handler.setFormatter(formatter)
12
- logger.addHandler(console_handler)
13
-
14
- # File Handler - Error Level
15
- error_file_handler = logging.FileHandler('error.log')
16
- error_file_handler.setFormatter(formatter)
17
- error_file_handler.setLevel(logging.ERROR)
18
- logger.addHandler(error_file_handler)
19
-
20
- return logger
 
1
+ import logging
2
+
3
+ def setup_logger(name):
4
+ logger = logging.getLogger(name)
5
+ if not logger.handlers:
6
+ logger.setLevel(logging.INFO)
7
+ formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
8
+
9
+ # Console handler
10
+ console_handler = logging.StreamHandler()
11
+ console_handler.setFormatter(formatter)
12
+ logger.addHandler(console_handler)
13
+
14
+ # File Handler - Error Level
15
+ # error_file_handler = logging.FileHandler('error.log')
16
+ # error_file_handler.setFormatter(formatter)
17
+ # error_file_handler.setLevel(logging.ERROR)
18
+ # logger.addHandler(error_file_handler)
19
+
20
+ return logger
api/models.py CHANGED
@@ -1,14 +1,14 @@
1
- from typing import List, Optional, Union, Dict, Any
2
- from pydantic import BaseModel
3
-
4
- class Message(BaseModel):
5
- role: str
6
- content: Union[str, List[Dict[str, Any]]]
7
-
8
- class ChatRequest(BaseModel):
9
- model: str
10
- messages: List[Message]
11
- stream: Optional[bool] = False
12
- temperature: Optional[float] = 0.7
13
- top_p: Optional[float] = 0.9
14
- max_tokens: Optional[int] = 8192
 
1
+ from typing import List, Optional
2
+ from pydantic import BaseModel
3
+
4
+ class Message(BaseModel):
5
+ role: str
6
+ content: str | list
7
+
8
+ class ChatRequest(BaseModel):
9
+ model: str
10
+ messages: List[Message]
11
+ stream: Optional[bool] = False
12
+ temperature: Optional[float] = 0.7
13
+ top_p: Optional[float] = 0.9
14
+ max_tokens: Optional[int] = 8192
api/routes.py CHANGED
@@ -1,60 +1,60 @@
1
- import json
2
- from fastapi import APIRouter, Depends, HTTPException, Request, Response
3
- from fastapi.responses import StreamingResponse
4
- from api.auth import verify_app_secret
5
- from api.config import ALLOWED_MODELS
6
- from api.models import ChatRequest
7
- from api.utils import process_non_streaming_response, process_streaming_response
8
- from api.logger import setup_logger
9
-
10
- logger = setup_logger(__name__)
11
-
12
- router = APIRouter()
13
-
14
- @router.options("/v1/chat/completions")
15
- @router.options("/api/v1/chat/completions")
16
- async def chat_completions_options():
17
- return Response(
18
- status_code=200,
19
- headers={
20
- "Access-Control-Allow-Origin": "*",
21
- "Access-Control-Allow-Methods": "POST, OPTIONS",
22
- "Access-Control-Allow-Headers": "Content-Type, Authorization",
23
- },
24
- )
25
-
26
- @router.get("/v1/models")
27
- @router.get("/api/v1/models")
28
- async def list_models():
29
- return {"object": "list", "data": ALLOWED_MODELS}
30
-
31
- @router.post("/v1/chat/completions")
32
- @router.post("/api/v1/chat/completions")
33
- async def chat_completions(
34
- request: ChatRequest, app_secret: str = Depends(verify_app_secret)
35
- ):
36
- logger.info("Entering chat_completions route")
37
- logger.info(f"Processing chat completion request for model: {request.model}")
38
-
39
- if request.model not in [model["id"] for model in ALLOWED_MODELS]:
40
- raise HTTPException(
41
- status_code=400,
42
- detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
43
- )
44
-
45
- if request.stream:
46
- logger.info("Streaming response")
47
- return StreamingResponse(process_streaming_response(request), media_type="text/event-stream")
48
- else:
49
- logger.info("Non-streaming response")
50
- return await process_non_streaming_response(request)
51
-
52
-
53
- @router.route('/')
54
- @router.route('/healthz')
55
- @router.route('/ready')
56
- @router.route('/alive')
57
- @router.route('/status')
58
- @router.get("/health")
59
- def health_check(request: Request):
60
- return Response(content=json.dumps({"status": "ok"}), media_type="application/json")
 
1
+ import json
2
+ from fastapi import APIRouter, Depends, HTTPException, Request, Response
3
+ from fastapi.responses import StreamingResponse
4
+ from api.auth import verify_app_secret
5
+ from api.config import ALLOWED_MODELS
6
+ from api.models import ChatRequest
7
+ from api.utils import process_non_streaming_response, process_streaming_response
8
+ from api.logger import setup_logger
9
+
10
+ logger = setup_logger(__name__)
11
+
12
+ router = APIRouter()
13
+
14
+ @router.options("/v1/chat/completions")
15
+ @router.options("/api/v1/chat/completions")
16
+ async def chat_completions_options():
17
+ return Response(
18
+ status_code=200,
19
+ headers={
20
+ "Access-Control-Allow-Origin": "*",
21
+ "Access-Control-Allow-Methods": "POST, OPTIONS",
22
+ "Access-Control-Allow-Headers": "Content-Type, Authorization",
23
+ },
24
+ )
25
+
26
+ @router.get("/v1/models")
27
+ @router.get("/api/v1/models")
28
+ async def list_models():
29
+ return {"object": "list", "data": ALLOWED_MODELS}
30
+
31
+ @router.post("/v1/chat/completions")
32
+ @router.post("/api/v1/chat/completions")
33
+ async def chat_completions(
34
+ request: ChatRequest, app_secret: str = Depends(verify_app_secret)
35
+ ):
36
+ logger.info("Entering chat_completions route")
37
+ logger.info(f"Processing chat completion request for model: {request.model}")
38
+
39
+ if request.model not in [model["id"] for model in ALLOWED_MODELS]:
40
+ raise HTTPException(
41
+ status_code=400,
42
+ detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
43
+ )
44
+
45
+ if request.stream:
46
+ logger.info("Streaming response")
47
+ return StreamingResponse(process_streaming_response(request), media_type="text/event-stream")
48
+ else:
49
+ logger.info("Non-streaming response")
50
+ return await process_non_streaming_response(request)
51
+
52
+
53
+ @router.route('/')
54
+ @router.route('/healthz')
55
+ @router.route('/ready')
56
+ @router.route('/alive')
57
+ @router.route('/status')
58
+ @router.get("/health")
59
+ def health_check(request: Request):
60
+ return Response(content=json.dumps({"status": "ok"}), media_type="application/json")
api/utils.py CHANGED
@@ -1,241 +1,480 @@
1
- from datetime import datetime
2
- import json
3
- from typing import Any, Dict, Optional
4
- import uuid
5
- import re
6
-
7
- import httpx
8
- from api.config import (
9
- MODEL_MAPPING,
10
- USER_SELECTED_MODEL,
11
- MODEL_PREFIXES,
12
- MODEL_REFERERS,
13
- MODEL_ALIASES,
14
- headers,
15
- AGENT_MODE,
16
- TRENDING_AGENT_MODE,
17
- BASE_URL,
18
- )
19
- from fastapi import HTTPException
20
- from api.models import ChatRequest
21
- from api.logger import setup_logger
22
-
23
- logger = setup_logger(__name__)
24
-
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
-
45
- def message_to_dict(message):
46
- if isinstance(message.content, str):
47
- return {"role": message.role, "content": message.content}
48
- elif isinstance(message.content, list) and len(message.content) == 2:
49
- return {
50
- "role": message.role,
51
- "content": message.content[0]["text"],
52
- "data": {
53
- "imageBase64": message.content[1]["image_url"]["url"],
54
- "fileText": "",
55
- "title": "snapshot",
56
- },
57
- }
58
- else:
59
- return {"role": message.role, "content": message.content}
60
-
61
-
62
- def get_full_model_name(model: str) -> str:
63
- # Handle aliases
64
- return MODEL_ALIASES.get(model, model)
65
-
66
-
67
- def get_model_prefix(model: str) -> str:
68
- return MODEL_PREFIXES.get(model, "")
69
-
70
-
71
- def get_referer_url(model: str) -> str:
72
- referer_path = MODEL_REFERERS.get(model, f"/?model={model}")
73
- return f"{BASE_URL}{referer_path}"
74
-
75
-
76
- async def process_streaming_response(request: ChatRequest):
77
- model = get_full_model_name(request.model)
78
- agent_mode = AGENT_MODE.get(model, {})
79
- trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
80
-
81
- prefix = get_model_prefix(model)
82
-
83
- # Construct formatted prompt
84
- formatted_prompt = ""
85
- for msg in request.messages:
86
- role = msg.role.capitalize()
87
- content = msg.content
88
- if isinstance(content, list) and len(content) == 2:
89
- # Handle image content
90
- content = f"FILE:BB\n$#$\n\n$#$\n{msg.content[0]['text']}"
91
- if role and content:
92
- formatted_prompt += f"{role}: {content}\n"
93
-
94
- if prefix:
95
- formatted_prompt = f"{prefix} {formatted_prompt}".strip()
96
-
97
- json_data = {
98
- "messages": [
99
- {
100
- "role": msg.role,
101
- "content": msg.content[0]["text"] if isinstance(msg.content, list) else msg.content,
102
- "data": msg.content[1]["image_url"]["url"] if isinstance(msg.content, list) and len(msg.content) == 2 else None,
103
- }
104
- for msg in request.messages
105
- ],
106
- "previewToken": None,
107
- "userId": None,
108
- "codeModelMode": True,
109
- "agentMode": agent_mode,
110
- "trendingAgentMode": trending_agent_mode,
111
- "isMicMode": False,
112
- "userSystemPrompt": None,
113
- "maxTokens": request.max_tokens,
114
- "playgroundTopP": request.top_p,
115
- "playgroundTemperature": request.temperature,
116
- "isChromeExt": False,
117
- "githubToken": None,
118
- "clickedAnswer2": False,
119
- "clickedAnswer3": False,
120
- "clickedForceWebSearch": False,
121
- "visitFromDelta": False,
122
- "mobileClient": False,
123
- "userSelectedModel": USER_SELECTED_MODEL.get(model, model),
124
- }
125
-
126
- async with httpx.AsyncClient() as client:
127
- try:
128
- async with client.stream(
129
- "POST",
130
- f"{BASE_URL}/api/chat",
131
- headers=headers,
132
- json=json_data,
133
- timeout=100,
134
- ) as response:
135
- response.raise_for_status()
136
- async for line in response.aiter_lines():
137
- timestamp = int(datetime.now().timestamp())
138
- if line:
139
- content = line
140
- # Clean the response if necessary
141
- if content.startswith("$@$v=undefined-rv1$@$"):
142
- content = content[21:]
143
- yield f"data: {json.dumps(create_chat_completion_data(content, model, timestamp))}\n\n"
144
-
145
- # Indicate the end of the stream
146
- timestamp = int(datetime.now().timestamp())
147
- yield f"data: {json.dumps(create_chat_completion_data('', model, timestamp, 'stop'))}\n\n"
148
- yield "data: [DONE]\n\n"
149
- except httpx.HTTPStatusError as e:
150
- logger.error(f"HTTP error occurred: {e}")
151
- raise HTTPException(status_code=e.response.status_code, detail=str(e))
152
- except httpx.RequestError as e:
153
- logger.error(f"Error occurred during request: {e}")
154
- raise HTTPException(status_code=500, detail=str(e))
155
-
156
-
157
- async def process_non_streaming_response(request: ChatRequest):
158
- model = get_full_model_name(request.model)
159
- agent_mode = AGENT_MODE.get(model, {})
160
- trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
161
-
162
- prefix = get_model_prefix(model)
163
-
164
- # Construct formatted prompt
165
- formatted_prompt = ""
166
- for msg in request.messages:
167
- role = msg.role.capitalize()
168
- content = msg.content
169
- if isinstance(content, list) and len(content) == 2:
170
- # Handle image content
171
- content = f"FILE:BB\n$#$\n\n$#$\n{msg.content[0]['text']}"
172
- if role and content:
173
- formatted_prompt += f"{role}: {content}\n"
174
-
175
- if prefix:
176
- formatted_prompt = f"{prefix} {formatted_prompt}".strip()
177
-
178
- json_data = {
179
- "messages": [
180
- {
181
- "role": msg.role,
182
- "content": msg.content[0]["text"] if isinstance(msg.content, list) else msg.content,
183
- "data": msg.content[1]["image_url"]["url"] if isinstance(msg.content, list) and len(msg.content) == 2 else None,
184
- }
185
- for msg in request.messages
186
- ],
187
- "previewToken": None,
188
- "userId": None,
189
- "codeModelMode": True,
190
- "agentMode": agent_mode,
191
- "trendingAgentMode": trending_agent_mode,
192
- "isMicMode": False,
193
- "userSystemPrompt": None,
194
- "maxTokens": request.max_tokens,
195
- "playgroundTopP": request.top_p,
196
- "playgroundTemperature": request.temperature,
197
- "isChromeExt": False,
198
- "githubToken": None,
199
- "clickedAnswer2": False,
200
- "clickedAnswer3": False,
201
- "clickedForceWebSearch": False,
202
- "visitFromDelta": False,
203
- "mobileClient": False,
204
- "userSelectedModel": USER_SELECTED_MODEL.get(model, model),
205
- }
206
-
207
- async with httpx.AsyncClient() as client:
208
- try:
209
- response = await client.post(
210
- f"{BASE_URL}/api/chat",
211
- headers=headers,
212
- json=json_data,
213
- timeout=100,
214
- )
215
- response.raise_for_status()
216
- full_response = response.text
217
-
218
- # Clean the response if necessary
219
- if full_response.startswith("$@$v=undefined-rv1$@$"):
220
- full_response = full_response[21:]
221
-
222
- return {
223
- "id": f"chatcmpl-{uuid.uuid4()}",
224
- "object": "chat.completion",
225
- "created": int(datetime.now().timestamp()),
226
- "model": model,
227
- "choices": [
228
- {
229
- "index": 0,
230
- "message": {"role": "assistant", "content": full_response},
231
- "finish_reason": "stop",
232
- }
233
- ],
234
- "usage": None,
235
- }
236
- except httpx.HTTPStatusError as e:
237
- logger.error(f"HTTP error occurred: {e}")
238
- raise HTTPException(status_code=e.response.status_code, detail=str(e))
239
- except httpx.RequestError as e:
240
- logger.error(f"Error occurred during request: {e}")
241
- raise HTTPException(status_code=500, detail=str(e))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datetime import datetime
2
+ import json
3
+ from typing import Any, Dict, Optional
4
+ import uuid
5
+
6
+ import httpx
7
+ from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE, BASE_URL
8
+ from fastapi import HTTPException
9
+ from api.models import ChatRequest
10
+
11
+ from api.logger import setup_logger
12
+
13
+ logger = setup_logger(__name__)
14
+
15
+
16
+ def create_chat_completion_data(
17
+ content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
18
+ ) -> Dict[str, Any]:
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": content, "role": "assistant"},
28
+ "finish_reason": finish_reason,
29
+ }
30
+ ],
31
+ "usage": None,
32
+ }
33
+
34
+
35
+ def message_to_dict(message):
36
+ if isinstance(message.content, str):
37
+ return {"role": message.role, "content": message.content}
38
+ elif isinstance(message.content, list) and len(message.content) == 2:
39
+ return {
40
+ "role": message.role,
41
+ "content": message.content[0]["text"],
42
+ "data": {
43
+ "imageBase64": message.content[1]["image_url"]["url"],
44
+ "fileText": "",
45
+ "title": "snapshot",
46
+ },
47
+ }
48
+ else:
49
+ return {"role": message.role, "content": message.content}
50
+
51
+
52
+ async def process_streaming_response(request: ChatRequest):
53
+ agent_mode = AGENT_MODE.get(request.model, {})
54
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
55
+ json_data = {
56
+ "messages": [message_to_dict(msg) for msg in request.messages],
57
+ "previewToken": None,
58
+ "userId": None,
59
+ "codeModelMode": True,
60
+ "agentMode": agent_mode,
61
+ "trendingAgentMode": trending_agent_mode,
62
+ "isMicMode": False,
63
+ "userSystemPrompt": None,
64
+ "maxTokens": request.max_tokens,
65
+ "playgroundTopP": request.top_p,
66
+ "playgroundTemperature": request.temperature,
67
+ "isChromeExt": False,
68
+ "githubToken": None,
69
+ "clickedAnswer2": False,
70
+ "clickedAnswer3": False,
71
+ "clickedForceWebSearch": False,
72
+ "visitFromDelta": False,
73
+ "mobileClient": False,
74
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
75
+ }
76
+
77
+ async with httpx.AsyncClient() as client:
78
+ try:
79
+ async with client.stream(
80
+ "POST",
81
+ f"{BASE_URL}/api/chat",
82
+ headers=headers,
83
+ json=json_data,
84
+ timeout=100,
85
+ ) as response:
86
+ response.raise_for_status()
87
+ async for line in response.aiter_lines():
88
+ timestamp = int(datetime.now().timestamp())
89
+ if line:
90
+ content = line
91
+ if content.startswith("$@$v=undefined-rv1$@$"):
92
+ content = content[21:]
93
+ yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
94
+
95
+ yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
96
+ yield "data: [DONE]\n\n"
97
+ except httpx.HTTPStatusError as e:
98
+ logger.error(f"HTTP error occurred: {e}")
99
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
100
+ except httpx.RequestError as e:
101
+ logger.error(f"Error occurred during request: {e}")
102
+ raise HTTPException(status_code=500, detail=str(e))
103
+
104
+
105
+ async def process_non_streaming_response(request: ChatRequest):
106
+ agent_mode = AGENT_MODE.get(request.model, {})
107
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
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,
114
+ "trendingAgentMode": trending_agent_mode,
115
+ "isMicMode": False,
116
+ "userSystemPrompt": None,
117
+ "maxTokens": request.max_tokens,
118
+ "playgroundTopP": request.top_p,
119
+ "playgroundTemperature": request.temperature,
120
+ "isChromeExt": False,
121
+ "githubToken": None,
122
+ "clickedAnswer2": False,
123
+ "clickedAnswer3": False,
124
+ "clickedForceWebSearch": False,
125
+ "visitFromDelta": False,
126
+ "mobileClient": False,
127
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
128
+ }
129
+ full_response = ""
130
+ async with httpx.AsyncClient() as client:
131
+ try:
132
+ async with client.stream(
133
+ method="POST", url=f"{BASE_URL}/api/chat", headers=headers, json=json_data
134
+ ) as response:
135
+ response.raise_for_status()
136
+ async for chunk in response.aiter_text():
137
+ full_response += chunk
138
+ except httpx.HTTPStatusError as e:
139
+ logger.error(f"HTTP error occurred: {e}")
140
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
141
+ except httpx.RequestError as e:
142
+ logger.error(f"Error occurred during request: {e}")
143
+ raise HTTPException(status_code=500, detail=str(e))
144
+ if full_response.startswith("$@$v=undefined-rv1$@$"):
145
+ full_response = full_response[21:]
146
+
147
+ return {
148
+ "id": f"chatcmpl-{uuid.uuid4()}",
149
+ "object": "chat.completion",
150
+ "created": int(datetime.now().timestamp()),
151
+ "model": request.model,
152
+ "choices": [
153
+ {
154
+ "index": 0,
155
+ "message": {"role": "assistant", "content": full_response},
156
+ "finish_reason": "stop",
157
+ }
158
+ ],
159
+ "usage": None,
160
+ }
161
+ from datetime import datetime
162
+ import json
163
+ from typing import Any, Dict, Optional
164
+ import uuid
165
+
166
+ import httpx
167
+ from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE, BASE_URL
168
+ from fastapi import HTTPException
169
+ from api.models import ChatRequest
170
+
171
+ from api.logger import setup_logger
172
+
173
+ logger = setup_logger(__name__)
174
+
175
+
176
+ def create_chat_completion_data(
177
+ content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
178
+ ) -> Dict[str, Any]:
179
+ return {
180
+ "id": f"chatcmpl-{uuid.uuid4()}",
181
+ "object": "chat.completion.chunk",
182
+ "created": timestamp,
183
+ "model": model,
184
+ "choices": [
185
+ {
186
+ "index": 0,
187
+ "delta": {"content": content, "role": "assistant"},
188
+ "finish_reason": finish_reason,
189
+ }
190
+ ],
191
+ "usage": None,
192
+ }
193
+
194
+
195
+ def message_to_dict(message):
196
+ if isinstance(message.content, str):
197
+ return {"role": message.role, "content": message.content}
198
+ elif isinstance(message.content, list) and len(message.content) == 2:
199
+ return {
200
+ "role": message.role,
201
+ "content": message.content[0]["text"],
202
+ "data": {
203
+ "imageBase64": message.content[1]["image_url"]["url"],
204
+ "fileText": "",
205
+ "title": "snapshot",
206
+ },
207
+ }
208
+ else:
209
+ return {"role": message.role, "content": message.content}
210
+
211
+
212
+ async def process_streaming_response(request: ChatRequest):
213
+ agent_mode = AGENT_MODE.get(request.model, {})
214
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
215
+ json_data = {
216
+ "messages": [message_to_dict(msg) for msg in request.messages],
217
+ "previewToken": None,
218
+ "userId": None,
219
+ "codeModelMode": True,
220
+ "agentMode": agent_mode,
221
+ "trendingAgentMode": trending_agent_mode,
222
+ "isMicMode": False,
223
+ "userSystemPrompt": None,
224
+ "maxTokens": request.max_tokens,
225
+ "playgroundTopP": request.top_p,
226
+ "playgroundTemperature": request.temperature,
227
+ "isChromeExt": False,
228
+ "githubToken": None,
229
+ "clickedAnswer2": False,
230
+ "clickedAnswer3": False,
231
+ "clickedForceWebSearch": False,
232
+ "visitFromDelta": False,
233
+ "mobileClient": False,
234
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
235
+ }
236
+
237
+ async with httpx.AsyncClient() as client:
238
+ try:
239
+ async with client.stream(
240
+ "POST",
241
+ f"{BASE_URL}/api/chat",
242
+ headers=headers,
243
+ json=json_data,
244
+ timeout=100,
245
+ ) as response:
246
+ response.raise_for_status()
247
+ async for line in response.aiter_lines():
248
+ timestamp = int(datetime.now().timestamp())
249
+ if line:
250
+ content = line
251
+ if content.startswith("$@$v=undefined-rv1$@$"):
252
+ content = content[21:]
253
+ yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
254
+
255
+ yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
256
+ yield "data: [DONE]\n\n"
257
+ except httpx.HTTPStatusError as e:
258
+ logger.error(f"HTTP error occurred: {e}")
259
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
260
+ except httpx.RequestError as e:
261
+ logger.error(f"Error occurred during request: {e}")
262
+ raise HTTPException(status_code=500, detail=str(e))
263
+
264
+
265
+ async def process_non_streaming_response(request: ChatRequest):
266
+ agent_mode = AGENT_MODE.get(request.model, {})
267
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
268
+ json_data = {
269
+ "messages": [message_to_dict(msg) for msg in request.messages],
270
+ "previewToken": None,
271
+ "userId": None,
272
+ "codeModelMode": True,
273
+ "agentMode": agent_mode,
274
+ "trendingAgentMode": trending_agent_mode,
275
+ "isMicMode": False,
276
+ "userSystemPrompt": None,
277
+ "maxTokens": request.max_tokens,
278
+ "playgroundTopP": request.top_p,
279
+ "playgroundTemperature": request.temperature,
280
+ "isChromeExt": False,
281
+ "githubToken": None,
282
+ "clickedAnswer2": False,
283
+ "clickedAnswer3": False,
284
+ "clickedForceWebSearch": False,
285
+ "visitFromDelta": False,
286
+ "mobileClient": False,
287
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
288
+ }
289
+ full_response = ""
290
+ async with httpx.AsyncClient() as client:
291
+ try:
292
+ async with client.stream(
293
+ method="POST", url=f"{BASE_URL}/api/chat", headers=headers, json=json_data
294
+ ) as response:
295
+ response.raise_for_status()
296
+ async for chunk in response.aiter_text():
297
+ full_response += chunk
298
+ except httpx.HTTPStatusError as e:
299
+ logger.error(f"HTTP error occurred: {e}")
300
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
301
+ except httpx.RequestError as e:
302
+ logger.error(f"Error occurred during request: {e}")
303
+ raise HTTPException(status_code=500, detail=str(e))
304
+ if full_response.startswith("$@$v=undefined-rv1$@$"):
305
+ full_response = full_response[21:]
306
+
307
+ return {
308
+ "id": f"chatcmpl-{uuid.uuid4()}",
309
+ "object": "chat.completion",
310
+ "created": int(datetime.now().timestamp()),
311
+ "model": request.model,
312
+ "choices": [
313
+ {
314
+ "index": 0,
315
+ "message": {"role": "assistant", "content": full_response},
316
+ "finish_reason": "stop",
317
+ }
318
+ ],
319
+ "usage": None,
320
+ }
321
+ from datetime import datetime
322
+ import json
323
+ from typing import Any, Dict, Optional
324
+ import uuid
325
+
326
+ import httpx
327
+ from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE, BASE_URL
328
+ from fastapi import HTTPException
329
+ from api.models import ChatRequest
330
+
331
+ from api.logger import setup_logger
332
+
333
+ logger = setup_logger(__name__)
334
+
335
+
336
+ def create_chat_completion_data(
337
+ content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
338
+ ) -> Dict[str, Any]:
339
+ return {
340
+ "id": f"chatcmpl-{uuid.uuid4()}",
341
+ "object": "chat.completion.chunk",
342
+ "created": timestamp,
343
+ "model": model,
344
+ "choices": [
345
+ {
346
+ "index": 0,
347
+ "delta": {"content": content, "role": "assistant"},
348
+ "finish_reason": finish_reason,
349
+ }
350
+ ],
351
+ "usage": None,
352
+ }
353
+
354
+
355
+ def message_to_dict(message):
356
+ if isinstance(message.content, str):
357
+ return {"role": message.role, "content": message.content}
358
+ elif isinstance(message.content, list) and len(message.content) == 2:
359
+ return {
360
+ "role": message.role,
361
+ "content": message.content[0]["text"],
362
+ "data": {
363
+ "imageBase64": message.content[1]["image_url"]["url"],
364
+ "fileText": "",
365
+ "title": "snapshot",
366
+ },
367
+ }
368
+ else:
369
+ return {"role": message.role, "content": message.content}
370
+
371
+
372
+ async def process_streaming_response(request: ChatRequest):
373
+ agent_mode = AGENT_MODE.get(request.model, {})
374
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
375
+ json_data = {
376
+ "messages": [message_to_dict(msg) for msg in request.messages],
377
+ "previewToken": None,
378
+ "userId": None,
379
+ "codeModelMode": True,
380
+ "agentMode": agent_mode,
381
+ "trendingAgentMode": trending_agent_mode,
382
+ "isMicMode": False,
383
+ "userSystemPrompt": None,
384
+ "maxTokens": request.max_tokens,
385
+ "playgroundTopP": request.top_p,
386
+ "playgroundTemperature": request.temperature,
387
+ "isChromeExt": False,
388
+ "githubToken": None,
389
+ "clickedAnswer2": False,
390
+ "clickedAnswer3": False,
391
+ "clickedForceWebSearch": False,
392
+ "visitFromDelta": False,
393
+ "mobileClient": False,
394
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
395
+ }
396
+
397
+ async with httpx.AsyncClient() as client:
398
+ try:
399
+ async with client.stream(
400
+ "POST",
401
+ f"{BASE_URL}/api/chat",
402
+ headers=headers,
403
+ json=json_data,
404
+ timeout=100,
405
+ ) as response:
406
+ response.raise_for_status()
407
+ async for line in response.aiter_lines():
408
+ timestamp = int(datetime.now().timestamp())
409
+ if line:
410
+ content = line
411
+ if content.startswith("$@$v=undefined-rv1$@$"):
412
+ content = content[21:]
413
+ yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
414
+
415
+ yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
416
+ yield "data: [DONE]\n\n"
417
+ except httpx.HTTPStatusError as e:
418
+ logger.error(f"HTTP error occurred: {e}")
419
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
420
+ except httpx.RequestError as e:
421
+ logger.error(f"Error occurred during request: {e}")
422
+ raise HTTPException(status_code=500, detail=str(e))
423
+
424
+
425
+ async def process_non_streaming_response(request: ChatRequest):
426
+ agent_mode = AGENT_MODE.get(request.model, {})
427
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
428
+ json_data = {
429
+ "messages": [message_to_dict(msg) for msg in request.messages],
430
+ "previewToken": None,
431
+ "userId": None,
432
+ "codeModelMode": True,
433
+ "agentMode": agent_mode,
434
+ "trendingAgentMode": trending_agent_mode,
435
+ "isMicMode": False,
436
+ "userSystemPrompt": None,
437
+ "maxTokens": request.max_tokens,
438
+ "playgroundTopP": request.top_p,
439
+ "playgroundTemperature": request.temperature,
440
+ "isChromeExt": False,
441
+ "githubToken": None,
442
+ "clickedAnswer2": False,
443
+ "clickedAnswer3": False,
444
+ "clickedForceWebSearch": False,
445
+ "visitFromDelta": False,
446
+ "mobileClient": False,
447
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
448
+ }
449
+ full_response = ""
450
+ async with httpx.AsyncClient() as client:
451
+ try:
452
+ async with client.stream(
453
+ method="POST", url=f"{BASE_URL}/api/chat", headers=headers, json=json_data
454
+ ) as response:
455
+ response.raise_for_status()
456
+ async for chunk in response.aiter_text():
457
+ full_response += chunk
458
+ except httpx.HTTPStatusError as e:
459
+ logger.error(f"HTTP error occurred: {e}")
460
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
461
+ except httpx.RequestError as e:
462
+ logger.error(f"Error occurred during request: {e}")
463
+ raise HTTPException(status_code=500, detail=str(e))
464
+ if full_response.startswith("$@$v=undefined-rv1$@$"):
465
+ full_response = full_response[21:]
466
+
467
+ return {
468
+ "id": f"chatcmpl-{uuid.uuid4()}",
469
+ "object": "chat.completion",
470
+ "created": int(datetime.now().timestamp()),
471
+ "model": request.model,
472
+ "choices": [
473
+ {
474
+ "index": 0,
475
+ "message": {"role": "assistant", "content": full_response},
476
+ "finish_reason": "stop",
477
+ }
478
+ ],
479
+ "usage": None,
480
+ }
requirements.txt CHANGED
@@ -1,7 +1,6 @@
1
- fastapi==0.95.2
2
- httpx==0.23.3
3
- pydantic==1.10.4
4
- python-dotenv==0.21.0
5
- uvicorn==0.21.1
6
- gunicorn==20.1.0
7
- mypy==1.5.1
 
1
+ fastapi==0.95.2
2
+ httpx==0.23.3
3
+ pydantic==1.10.4
4
+ python-dotenv==0.21.0
5
+ uvicorn==0.21.1
6
+ gunicorn==20.1.0