Niansuh commited on
Commit
d96bc9d
·
verified ·
1 Parent(s): d2222ab

Upload 13 files

Browse files
Files changed (11) hide show
  1. Dockerfile +1 -1
  2. api/__init__.py +0 -0
  3. api/__pycache__/dummy.txt +1 -0
  4. api/app.py +41 -0
  5. api/auth.py +10 -0
  6. api/config.py +230 -0
  7. api/logger.py +20 -0
  8. api/models.py +14 -0
  9. api/routes.py +59 -0
  10. api/utils.py +216 -0
  11. requirements.txt +6 -7
Dockerfile CHANGED
@@ -23,4 +23,4 @@ COPY . /app
23
  EXPOSE 8001
24
 
25
  # Command to run the app with Gunicorn and Uvicorn workers
26
- CMD ["gunicorn", "-k", "uvicorn.workers.UvicornWorker", "--workers", "4", "--bind", "0.0.0.0:8001", "main:app"]
 
23
  EXPOSE 8001
24
 
25
  # Command to run the app with Gunicorn and Uvicorn workers
26
+ CMD ["gunicorn", "-k", "uvicorn.workers.UvicornWorker", "--workers", "4", "--bind", "0.0.0.0:8001", "main:app"]
api/__init__.py ADDED
File without changes
api/__pycache__/dummy.txt ADDED
@@ -0,0 +1 @@
 
 
1
+
api/app.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, Request
2
+ from starlette.middleware.cors import CORSMiddleware
3
+ from fastapi.responses import JSONResponse
4
+ from api.logger import setup_logger
5
+ from api.routes import router
6
+
7
+
8
+ logger = setup_logger(__name__)
9
+
10
+ def create_app():
11
+ app = FastAPI(
12
+ title="NiansuhAI API Gateway",
13
+ docs_url=None, # Disable Swagger UI
14
+ redoc_url=None, # Disable ReDoc
15
+ openapi_url=None, # Disable OpenAPI schema
16
+ )
17
+
18
+ # CORS settings
19
+ app.add_middleware(
20
+ CORSMiddleware,
21
+ allow_origins=["*"], # Adjust as needed for security
22
+ allow_credentials=True,
23
+ allow_methods=["*"],
24
+ allow_headers=["*"],
25
+ )
26
+
27
+ # Include routes
28
+ app.include_router(router)
29
+
30
+ # Global exception handler for better error reporting
31
+ @app.exception_handler(Exception)
32
+ async def global_exception_handler(request: Request, exc: Exception):
33
+ logger.error(f"An error occurred: {str(exc)}")
34
+ return JSONResponse(
35
+ status_code=500,
36
+ content={"message": "An internal server error occurred."},
37
+ )
38
+
39
+ return app
40
+
41
+ app = create_app()
api/auth.py ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import Depends, HTTPException
2
+ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
3
+ from api.config import APP_SECRET
4
+
5
+ security = HTTPBearer()
6
+
7
+ def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)):
8
+ if credentials.credentials != APP_SECRET:
9
+ raise HTTPException(status_code=403, detail="Invalid APP_SECRET")
10
+ return credentials.credentials
api/config.py ADDED
@@ -0,0 +1,230 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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': '"Chromium";v="130", "Google Chrome";v="130", "Not?A_Brand";v="99"',
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/130.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": "HerokuAgent", "name": "HerokuAgent"},
49
+ {"id": "GodotAgent", "name": "GodotAgent"},
50
+ {"id": "GoAgent", "name": "GoAgent"},
51
+ {"id": "GitlabAgent", "name": "GitlabAgent"},
52
+ {"id": "GitAgent", "name": "GitAgent"},
53
+ {"id": "RepoMap", "name": "RepoMap"},
54
+ {"id": "gemini-1.5-pro-latest", "name": "gemini-pro"},
55
+ {"id": "gemini-1.5-pro", "name": "gemini-1.5-pro"},
56
+ {"id": "claude-3-5-sonnet-20240620", "name": "claude-sonnet-3.5"},
57
+ {"id": "claude-3-5-sonnet", "name": "claude-sonnet-3.5"},
58
+ {"id": "Niansuh", "name": "Niansuh"},
59
+ {"id": "o1-preview", "name": "o1-preview"},
60
+
61
+ # Added New Agents
62
+ {"id": "FlaskAgent", "name": "FlaskAgent"},
63
+ {"id": "FirebaseAgent", "name": "FirebaseAgent"},
64
+ {"id": "FastAPIAgent", "name": "FastAPIAgent"},
65
+ {"id": "ErlangAgent", "name": "ErlangAgent"},
66
+ {"id": "ElectronAgent", "name": "ElectronAgent"},
67
+ {"id": "DockerAgent", "name": "DockerAgent"},
68
+ {"id": "DigitalOceanAgent", "name": "DigitalOceanAgent"},
69
+ {"id": "BitbucketAgent", "name": "BitbucketAgent"},
70
+ {"id": "AzureAgent", "name": "AzureAgent"},
71
+ {"id": "FlutterAgent", "name": "FlutterAgent"},
72
+ {"id": "YoutubeAgent", "name": "YoutubeAgent"},
73
+ {"id": "builderAgent", "name": "builderAgent"},
74
+ ]
75
+
76
+ MODEL_MAPPING = {
77
+ "blackboxai": "blackboxai",
78
+ "blackboxai-pro": "blackboxai-pro",
79
+ "flux": "flux",
80
+ "ImageGeneration": "flux",
81
+ "llama-3.1-8b": "llama-3.1-8b",
82
+ "llama-3.1-70b": "llama-3.1-70b",
83
+ "llama-3.1-405b": "llama-3.1-405b",
84
+ "gpt-4o": "gpt-4o",
85
+ "gemini-pro": "gemini-pro",
86
+ "gemini-1.5-flash": "gemini-1.5-flash",
87
+ "claude-sonnet-3.5": "claude-sonnet-3.5",
88
+ "PythonAgent": "PythonAgent",
89
+ "JavaAgent": "JavaAgent",
90
+ "JavaScriptAgent": "JavaScriptAgent",
91
+ "HTMLAgent": "HTMLAgent",
92
+ "GoogleCloudAgent": "GoogleCloudAgent",
93
+ "AndroidDeveloper": "AndroidDeveloper",
94
+ "SwiftDeveloper": "SwiftDeveloper",
95
+ "Next.jsAgent": "Next.jsAgent",
96
+ "MongoDBAgent": "MongoDBAgent",
97
+ "PyTorchAgent": "PyTorchAgent",
98
+ "ReactAgent": "ReactAgent",
99
+ "XcodeAgent": "XcodeAgent",
100
+ "AngularJSAgent": "AngularJSAgent",
101
+ "HerokuAgent": "HerokuAgent",
102
+ "GodotAgent": "GodotAgent",
103
+ "GoAgent": "GoAgent",
104
+ "GitlabAgent": "GitlabAgent",
105
+ "GitAgent": "GitAgent",
106
+ "RepoMap": "RepoMap",
107
+ # Additional mappings
108
+ "gemini-flash": "gemini-1.5-flash",
109
+ "claude-3.5-sonnet": "claude-sonnet-3.5",
110
+ "flux": "flux",
111
+ "gemini-1.5-pro-latest": "gemini-pro",
112
+ "gemini-1.5-pro": "gemini-1.5-pro",
113
+ "claude-3-5-sonnet-20240620": "claude-sonnet-3.5",
114
+ "claude-3-5-sonnet": "claude-sonnet-3.5",
115
+ "Niansuh": "Niansuh",
116
+ "o1-preview": "o1-preview",
117
+
118
+ # Added New Agents
119
+ "FlaskAgent": "FlaskAgent",
120
+ "FirebaseAgent": "FirebaseAgent",
121
+ "FastAPIAgent": "FastAPIAgent",
122
+ "ErlangAgent": "ErlangAgent",
123
+ "ElectronAgent": "ElectronAgent",
124
+ "DockerAgent": "DockerAgent",
125
+ "DigitalOceanAgent": "DigitalOceanAgent",
126
+ "BitbucketAgent": "BitbucketAgent",
127
+ "AzureAgent": "AzureAgent",
128
+ "FlutterAgent": "FlutterAgent",
129
+ "YoutubeAgent": "YoutubeAgent",
130
+ "builderAgent": "builderAgent",
131
+ }
132
+
133
+ # Agent modes
134
+ AGENT_MODE = {
135
+ 'flux': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "flux"},
136
+ 'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
137
+ 'o1-preview': {'mode': True, 'id': "o1Dst8La8", 'name': "o1-preview"},
138
+ }
139
+
140
+ TRENDING_AGENT_MODE = {
141
+ "blackboxai": {},
142
+ "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
143
+ "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
144
+ 'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
145
+ 'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
146
+ 'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
147
+ 'PythonAgent': {'mode': True, 'id': "Python Agent"},
148
+ 'JavaAgent': {'mode': True, 'id': "Java Agent"},
149
+ 'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
150
+ 'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
151
+ 'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
152
+ 'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
153
+ 'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
154
+ 'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
155
+ 'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
156
+ 'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
157
+ 'ReactAgent': {'mode': True, 'id': "React Agent"},
158
+ 'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
159
+ 'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
160
+ 'HerokuAgent': {'mode': True, 'id': "HerokuAgent"},
161
+ 'GodotAgent': {'mode': True, 'id': "GodotAgent"},
162
+ 'GoAgent': {'mode': True, 'id': "GoAgent"},
163
+ 'GitlabAgent': {'mode': True, 'id': "GitlabAgent"},
164
+ 'GitAgent': {'mode': True, 'id': "GitAgent"},
165
+ 'RepoMap': {'mode': True, 'id': "repomap"},
166
+
167
+ # Added New Agents
168
+ 'FlaskAgent': {'mode': True, 'id': "FlaskAgentTrendID"},
169
+ 'FirebaseAgent': {'mode': True, 'id': "FirebaseAgentTrendID"},
170
+ 'FastAPIAgent': {'mode': True, 'id': "FastAPIAgentTrendID"},
171
+ 'ErlangAgent': {'mode': True, 'id': "ErlangAgentTrendID"},
172
+ 'ElectronAgent': {'mode': True, 'id': "ElectronAgentTrendID"},
173
+ 'DockerAgent': {'mode': True, 'id': "DockerAgentTrendID"},
174
+ 'DigitalOceanAgent': {'mode': True, 'id': "DigitalOceanAgentTrendID"},
175
+ 'BitbucketAgent': {'mode': True, 'id': "BitbucketAgentTrendID"},
176
+ 'AzureAgent': {'mode': True, 'id': "AzureAgentTrendID"},
177
+ 'FlutterAgent': {'mode': True, 'id': "FlutterAgentTrendID"},
178
+ 'YoutubeAgent': {'mode': True, 'id': "YoutubeAgentTrendID"},
179
+ 'builderAgent': {'mode': True, 'id': "builderAgentTrendID"},
180
+ }
181
+
182
+ # Model prefixes
183
+ MODEL_PREFIXES = {
184
+ 'gpt-4o': '@GPT-4o',
185
+ 'gemini-pro': '@Gemini-PRO',
186
+ 'PythonAgent': '@Python Agent',
187
+ 'JavaAgent': '@Java Agent',
188
+ 'JavaScriptAgent': '@JavaScript Agent',
189
+ 'HTMLAgent': '@HTML Agent',
190
+ 'GoogleCloudAgent': '@Google Cloud Agent',
191
+ 'AndroidDeveloper': '@Android Developer',
192
+ 'SwiftDeveloper': '@Swift Developer',
193
+ 'Next.jsAgent': '@Next.js Agent',
194
+ 'MongoDBAgent': '@MongoDB Agent',
195
+ 'PyTorchAgent': '@PyTorch Agent',
196
+ 'ReactAgent': '@React Agent',
197
+ 'XcodeAgent': '@Xcode Agent',
198
+ 'AngularJSAgent': '@AngularJS Agent',
199
+ 'HerokuAgent': '@Heroku Agent',
200
+ 'GodotAgent': '@Godot Agent',
201
+ 'GoAgent': '@Go Agent',
202
+ 'GitlabAgent': '@Gitlab Agent',
203
+ 'GitAgent': '@Gitlab Agent',
204
+ 'blackboxai-pro': '@BLACKBOXAI-PRO',
205
+ 'flux': '@Image Generation',
206
+ # Add any additional prefixes if necessary
207
+
208
+ # Added New Agents
209
+ 'FlaskAgent': '@Flask Agent',
210
+ 'FirebaseAgent': '@Firebase Agent',
211
+ 'FastAPIAgent': '@FastAPI Agent',
212
+ 'ErlangAgent': '@Erlang Agent',
213
+ 'ElectronAgent': '@Electron Agent',
214
+ 'DockerAgent': '@Docker Agent',
215
+ 'DigitalOceanAgent': '@DigitalOcean Agent',
216
+ 'BitbucketAgent': '@Bitbucket Agent',
217
+ 'AzureAgent': '@Azure Agent',
218
+ 'FlutterAgent': '@Flutter Agent',
219
+ 'YoutubeAgent': '@Youtube Agent',
220
+ 'builderAgent': '@builder Agent',
221
+ }
222
+
223
+ # Model referers
224
+ MODEL_REFERERS = {
225
+ "blackboxai": "/?model=blackboxai",
226
+ "gpt-4o": "/?model=gpt-4o",
227
+ "gemini-pro": "/?model=gemini-pro",
228
+ "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
229
+ # Add any additional referers if necessary
230
+ }
api/logger.py ADDED
@@ -0,0 +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
api/models.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.5
13
+ top_p: Optional[float] = 0.9
14
+ max_tokens: Optional[int] = 1024
api/routes.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ @router.route('/')
53
+ @router.route('/healthz')
54
+ @router.route('/ready')
55
+ @router.route('/alive')
56
+ @router.route('/status')
57
+ @router.get("/health")
58
+ def health_check(request: Request):
59
+ return Response(content=json.dumps({"status": "ok"}), media_type="application/json")
api/utils.py ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datetime import datetime
2
+ import json
3
+ from typing import Any, Dict, Optional
4
+
5
+ import httpx
6
+ from api.config import (
7
+ MODEL_MAPPING,
8
+ headers,
9
+ AGENT_MODE,
10
+ TRENDING_AGENT_MODE,
11
+ BASE_URL,
12
+ MODEL_PREFIXES,
13
+ MODEL_REFERERS
14
+ )
15
+ from fastapi import HTTPException
16
+ from api.models import ChatRequest
17
+
18
+ from api.logger import setup_logger
19
+
20
+ import uuid
21
+ import asyncio
22
+ import random # Newly added imports
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: Optional[str] = None) -> str:
67
+ """Remove the model prefix from the response content if present."""
68
+ if model_prefix and content.startswith(model_prefix):
69
+ logger.debug(f"Stripping prefix '{model_prefix}' from content.")
70
+ return content[len(model_prefix):].strip()
71
+ logger.debug("No prefix to strip from content.")
72
+ return content
73
+
74
+ async def process_streaming_response(request: ChatRequest):
75
+ agent_mode = AGENT_MODE.get(request.model, {})
76
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
77
+ model_prefix = MODEL_PREFIXES.get(request.model, "")
78
+ referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
79
+ referer_url = f"{BASE_URL}{referer_path}"
80
+
81
+ # Update headers with dynamic Referer
82
+ dynamic_headers = headers.copy()
83
+ dynamic_headers['Referer'] = referer_url
84
+
85
+ # Introduce delay for 'o1-preview' model
86
+ if request.model == 'o1-preview':
87
+ delay_seconds = random.randint(1, 60)
88
+ logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview'")
89
+ await asyncio.sleep(delay_seconds)
90
+
91
+ json_data = {
92
+ "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
93
+ "previewToken": None,
94
+ "userId": None,
95
+ "codeModelMode": True,
96
+ "agentMode": agent_mode,
97
+ "trendingAgentMode": trending_agent_mode,
98
+ "isMicMode": False,
99
+ "userSystemPrompt": None,
100
+ "maxTokens": request.max_tokens,
101
+ "playgroundTopP": request.top_p,
102
+ "playgroundTemperature": request.temperature,
103
+ "isChromeExt": False,
104
+ "githubToken": None,
105
+ "clickedAnswer2": False,
106
+ "clickedAnswer3": False,
107
+ "clickedForceWebSearch": False,
108
+ "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
109
+ "visitFromDelta": False,
110
+ "mobileClient": False,
111
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
112
+ }
113
+
114
+ async with httpx.AsyncClient() as client:
115
+ try:
116
+ async with client.stream(
117
+ "POST",
118
+ f"{BASE_URL}/api/chat",
119
+ headers=dynamic_headers,
120
+ json=json_data,
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 from the response content
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))
139
+ except httpx.RequestError as e:
140
+ logger.error(f"Error occurred during request: {e}")
141
+ raise HTTPException(status_code=500, detail=str(e))
142
+
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
+
150
+ # Update headers with dynamic Referer
151
+ dynamic_headers = headers.copy()
152
+ dynamic_headers['Referer'] = referer_url
153
+
154
+ # Introduce delay for 'o1-preview' model
155
+ if request.model == 'o1-preview':
156
+ delay_seconds = random.randint(20, 60)
157
+ logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview'")
158
+ await asyncio.sleep(delay_seconds)
159
+
160
+ json_data = {
161
+ "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
162
+ "previewToken": None,
163
+ "userId": None,
164
+ "codeModelMode": True,
165
+ "agentMode": agent_mode,
166
+ "trendingAgentMode": trending_agent_mode,
167
+ "isMicMode": False,
168
+ "userSystemPrompt": None,
169
+ "maxTokens": request.max_tokens,
170
+ "playgroundTopP": request.top_p,
171
+ "playgroundTemperature": request.temperature,
172
+ "isChromeExt": False,
173
+ "githubToken": None,
174
+ "clickedAnswer2": False,
175
+ "clickedAnswer3": False,
176
+ "clickedForceWebSearch": False,
177
+ "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
178
+ "visitFromDelta": False,
179
+ "mobileClient": False,
180
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
181
+ }
182
+ full_response = ""
183
+ async with httpx.AsyncClient() as client:
184
+ try:
185
+ async with client.stream(
186
+ method="POST", url=f"{BASE_URL}/api/chat", headers=dynamic_headers, json=json_data
187
+ ) as response:
188
+ response.raise_for_status()
189
+ async for chunk in response.aiter_text():
190
+ full_response += chunk
191
+ except httpx.HTTPStatusError as e:
192
+ logger.error(f"HTTP error occurred: {e}")
193
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
194
+ except httpx.RequestError as e:
195
+ logger.error(f"Error occurred during request: {e}")
196
+ raise HTTPException(status_code=500, detail=str(e))
197
+ if full_response.startswith("$@$v=undefined-rv1$@$"):
198
+ full_response = full_response[21:]
199
+
200
+ # Strip the model prefix from the full response
201
+ cleaned_full_response = strip_model_prefix(full_response, model_prefix)
202
+
203
+ return {
204
+ "id": f"chatcmpl-{uuid.uuid4()}",
205
+ "object": "chat.completion",
206
+ "created": int(datetime.now().timestamp()),
207
+ "model": request.model,
208
+ "choices": [
209
+ {
210
+ "index": 0,
211
+ "message": {"role": "assistant", "content": cleaned_full_response},
212
+ "finish_reason": "stop",
213
+ }
214
+ ],
215
+ "usage": None,
216
+ }
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
- aiohttp==3.8.4
 
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