Niansuh commited on
Commit
552bcd7
·
verified ·
1 Parent(s): 923acec

Upload 13 files

Browse files
.github/workflows/docker-deploy.yml CHANGED
@@ -36,7 +36,7 @@ jobs:
36
  with:
37
  context: . # The context is the root of your repository
38
  push: true # Automatically push the image after building
39
- tags: ${{ secrets.DOCKER_USERNAME }}/blackboxv2:v2 # Replace 'your-app-name' with your desired Docker image name
40
 
41
  # Step 5: Log out of Docker Hub
42
  - name: Log out of Docker Hub
 
36
  with:
37
  context: . # The context is the root of your repository
38
  push: true # Automatically push the image after building
39
+ tags: ${{ secrets.DOCKER_USERNAME }}/blackboxv2:v1 # Replace 'your-app-name' with your desired Docker image name
40
 
41
  # Step 5: Log out of Docker Hub
42
  - name: Log out of Docker Hub
Dockerfile CHANGED
@@ -1,5 +1,3 @@
1
- # Dockerfile
2
-
3
  # Use an official Python runtime as a parent image
4
  FROM python:3.10-slim
5
 
 
 
 
1
  # Use an official Python runtime as a parent image
2
  FROM python:3.10-slim
3
 
api/__init__.py ADDED
File without changes
api/__pycache__/dummy.txt ADDED
@@ -0,0 +1 @@
 
 
1
+
api/app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ logger = setup_logger(__name__)
8
+
9
+ def create_app():
10
+ app = FastAPI(
11
+ title="NiansuhAI API Gateway",
12
+ docs_url=None, # Disable Swagger UI
13
+ redoc_url=None, # Disable ReDoc
14
+ openapi_url=None, # Disable OpenAPI schema
15
+ )
16
+
17
+ # CORS settings
18
+ app.add_middleware(
19
+ CORSMiddleware,
20
+ allow_origins=["*"], # Adjust as needed for security
21
+ allow_credentials=True,
22
+ allow_methods=["*"],
23
+ allow_headers=["*"],
24
+ )
25
+
26
+ # Include routes
27
+ app.include_router(router)
28
+
29
+ # Global exception handler for better error reporting
30
+ @app.exception_handler(Exception)
31
+ async def global_exception_handler(request: Request, exc: Exception):
32
+ logger.error(f"An error occurred: {str(exc)}")
33
+ return JSONResponse(
34
+ status_code=500,
35
+ content={"message": "An internal server error occurred."},
36
+ )
37
+
38
+ return app
39
+
40
+ 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,231 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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': "NiansuhAIk1HgESy", '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
+ 'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
187
+ 'PythonAgent': '@Python Agent',
188
+ 'JavaAgent': '@Java Agent',
189
+ 'JavaScriptAgent': '@JavaScript Agent',
190
+ 'HTMLAgent': '@HTML Agent',
191
+ 'GoogleCloudAgent': '@Google Cloud Agent',
192
+ 'AndroidDeveloper': '@Android Developer',
193
+ 'SwiftDeveloper': '@Swift Developer',
194
+ 'Next.jsAgent': '@Next.js Agent',
195
+ 'MongoDBAgent': '@MongoDB Agent',
196
+ 'PyTorchAgent': '@PyTorch Agent',
197
+ 'ReactAgent': '@React Agent',
198
+ 'XcodeAgent': '@Xcode Agent',
199
+ 'AngularJSAgent': '@AngularJS Agent',
200
+ 'HerokuAgent': '@Heroku Agent',
201
+ 'GodotAgent': '@Godot Agent',
202
+ 'GoAgent': '@Go Agent',
203
+ 'GitlabAgent': '@Gitlab Agent',
204
+ 'GitAgent': '@Gitlab Agent',
205
+ 'blackboxai-pro': '@BLACKBOXAI-PRO',
206
+ 'flux': '@Image Generation',
207
+ # Add any additional prefixes if necessary
208
+
209
+ # Added New Agents
210
+ 'FlaskAgent': '@Flask Agent',
211
+ 'FirebaseAgent': '@Firebase Agent',
212
+ 'FastAPIAgent': '@FastAPI Agent',
213
+ 'ErlangAgent': '@Erlang Agent',
214
+ 'ElectronAgent': '@Electron Agent',
215
+ 'DockerAgent': '@Docker Agent',
216
+ 'DigitalOceanAgent': '@DigitalOcean Agent',
217
+ 'BitbucketAgent': '@Bitbucket Agent',
218
+ 'AzureAgent': '@Azure Agent',
219
+ 'FlutterAgent': '@Flutter Agent',
220
+ 'YoutubeAgent': '@Youtube Agent',
221
+ 'builderAgent': '@builder Agent',
222
+ }
223
+
224
+ # Model referers
225
+ MODEL_REFERERS = {
226
+ "blackboxai": "/?model=blackboxai",
227
+ "gpt-4o": "/?model=gpt-4o",
228
+ "gemini-pro": "/?model=gemini-pro",
229
+ "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
230
+ # Add any additional referers if necessary
231
+ }
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.7
13
+ top_p: Optional[float] = 0.9
14
+ max_tokens: Optional[int] = 99999999
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,200 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
22
+ logger = setup_logger(__name__)
23
+
24
+ def create_chat_completion_data(
25
+ content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
26
+ ) -> Dict[str, Any]:
27
+ return {
28
+ "id": f"chatcmpl-{uuid.uuid4()}",
29
+ "object": "chat.completion.chunk",
30
+ "created": timestamp,
31
+ "model": model,
32
+ "choices": [
33
+ {
34
+ "index": 0,
35
+ "delta": {"content": content, "role": "assistant"},
36
+ "finish_reason": finish_reason,
37
+ }
38
+ ],
39
+ "usage": None,
40
+ }
41
+
42
+ def message_to_dict(message, model_prefix: Optional[str] = None):
43
+ if isinstance(message.content, str):
44
+ content = message.content
45
+ if model_prefix:
46
+ content = f"{model_prefix} {content}"
47
+ return {"role": message.role, "content": content}
48
+ elif isinstance(message.content, list) and len(message.content) == 2:
49
+ content = message.content[0]["text"]
50
+ if model_prefix:
51
+ content = f"{model_prefix} {content}"
52
+ return {
53
+ "role": message.role,
54
+ "content": content,
55
+ "data": {
56
+ "imageBase64": message.content[1]["image_url"]["url"],
57
+ "fileText": "",
58
+ "title": "snapshot",
59
+ },
60
+ }
61
+ else:
62
+ return {"role": message.role, "content": message.content}
63
+
64
+ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
65
+ """Remove the model prefix from the response content if present."""
66
+ if model_prefix and content.startswith(model_prefix):
67
+ logger.debug(f"Stripping prefix '{model_prefix}' from content.")
68
+ return content[len(model_prefix):].strip()
69
+ logger.debug("No prefix to strip from content.")
70
+ return content
71
+
72
+ async def process_streaming_response(request: ChatRequest):
73
+ agent_mode = AGENT_MODE.get(request.model, {})
74
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
75
+ model_prefix = MODEL_PREFIXES.get(request.model, "")
76
+ referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
77
+ referer_url = f"{BASE_URL}{referer_path}"
78
+
79
+ # Update headers with dynamic Referer
80
+ dynamic_headers = headers.copy()
81
+ dynamic_headers['Referer'] = referer_url
82
+
83
+ json_data = {
84
+ "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
85
+ "previewToken": None,
86
+ "userId": None,
87
+ "codeModelMode": True,
88
+ "agentMode": agent_mode,
89
+ "trendingAgentMode": trending_agent_mode,
90
+ "isMicMode": False,
91
+ "userSystemPrompt": None,
92
+ "maxTokens": request.max_tokens,
93
+ "playgroundTopP": request.top_p,
94
+ "playgroundTemperature": request.temperature,
95
+ "isChromeExt": False,
96
+ "githubToken": None,
97
+ "clickedAnswer2": False,
98
+ "clickedAnswer3": False,
99
+ "clickedForceWebSearch": False,
100
+ "visitFromDelta": False,
101
+ "mobileClient": False,
102
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
103
+ }
104
+
105
+ async with httpx.AsyncClient() as client:
106
+ try:
107
+ async with client.stream(
108
+ "POST",
109
+ f"{BASE_URL}/api/chat",
110
+ headers=dynamic_headers,
111
+ json=json_data,
112
+ timeout=100,
113
+ ) as response:
114
+ response.raise_for_status()
115
+ async for line in response.aiter_lines():
116
+ timestamp = int(datetime.now().timestamp())
117
+ if line:
118
+ content = line
119
+ if content.startswith("$@$v=undefined-rv1$@$"):
120
+ content = content[21:]
121
+ # Strip the model prefix from the response content
122
+ cleaned_content = strip_model_prefix(content, model_prefix)
123
+ yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
124
+
125
+ yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
126
+ yield "data: [DONE]\n\n"
127
+ except httpx.HTTPStatusError as e:
128
+ logger.error(f"HTTP error occurred: {e}")
129
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
130
+ except httpx.RequestError as e:
131
+ logger.error(f"Error occurred during request: {e}")
132
+ raise HTTPException(status_code=500, detail=str(e))
133
+
134
+ async def process_non_streaming_response(request: ChatRequest):
135
+ agent_mode = AGENT_MODE.get(request.model, {})
136
+ trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
137
+ model_prefix = MODEL_PREFIXES.get(request.model, "")
138
+ referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
139
+ referer_url = f"{BASE_URL}{referer_path}"
140
+
141
+ # Update headers with dynamic Referer
142
+ dynamic_headers = headers.copy()
143
+ dynamic_headers['Referer'] = referer_url
144
+
145
+ json_data = {
146
+ "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
147
+ "previewToken": None,
148
+ "userId": None,
149
+ "codeModelMode": True,
150
+ "agentMode": agent_mode,
151
+ "trendingAgentMode": trending_agent_mode,
152
+ "isMicMode": False,
153
+ "userSystemPrompt": None,
154
+ "maxTokens": request.max_tokens,
155
+ "playgroundTopP": request.top_p,
156
+ "playgroundTemperature": request.temperature,
157
+ "isChromeExt": False,
158
+ "githubToken": None,
159
+ "clickedAnswer2": False,
160
+ "clickedAnswer3": False,
161
+ "clickedForceWebSearch": False,
162
+ "visitFromDelta": False,
163
+ "mobileClient": False,
164
+ "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
165
+ }
166
+ full_response = ""
167
+ async with httpx.AsyncClient() as client:
168
+ try:
169
+ async with client.stream(
170
+ method="POST", url=f"{BASE_URL}/api/chat", headers=dynamic_headers, json=json_data
171
+ ) as response:
172
+ response.raise_for_status()
173
+ async for chunk in response.aiter_text():
174
+ full_response += chunk
175
+ except httpx.HTTPStatusError as e:
176
+ logger.error(f"HTTP error occurred: {e}")
177
+ raise HTTPException(status_code=e.response.status_code, detail=str(e))
178
+ except httpx.RequestError as e:
179
+ logger.error(f"Error occurred during request: {e}")
180
+ raise HTTPException(status_code=500, detail=str(e))
181
+ if full_response.startswith("$@$v=undefined-rv1$@$"):
182
+ full_response = full_response[21:]
183
+
184
+ # Strip the model prefix from the full response
185
+ cleaned_full_response = strip_model_prefix(full_response, model_prefix)
186
+
187
+ return {
188
+ "id": f"chatcmpl-{uuid.uuid4()}",
189
+ "object": "chat.completion",
190
+ "created": int(datetime.now().timestamp()),
191
+ "model": request.model,
192
+ "choices": [
193
+ {
194
+ "index": 0,
195
+ "message": {"role": "assistant", "content": cleaned_full_response},
196
+ "finish_reason": "stop",
197
+ }
198
+ ],
199
+ "usage": None,
200
+ }
main.py CHANGED
@@ -1,7 +1,5 @@
1
- # main.py
2
-
3
- import uvicorn
4
- from api.app import app
5
-
6
- if __name__ == "__main__":
7
- uvicorn.run(app, host="0.0.0.0", port=8001)
 
1
+ import uvicorn
2
+ from api.app import app
3
+
4
+ if __name__ == "__main__":
5
+ uvicorn.run(app, host="0.0.0.0", port=8001)
 
 
requirements.txt CHANGED
@@ -1,6 +1,6 @@
1
- fastapi==0.95.2
2
- pydantic==1.10.4
3
- python-dotenv==0.21.0
4
- uvicorn==0.21.1
5
- gunicorn==20.1.0
6
- aiohttp==3.8.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