Niansuh commited on
Commit
ec71f25
·
verified ·
1 Parent(s): 40c5921

Upload 20 files

Browse files
Files changed (10) hide show
  1. Dockerfile +8 -39
  2. api/app.py +40 -40
  3. api/auth.py +10 -10
  4. api/config.py +115 -173
  5. api/logger.py +20 -27
  6. api/models.py +14 -14
  7. api/routes.py +60 -71
  8. api/utils.py +480 -197
  9. main.py +5 -5
  10. requirements.txt +6 -6
Dockerfile CHANGED
@@ -1,57 +1,26 @@
1
- # Stage 1: Builder
2
- FROM python:3.10-slim AS builder
3
 
4
  # Install system dependencies
5
  RUN apt-get update && apt-get install -y --no-install-recommends \
6
  build-essential \
7
  && rm -rf /var/lib/apt/lists/*
8
 
9
- # Set environment variables
10
- ENV PYTHONUNBUFFERED=1
11
- ENV PYTHONDONTWRITEBYTECODE=1
12
-
13
- # Set the working directory
14
  WORKDIR /app
15
 
16
- # Copy the requirements file first for better caching
17
  COPY requirements.txt /app/
18
 
19
- # Install dependencies system-wide
20
  RUN pip install --no-cache-dir --upgrade pip
21
  RUN pip install --no-cache-dir -r requirements.txt
22
 
23
- # Stage 2: Production
24
- FROM python:3.10-slim
25
-
26
- # Install system dependencies required for production
27
- RUN apt-get update && apt-get install -y --no-install-recommends \
28
- libpq-dev \
29
- && rm -rf /var/lib/apt/lists/*
30
-
31
- # Create a non-root user and group
32
- RUN groupadd -r appuser && useradd -r -g appuser appuser
33
-
34
- # Set environment variables
35
- ENV PYTHONUNBUFFERED=1
36
- ENV PYTHONDONTWRITEBYTECODE=1
37
-
38
- # Set the working directory
39
- WORKDIR /app
40
-
41
- # Copy installed Python packages from the builder stage
42
- COPY --from=builder /usr/local /usr/local
43
-
44
  # Copy the current directory contents into the container
45
  COPY . /app
46
 
47
- # Change ownership to the non-root user
48
- RUN chown -R appuser:appuser /app
49
-
50
- # Switch to the non-root user
51
- USER appuser
52
-
53
- # Expose the port
54
  EXPOSE 8001
55
 
56
- # Run gunicorn
57
- CMD ["gunicorn", "-k", "uvicorn.workers.UvicornWorker", "--workers", "4", "--bind", "0.0.0.0:8001", "main:app"]
 
1
+ # Use an official Python runtime as a parent image
2
+ FROM python:3.10-slim
3
 
4
  # Install system dependencies
5
  RUN apt-get update && apt-get install -y --no-install-recommends \
6
  build-essential \
7
  && rm -rf /var/lib/apt/lists/*
8
 
9
+ # Set the working directory in the container
 
 
 
 
10
  WORKDIR /app
11
 
12
+ # Copy the requirements file first
13
  COPY requirements.txt /app/
14
 
15
+ # Install dependencies
16
  RUN pip install --no-cache-dir --upgrade pip
17
  RUN pip install --no-cache-dir -r requirements.txt
18
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  # Copy the current directory contents into the container
20
  COPY . /app
21
 
22
+ # Expose the port that the FastAPI app runs on
 
 
 
 
 
 
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/app.py CHANGED
@@ -1,40 +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()
 
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 CHANGED
@@ -1,10 +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
 
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 CHANGED
@@ -1,173 +1,115 @@
1
- # api/config.py
2
-
3
- import os
4
- from dotenv import load_dotenv
5
-
6
- load_dotenv()
7
-
8
- BASE_URL = "https://www.blackbox.ai"
9
- headers = {
10
- 'accept': '*/*',
11
- 'accept-language': 'en-US,en;q=0.9',
12
- 'origin': 'https://www.blackbox.ai',
13
- 'cache-control': 'no-cache',
14
- 'pragma': 'no-cache',
15
- 'priority': 'u=1, i',
16
- 'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
17
- 'sec-ch-ua-mobile': '?0',
18
- 'sec-ch-ua-platform': '"Linux"',
19
- 'sec-fetch-dest': 'empty',
20
- 'sec-fetch-mode': 'cors',
21
- 'sec-fetch-site': 'same-origin',
22
- 'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
23
- 'AppleWebKit/537.36 (KHTML, like Gecko) '
24
- 'Chrome/129.0.0.0 Safari/537.36'
25
- }
26
-
27
- APP_SECRET = os.getenv("APP_SECRET")
28
-
29
- # Define all allowed models, including agent models
30
- ALLOWED_MODELS = [
31
- {"id": "blackboxai", "name": "blackboxai"},
32
- {"id": "blackboxai-pro", "name": "blackboxai-pro"},
33
- {"id": "ImageGeneration", "name": "ImageGeneration"},
34
- {"id": "llama-3.1-8b", "name": "llama-3.1-8b"},
35
- {"id": "llama-3.1-70b", "name": "llama-3.1-70b"},
36
- {"id": "llama-3.1-405b", "name": "llama-3.1-405b"},
37
- {"id": "gpt-4o", "name": "gpt-4o"},
38
- {"id": "gemini-pro", "name": "gemini-pro"},
39
- {"id": "gemini-1.5-flash", "name": "gemini-1.5-flash"},
40
- {"id": "claude-sonnet-3.5", "name": "claude-sonnet-3.5"},
41
- {"id": "PythonAgent", "name": "PythonAgent"},
42
- {"id": "JavaAgent", "name": "JavaAgent"},
43
- {"id": "JavaScriptAgent", "name": "JavaScriptAgent"},
44
- {"id": "HTMLAgent", "name": "HTMLAgent"},
45
- {"id": "GoogleCloudAgent", "name": "GoogleCloudAgent"},
46
- {"id": "AndroidDeveloper", "name": "AndroidDeveloper"},
47
- {"id": "SwiftDeveloper", "name": "SwiftDeveloper"},
48
- {"id": "Next.jsAgent", "name": "Next.jsAgent"},
49
- {"id": "MongoDBAgent", "name": "MongoDBAgent"},
50
- {"id": "PyTorchAgent", "name": "PyTorchAgent"},
51
- {"id": "ReactAgent", "name": "ReactAgent"},
52
- {"id": "XcodeAgent", "name": "XcodeAgent"},
53
- {"id": "AngularJSAgent", "name": "AngularJSAgent"},
54
- {"id": "RepoMap", "name": "RepoMap"},
55
- {"id": "gemini-1.5-pro-latest", "name": "gemini-pro"},
56
- {"id": "gemini-1.5-pro", "name": "gemini-1.5-pro"},
57
- {"id": "claude-3-5-sonnet-20240620", "name": "claude-sonnet-3.5"},
58
- {"id": "claude-3-5-sonnet", "name": "claude-sonnet-3.5"},
59
- ]
60
-
61
- # Mapping of requested model names to actual API model identifiers
62
- MODEL_MAPPING = {
63
- "blackboxai": "blackboxai",
64
- "blackboxai-pro": "blackboxai-pro",
65
- "ImageGeneration": "ImageGeneration",
66
- "llama-3.1-8b": "llama-3.1-8b",
67
- "llama-3.1-70b": "llama-3.1-70b",
68
- "llama-3.1-405b": "llama-3.1-405b",
69
- "gpt-4o": "gpt-4o",
70
- "gemini-pro": "gemini-pro",
71
- "gemini-1.5-flash": "gemini-1.5-flash",
72
- "claude-sonnet-3.5": "claude-sonnet-3.5",
73
- "PythonAgent": "PythonAgent",
74
- "JavaAgent": "JavaAgent",
75
- "JavaScriptAgent": "JavaScriptAgent",
76
- "HTMLAgent": "HTMLAgent",
77
- "GoogleCloudAgent": "GoogleCloudAgent",
78
- "AndroidDeveloper": "AndroidDeveloper",
79
- "SwiftDeveloper": "SwiftDeveloper",
80
- "Next.jsAgent": "Next.jsAgent",
81
- "MongoDBAgent": "MongoDBAgent",
82
- "PyTorchAgent": "PyTorchAgent",
83
- "ReactAgent": "ReactAgent",
84
- "XcodeAgent": "XcodeAgent",
85
- "AngularJSAgent": "AngularJSAgent",
86
- "RepoMap": "RepoMap",
87
- # Additional mappings
88
- "gemini-flash": "gemini-1.5-flash",
89
- "claude-3.5-sonnet": "claude-sonnet-3.5",
90
- "flux": "ImageGeneration",
91
- "gemini-1.5-pro-latest": "gemini-pro",
92
- "gemini-1.5-pro": "gemini-1.5-pro",
93
- "claude-3-5-sonnet-20240620": "claude-sonnet-3.5",
94
- "claude-3-5-sonnet": "claude-sonnet-3.5",
95
- }
96
-
97
- # Agent modes configuration
98
- AGENT_MODE = {
99
- 'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
100
- # Add other agent-specific modes if necessary
101
- }
102
-
103
- # Trending agent modes configuration
104
- TRENDING_AGENT_MODE = {
105
- "blackboxai": {},
106
- "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
107
- "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
108
- "llama-3.1-70b": {'mode': True, 'id': "llama-3.1-70b"},
109
- "llama-3.1-405b": {'mode': True, 'id': "llama-3.1-405b"},
110
- "blackboxai-pro": {'mode': True, 'id': "BLACKBOXAI-PRO"},
111
- "PythonAgent": {'mode': True, 'id': "Python Agent"},
112
- "JavaAgent": {'mode': True, 'id': "Java Agent"},
113
- "JavaScriptAgent": {'mode': True, 'id': "JavaScript Agent"},
114
- "HTMLAgent": {'mode': True, 'id': "HTML Agent"},
115
- "GoogleCloudAgent": {'mode': True, 'id': "Google Cloud Agent"},
116
- "AndroidDeveloper": {'mode': True, 'id': "Android Developer"},
117
- "SwiftDeveloper": {'mode': True, 'id': "Swift Developer"},
118
- "Next.jsAgent": {'mode': True, 'id': "Next.js Agent"},
119
- "MongoDBAgent": {'mode': True, 'id': "MongoDB Agent"},
120
- "PyTorchAgent": {'mode': True, 'id': "PyTorch Agent"},
121
- "ReactAgent": {'mode': True, 'id': "React Agent"},
122
- "XcodeAgent": {'mode': True, 'id': "Xcode Agent"},
123
- "AngularJSAgent": {'mode': True, 'id': "AngularJS Agent"},
124
- "RepoMap": {'mode': True, 'id': "repomap"},
125
- }
126
-
127
- # User selected models mapping
128
- USER_SELECTED_MODEL = {
129
- "gpt-4o": "gpt-4o",
130
- "gemini-pro": "gemini-pro",
131
- "claude-sonnet-3.5": "claude-sonnet-3.5",
132
- }
133
-
134
- # Model prefixes for prompt formatting
135
- MODEL_PREFIXES = {
136
- 'gpt-4o': '@GPT-4o',
137
- 'gemini-pro': '@Gemini-PRO',
138
- 'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
139
- 'PythonAgent': '@Python Agent',
140
- 'JavaAgent': '@Java Agent',
141
- 'JavaScriptAgent': '@JavaScript Agent',
142
- 'HTMLAgent': '@HTML Agent',
143
- 'GoogleCloudAgent': '@Google Cloud Agent',
144
- 'AndroidDeveloper': '@Android Developer',
145
- 'SwiftDeveloper': '@Swift Developer',
146
- 'Next.jsAgent': '@Next.js Agent',
147
- 'MongoDBAgent': '@MongoDB Agent',
148
- 'PyTorchAgent': '@PyTorch Agent',
149
- 'ReactAgent': '@React Agent',
150
- 'XcodeAgent': '@Xcode Agent',
151
- 'AngularJSAgent': '@AngularJS Agent',
152
- 'blackboxai-pro': '@BLACKBOXAI-PRO',
153
- 'ImageGeneration': '@Image Generation',
154
- }
155
-
156
- # Model referers for API requests
157
- MODEL_REFERERS = {
158
- "blackboxai": "/?model=blackboxai",
159
- "gpt-4o": "/?model=gpt-4o",
160
- "gemini-pro": "/?model=gemini-pro",
161
- "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
162
- }
163
-
164
- # Model aliases for flexibility in model naming
165
- MODEL_ALIASES = {
166
- "gemini-flash": "gemini-1.5-flash",
167
- "claude-3.5-sonnet": "claude-sonnet-3.5",
168
- "flux": "ImageGeneration",
169
- "gemini-1.5-pro-latest": "gemini-pro",
170
- "gemini-1.5-pro": "gemini-1.5-pro",
171
- "claude-3-5-sonnet-20240620": "claude-sonnet-3.5",
172
- "claude-3-5-sonnet": "claude-sonnet-3.5",
173
- }
 
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,27 +1,20 @@
1
- # api/logger.py
2
-
3
- import logging
4
- import os
5
-
6
- def setup_logger(name):
7
- logger = logging.getLogger(name)
8
- if not logger.handlers:
9
- # Set logging level based on environment
10
- log_level = os.getenv("LOG_LEVEL", "INFO").upper()
11
- logger.setLevel(getattr(logging, log_level, logging.INFO))
12
-
13
- formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
14
-
15
- # Console handler
16
- console_handler = logging.StreamHandler()
17
- console_handler.setFormatter(formatter)
18
- logger.addHandler(console_handler)
19
-
20
- # File Handler - Error Level (Optional)
21
- # Uncomment the following lines to enable error logging to a file
22
- # error_file_handler = logging.FileHandler('error.log')
23
- # error_file_handler.setFormatter(formatter)
24
- # error_file_handler.setLevel(logging.ERROR)
25
- # logger.addHandler(error_file_handler)
26
-
27
- 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
2
- from pydantic import BaseModel
3
-
4
- class Message(BaseModel):
5
- role: str
6
- content: Union[str, List[dict]] # Allowing for both text and complex structures like images
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,71 +1,60 @@
1
- # api/routes.py
2
-
3
- import json
4
- from fastapi import APIRouter, Depends, HTTPException, Request, Response
5
- from fastapi.responses import StreamingResponse
6
- from api.auth import verify_app_secret
7
- from api.config import ALLOWED_MODELS
8
- from api.models import ChatRequest
9
- from api.utils import process_non_streaming_response, process_streaming_response
10
- from api.logger import setup_logger
11
-
12
- logger = setup_logger(__name__)
13
-
14
- router = APIRouter()
15
-
16
- @router.options("/v1/chat/completions")
17
- @router.options("/api/v1/chat/completions")
18
- async def chat_completions_options():
19
- return Response(
20
- status_code=200,
21
- headers={
22
- "Access-Control-Allow-Origin": "*",
23
- "Access-Control-Allow-Methods": "POST, OPTIONS",
24
- "Access-Control-Allow-Headers": "Content-Type, Authorization",
25
- },
26
- )
27
-
28
- @router.get("/v1/models")
29
- @router.get("/api/v1/models")
30
- async def list_models():
31
- return {"object": "list", "data": ALLOWED_MODELS}
32
-
33
- @router.post("/v1/chat/completions")
34
- @router.post("/api/v1/chat/completions")
35
- async def chat_completions(
36
- request: ChatRequest, app_secret: str = Depends(verify_app_secret)
37
- ):
38
- logger.info("Entering chat_completions route")
39
- logger.info(f"Processing chat completion request for model: {request.model}")
40
-
41
- # Validate if the requested model is allowed
42
- allowed_model_ids = [model["id"] for model in ALLOWED_MODELS]
43
- if request.model not in allowed_model_ids:
44
- logger.warning(f"Model {request.model} is not allowed.")
45
- raise HTTPException(
46
- status_code=400,
47
- detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(allowed_model_ids)}",
48
- )
49
-
50
- try:
51
- if request.stream:
52
- logger.info("Streaming response")
53
- return StreamingResponse(process_streaming_response(request), media_type="text/event-stream")
54
- else:
55
- logger.info("Non-streaming response")
56
- return await process_non_streaming_response(request)
57
- except HTTPException as e:
58
- logger.error(f"HTTPException: {e.detail}")
59
- raise e
60
- except Exception as e:
61
- logger.error(f"Unexpected error: {str(e)}")
62
- raise HTTPException(status_code=500, detail="Internal Server Error")
63
-
64
- @router.get("/health")
65
- @router.get("/")
66
- @router.get("/healthz")
67
- @router.get("/ready")
68
- @router.get("/alive")
69
- @router.get("/status")
70
- def health_check(request: Request):
71
- 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,197 +1,480 @@
1
- # api/utils.py
2
-
3
- from datetime import datetime
4
- import json
5
- from typing import Any, Dict, Optional
6
- import uuid
7
- import re
8
-
9
- import httpx
10
- from api.config import (
11
- MODEL_MAPPING,
12
- MODEL_ALIASES,
13
- headers,
14
- AGENT_MODE,
15
- TRENDING_AGENT_MODE,
16
- BASE_URL
17
- )
18
- from fastapi import HTTPException
19
- from api.models import ChatRequest
20
-
21
- from api.logger import setup_logger
22
-
23
- logger = setup_logger(__name__)
24
-
25
- def create_chat_completion_data(
26
- content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
27
- ) -> Dict[str, Any]:
28
- return {
29
- "id": f"chatcmpl-{uuid.uuid4()}",
30
- "object": "chat.completion.chunk",
31
- "created": timestamp,
32
- "model": model,
33
- "choices": [
34
- {
35
- "index": 0,
36
- "delta": {"content": content, "role": "assistant"},
37
- "finish_reason": finish_reason,
38
- }
39
- ],
40
- "usage": None,
41
- }
42
-
43
- def message_to_dict(message):
44
- if isinstance(message.content, str):
45
- return {"role": message.role, "content": message.content}
46
- elif isinstance(message.content, list) and len(message.content) == 2:
47
- return {
48
- "role": message.role,
49
- "content": message.content[0]["text"],
50
- "data": {
51
- "imageBase64": message.content[1]["image_url"]["url"],
52
- "fileText": "",
53
- "title": "snapshot",
54
- },
55
- }
56
- else:
57
- return {"role": message.role, "content": message.content}
58
-
59
- async def process_streaming_response(request: ChatRequest):
60
- # Map the requested model to the actual model used by the API
61
- model = MODEL_MAPPING.get(request.model, MODEL_ALIASES.get(request.model, "blackboxai"))
62
-
63
- logger.info(f"Using model: {model}")
64
-
65
- agent_mode = AGENT_MODE.get(model, {})
66
- trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
67
-
68
- json_data = {
69
- "messages": [message_to_dict(msg) for msg in request.messages],
70
- "previewToken": None,
71
- "userId": None,
72
- "codeModelMode": True,
73
- "agentMode": agent_mode,
74
- "trendingAgentMode": trending_agent_mode,
75
- "isMicMode": False,
76
- "userSystemPrompt": None,
77
- "maxTokens": request.max_tokens,
78
- "playgroundTopP": request.top_p,
79
- "playgroundTemperature": request.temperature,
80
- "isChromeExt": False,
81
- "githubToken": None,
82
- "clickedAnswer2": False,
83
- "clickedAnswer3": False,
84
- "clickedForceWebSearch": False,
85
- "visitFromDelta": False,
86
- "mobileClient": False,
87
- "userSelectedModel": model,
88
- }
89
-
90
- logger.debug(f"Payload for streaming request: {json.dumps(json_data)}")
91
-
92
- async with httpx.AsyncClient() as client:
93
- try:
94
- async with client.stream(
95
- "POST",
96
- f"{BASE_URL}/api/chat",
97
- headers=headers,
98
- json=json_data,
99
- timeout=100,
100
- ) as response:
101
- response.raise_for_status()
102
- logger.info(f"Received response status: {response.status_code}")
103
- async for line in response.aiter_lines():
104
- timestamp = int(datetime.now().timestamp())
105
- if line:
106
- content = line
107
- if content.startswith("$@$v=undefined-rv1$@$"):
108
- content = content[21:]
109
- logger.debug(f"Streaming content: {content}")
110
- yield f"data: {json.dumps(create_chat_completion_data(content, model, timestamp))}\n\n"
111
-
112
- # Indicate the end of the stream
113
- timestamp = int(datetime.now().timestamp())
114
- logger.debug("Ending streaming response")
115
- yield f"data: {json.dumps(create_chat_completion_data('', model, timestamp, 'stop'))}\n\n"
116
- yield "data: [DONE]\n\n"
117
- except httpx.HTTPStatusError as e:
118
- logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
119
- raise HTTPException(status_code=e.response.status_code, detail=str(e))
120
- except httpx.RequestError as e:
121
- logger.error(f"Request error occurred: {str(e)}")
122
- raise HTTPException(status_code=500, detail=str(e))
123
-
124
-
125
- async def process_non_streaming_response(request: ChatRequest):
126
- # Map the requested model to the actual model used by the API
127
- model = MODEL_MAPPING.get(request.model, MODEL_ALIASES.get(request.model, "blackboxai"))
128
-
129
- logger.info(f"Using model: {model}")
130
-
131
- agent_mode = AGENT_MODE.get(model, {})
132
- trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
133
-
134
- json_data = {
135
- "messages": [message_to_dict(msg) for msg in request.messages],
136
- "previewToken": None,
137
- "userId": None,
138
- "codeModelMode": True,
139
- "agentMode": agent_mode,
140
- "trendingAgentMode": trending_agent_mode,
141
- "isMicMode": False,
142
- "userSystemPrompt": None,
143
- "maxTokens": request.max_tokens,
144
- "playgroundTopP": request.top_p,
145
- "playgroundTemperature": request.temperature,
146
- "isChromeExt": False,
147
- "githubToken": None,
148
- "clickedAnswer2": False,
149
- "clickedAnswer3": False,
150
- "clickedForceWebSearch": False,
151
- "visitFromDelta": False,
152
- "mobileClient": False,
153
- "userSelectedModel": model,
154
- }
155
-
156
- logger.debug(f"Payload for non-streaming request: {json.dumps(json_data)}")
157
-
158
- full_response = ""
159
- async with httpx.AsyncClient() as client:
160
- try:
161
- async with client.stream(
162
- method="POST",
163
- url=f"{BASE_URL}/api/chat",
164
- headers=headers,
165
- json=json_data,
166
- ) as response:
167
- response.raise_for_status()
168
- logger.info(f"Received response status: {response.status_code}")
169
- async for chunk in response.aiter_text():
170
- full_response += chunk
171
- logger.debug(f"Received chunk: {chunk}")
172
- except httpx.HTTPStatusError as e:
173
- logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
174
- raise HTTPException(status_code=e.response.status_code, detail=str(e))
175
- except httpx.RequestError as e:
176
- logger.error(f"Request error occurred: {str(e)}")
177
- raise HTTPException(status_code=500, detail=str(e))
178
-
179
- if full_response.startswith("$@$v=undefined-rv1$@$"):
180
- full_response = full_response[21:]
181
-
182
- logger.debug(f"Full non-streaming response: {full_response}")
183
-
184
- return {
185
- "id": f"chatcmpl-{uuid.uuid4()}",
186
- "object": "chat.completion",
187
- "created": int(datetime.now().timestamp()),
188
- "model": model,
189
- "choices": [
190
- {
191
- "index": 0,
192
- "message": {"role": "assistant", "content": full_response},
193
- "finish_reason": "stop",
194
- }
195
- ],
196
- "usage": None,
197
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ }
main.py CHANGED
@@ -1,5 +1,5 @@
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)
 
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
- 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
 
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