Niansuh commited on
Commit
a5d5b47
·
verified ·
1 Parent(s): b786aea

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +607 -5
main.py CHANGED
@@ -1,5 +1,607 @@
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 os
2
+ import re
3
+ import random
4
+ import string
5
+ import uuid
6
+ import json
7
+ import logging
8
+ import asyncio
9
+ import time
10
+ from collections import defaultdict
11
+ from typing import List, Dict, Any, Optional, AsyncGenerator, Union
12
+ from datetime import datetime
13
+
14
+ from aiohttp import ClientSession, ClientTimeout, ClientError
15
+ from fastapi import FastAPI, HTTPException, Request, Depends, Header
16
+ from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
17
+ from pydantic import BaseModel
18
+
19
+ # Configure logging
20
+ logging.basicConfig(
21
+ level=logging.INFO,
22
+ format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
23
+ handlers=[logging.StreamHandler()]
24
+ )
25
+ logger = logging.getLogger(__name__)
26
+
27
+ # Load environment variables
28
+ API_KEYS = os.getenv('API_KEYS', '').split(',') # Comma-separated API keys
29
+ RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60')) # Requests per minute
30
+ AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '') # Comma-separated available models
31
+
32
+ if not API_KEYS or API_KEYS == ['']:
33
+ logger.error("No API keys found. Please set the API_KEYS environment variable.")
34
+ raise Exception("API_KEYS environment variable not set.")
35
+
36
+ # Process available models
37
+ if AVAILABLE_MODELS:
38
+ AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
39
+ else:
40
+ AVAILABLE_MODELS = [] # If empty, all models are available
41
+
42
+ # Simple in-memory rate limiter based solely on IP addresses
43
+ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
44
+
45
+ # Define cleanup interval and window
46
+ CLEANUP_INTERVAL = 60 # seconds
47
+ RATE_LIMIT_WINDOW = 60 # seconds
48
+
49
+ async def cleanup_rate_limit_stores():
50
+ """
51
+ Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
52
+ """
53
+ while True:
54
+ current_time = time.time()
55
+ ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
56
+ for ip in ips_to_delete:
57
+ del rate_limit_store[ip]
58
+ logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
59
+ await asyncio.sleep(CLEANUP_INTERVAL)
60
+
61
+ async def rate_limiter_per_ip(request: Request):
62
+ """
63
+ Rate limiter that enforces a limit based on the client's IP address.
64
+ """
65
+ client_ip = request.client.host
66
+ current_time = time.time()
67
+
68
+ # Initialize or update the count and timestamp
69
+ if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
70
+ rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
71
+ else:
72
+ if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
73
+ logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
74
+ raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
75
+ rate_limit_store[client_ip]["count"] += 1
76
+
77
+ async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
78
+ """
79
+ Dependency to extract and validate the API key from the Authorization header.
80
+ """
81
+ client_ip = request.client.host
82
+ if authorization is None or not authorization.startswith('Bearer '):
83
+ logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
84
+ raise HTTPException(status_code=401, detail='Invalid authorization header format')
85
+ api_key = authorization[7:]
86
+ if api_key not in API_KEYS:
87
+ logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
88
+ raise HTTPException(status_code=401, detail='Invalid API key')
89
+ return api_key
90
+
91
+ # Custom exception for model not working
92
+ class ModelNotWorkingException(Exception):
93
+ def __init__(self, model: str):
94
+ self.model = model
95
+ self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
96
+ super().__init__(self.message)
97
+
98
+ # Mock implementations for ImageResponse and to_data_uri
99
+ class ImageResponse:
100
+ def __init__(self, url: str, alt: str):
101
+ self.url = url
102
+ self.alt = alt
103
+
104
+ def to_data_uri(image: Any) -> str:
105
+ return "data:image/png;base64,..." # Replace with actual base64 data
106
+
107
+ class Blackbox:
108
+ url = "https://www.blackbox.ai"
109
+ api_endpoint = "https://www.blackbox.ai/api/chat"
110
+ working = True
111
+ supports_stream = True
112
+ supports_system_message = True
113
+ supports_message_history = True
114
+
115
+ default_model = 'blackboxai'
116
+ image_models = ['ImageGeneration']
117
+ models = [
118
+ default_model,
119
+ 'blackboxai-pro',
120
+ "llama-3.1-8b",
121
+ 'llama-3.1-70b',
122
+ 'llama-3.1-405b',
123
+ 'gpt-4o',
124
+ 'gemini-pro',
125
+ 'gemini-1.5-flash',
126
+ 'claude-sonnet-3.5',
127
+ 'PythonAgent',
128
+ 'JavaAgent',
129
+ 'JavaScriptAgent',
130
+ 'HTMLAgent',
131
+ 'GoogleCloudAgent',
132
+ 'AndroidDeveloper',
133
+ 'SwiftDeveloper',
134
+ 'Next.jsAgent',
135
+ 'MongoDBAgent',
136
+ 'PyTorchAgent',
137
+ 'ReactAgent',
138
+ 'XcodeAgent',
139
+ 'AngularJSAgent',
140
+ *image_models,
141
+ 'Niansuh',
142
+ ]
143
+
144
+ # Filter models based on AVAILABLE_MODELS
145
+ if AVAILABLE_MODELS:
146
+ models = [model for model in models if model in AVAILABLE_MODELS]
147
+
148
+ agentMode = {
149
+ 'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
150
+ 'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
151
+ }
152
+ trendingAgentMode = {
153
+ "blackboxai": {},
154
+ "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
155
+ "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
156
+ 'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
157
+ 'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
158
+ 'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
159
+ 'PythonAgent': {'mode': True, 'id': "Python Agent"},
160
+ 'JavaAgent': {'mode': True, 'id': "Java Agent"},
161
+ 'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
162
+ 'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
163
+ 'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
164
+ 'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
165
+ 'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
166
+ 'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
167
+ 'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
168
+ 'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
169
+ 'ReactAgent': {'mode': True, 'id': "React Agent"},
170
+ 'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
171
+ 'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
172
+ }
173
+
174
+ userSelectedModel = {
175
+ "gpt-4o": "gpt-4o",
176
+ "gemini-pro": "gemini-pro",
177
+ 'claude-sonnet-3.5': "claude-sonnet-3.5",
178
+ }
179
+
180
+ model_prefixes = {
181
+ 'gpt-4o': '@GPT-4o',
182
+ 'gemini-pro': '@Gemini-PRO',
183
+ 'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
184
+ 'PythonAgent': '@Python Agent',
185
+ 'JavaAgent': '@Java Agent',
186
+ 'JavaScriptAgent': '@JavaScript Agent',
187
+ 'HTMLAgent': '@HTML Agent',
188
+ 'GoogleCloudAgent': '@Google Cloud Agent',
189
+ 'AndroidDeveloper': '@Android Developer',
190
+ 'SwiftDeveloper': '@Swift Developer',
191
+ 'Next.jsAgent': '@Next.js Agent',
192
+ 'MongoDBAgent': '@MongoDB Agent',
193
+ 'PyTorchAgent': '@PyTorch Agent',
194
+ 'ReactAgent': '@React Agent',
195
+ 'XcodeAgent': '@Xcode Agent',
196
+ 'AngularJSAgent': '@AngularJS Agent',
197
+ 'blackboxai-pro': '@BLACKBOXAI-PRO',
198
+ 'ImageGeneration': '@Image Generation',
199
+ 'Niansuh': '@Niansuh',
200
+ }
201
+
202
+ model_referers = {
203
+ "blackboxai": f"{url}/?model=blackboxai",
204
+ "gpt-4o": f"{url}/?model=gpt-4o",
205
+ "gemini-pro": f"{url}/?model=gemini-pro",
206
+ "claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5"
207
+ }
208
+
209
+ model_aliases = {
210
+ "gemini-flash": "gemini-1.5-flash",
211
+ "claude-3.5-sonnet": "claude-sonnet-3.5",
212
+ "flux": "ImageGeneration",
213
+ "niansuh": "Niansuh",
214
+ }
215
+
216
+ @classmethod
217
+ def get_model(cls, model: str) -> Optional[str]:
218
+ if model in cls.models:
219
+ return model
220
+ elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
221
+ return model
222
+ elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
223
+ return cls.model_aliases[model]
224
+ else:
225
+ return cls.default_model if cls.default_model in cls.models else None
226
+
227
+ @classmethod
228
+ async def create_async_generator(
229
+ cls,
230
+ model: str,
231
+ messages: List[Dict[str, str]],
232
+ proxy: Optional[str] = None,
233
+ image: Any = None,
234
+ image_name: Optional[str] = None,
235
+ webSearchMode: bool = False,
236
+ **kwargs
237
+ ) -> AsyncGenerator[Any, None]:
238
+ model = cls.get_model(model)
239
+ if model is None:
240
+ logger.error(f"Model {model} is not available.")
241
+ raise ModelNotWorkingException(model)
242
+
243
+ logger.info(f"Selected model: {model}")
244
+
245
+ if not cls.working or model not in cls.models:
246
+ logger.error(f"Model {model} is not working or not supported.")
247
+ raise ModelNotWorkingException(model)
248
+
249
+ headers = {
250
+ "accept": "*/*",
251
+ "accept-language": "en-US,en;q=0.9",
252
+ "cache-control": "no-cache",
253
+ "content-type": "application/json",
254
+ "origin": cls.url,
255
+ "pragma": "no-cache",
256
+ "priority": "u=1, i",
257
+ "referer": cls.model_referers.get(model, cls.url),
258
+ "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
259
+ "sec-ch-ua-mobile": "?0",
260
+ "sec-ch-ua-platform": '"Linux"',
261
+ "sec-fetch-dest": "empty",
262
+ "sec-fetch-mode": "cors",
263
+ "sec-fetch-site": "same-origin",
264
+ "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
265
+ }
266
+
267
+ if model in cls.model_prefixes:
268
+ prefix = cls.model_prefixes[model]
269
+ if not messages[0]['content'].startswith(prefix):
270
+ logger.debug(f"Adding prefix '{prefix}' to the first message.")
271
+ messages[0]['content'] = f"{prefix} {messages[0]['content']}"
272
+
273
+ random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
274
+ messages[-1]['id'] = random_id
275
+ messages[-1]['role'] = 'user'
276
+
277
+ # Don't log the full message content for privacy
278
+ logger.debug(f"Generated message ID: {random_id} for model: {model}")
279
+
280
+ if image is not None:
281
+ messages[-1]['data'] = {
282
+ 'fileText': '',
283
+ 'imageBase64': to_data_uri(image),
284
+ 'title': image_name
285
+ }
286
+ messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
287
+ logger.debug("Image data added to the message.")
288
+
289
+ data = {
290
+ "messages": messages,
291
+ "id": random_id,
292
+ "previewToken": None,
293
+ "userId": None,
294
+ "codeModelMode": True,
295
+ "agentMode": {},
296
+ "trendingAgentMode": {},
297
+ "isMicMode": False,
298
+ "userSystemPrompt": None,
299
+ "maxTokens": 99999999,
300
+ "playgroundTopP": 0.9,
301
+ "playgroundTemperature": 0.5,
302
+ "isChromeExt": False,
303
+ "githubToken": None,
304
+ "clickedAnswer2": False,
305
+ "clickedAnswer3": False,
306
+ "clickedForceWebSearch": False,
307
+ "visitFromDelta": False,
308
+ "mobileClient": False,
309
+ "userSelectedModel": None,
310
+ "webSearchMode": webSearchMode,
311
+ "validated": "00f37b34-a166-4efb-bce5-1312d87f2f94"
312
+ }
313
+
314
+ if model in cls.agentMode:
315
+ data["agentMode"] = cls.agentMode[model]
316
+ elif model in cls.trendingAgentMode:
317
+ data["trendingAgentMode"] = cls.trendingAgentMode[model]
318
+ elif model in cls.userSelectedModel:
319
+ data["userSelectedModel"] = cls.userSelectedModel[model]
320
+ logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
321
+
322
+ timeout = ClientTimeout(total=60) # Set an appropriate timeout
323
+ retry_attempts = 10 # Set the number of retry attempts
324
+
325
+ for attempt in range(retry_attempts):
326
+ try:
327
+ async with ClientSession(headers=headers, timeout=timeout) as session:
328
+ async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
329
+ response.raise_for_status()
330
+ logger.info(f"Received response with status {response.status}")
331
+ if model == 'ImageGeneration':
332
+ response_text = await response.text()
333
+ url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
334
+ if url_match:
335
+ image_url = url_match.group(0)
336
+ logger.info(f"Image URL found.")
337
+ yield ImageResponse(image_url, alt=messages[-1]['content'])
338
+ else:
339
+ logger.error("Image URL not found in the response.")
340
+ raise Exception("Image URL not found in the response")
341
+ else:
342
+ full_response = ""
343
+ search_results_json = ""
344
+ try:
345
+ async for chunk, _ in response.content.iter_chunks():
346
+ if chunk:
347
+ decoded_chunk = chunk.decode(errors='ignore')
348
+ decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
349
+ if decoded_chunk.strip():
350
+ if '$~~~$' in decoded_chunk:
351
+ search_results_json += decoded_chunk
352
+ else:
353
+ full_response += decoded_chunk
354
+ yield decoded_chunk
355
+ logger.info("Finished streaming response chunks.")
356
+ except Exception as e:
357
+ logger.exception("Error while iterating over response chunks.")
358
+ raise e
359
+ if data["webSearchMode"] and search_results_json:
360
+ match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
361
+ if match:
362
+ try:
363
+ search_results = json.loads(match.group(1))
364
+ formatted_results = "\n\n**Sources:**\n"
365
+ for i, result in enumerate(search_results[:5], 1):
366
+ formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
367
+ logger.info("Formatted search results.")
368
+ yield formatted_results
369
+ except json.JSONDecodeError as je:
370
+ logger.error("Failed to parse search results JSON.")
371
+ raise je
372
+ break # Exit the retry loop if successful
373
+ except ClientError as ce:
374
+ logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
375
+ if attempt == retry_attempts - 1:
376
+ raise HTTPException(status_code=502, detail="Error communicating with the external API.")
377
+ except asyncio.TimeoutError:
378
+ logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
379
+ if attempt == retry_attempts - 1:
380
+ raise HTTPException(status_code=504, detail="External API request timed out.")
381
+ except Exception as e:
382
+ logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
383
+ if attempt == retry_attempts - 1:
384
+ raise HTTPException(status_code=500, detail=str(e))
385
+
386
+ # FastAPI app setup
387
+ app = FastAPI()
388
+
389
+ # Add the cleanup task when the app starts
390
+ @app.on_event("startup")
391
+ async def startup_event():
392
+ asyncio.create_task(cleanup_rate_limit_stores())
393
+ logger.info("Started rate limit store cleanup task.")
394
+
395
+ # Middleware to enhance security and enforce Content-Type for specific endpoints
396
+ @app.middleware("http")
397
+ async def security_middleware(request: Request, call_next):
398
+ client_ip = request.client.host
399
+ # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
400
+ if request.method == "POST" and request.url.path == "/v1/chat/completions":
401
+ content_type = request.headers.get("Content-Type")
402
+ if content_type != "application/json":
403
+ logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
404
+ return JSONResponse(
405
+ status_code=400,
406
+ content={
407
+ "error": {
408
+ "message": "Content-Type must be application/json",
409
+ "type": "invalid_request_error",
410
+ "param": None,
411
+ "code": None
412
+ }
413
+ },
414
+ )
415
+ response = await call_next(request)
416
+ return response
417
+
418
+ # Request Models
419
+ class Message(BaseModel):
420
+ role: str
421
+ content: str
422
+
423
+ class ChatRequest(BaseModel):
424
+ model: str
425
+ messages: List[Message]
426
+ temperature: Optional[float] = 1.0
427
+ top_p: Optional[float] = 1.0
428
+ n: Optional[int] = 1
429
+ stream: Optional[bool] = False
430
+ stop: Optional[Union[str, List[str]]] = None
431
+ max_tokens: Optional[int] = None
432
+ presence_penalty: Optional[float] = 0.0
433
+ frequency_penalty: Optional[float] = 0.0
434
+ logit_bias: Optional[Dict[str, float]] = None
435
+ user: Optional[str] = None
436
+ webSearchMode: Optional[bool] = False # Custom parameter
437
+
438
+ def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
439
+ return {
440
+ "id": f"chatcmpl-{uuid.uuid4()}",
441
+ "object": "chat.completion.chunk",
442
+ "created": int(datetime.now().timestamp()),
443
+ "model": model,
444
+ "choices": [
445
+ {
446
+ "index": 0,
447
+ "delta": {"content": content, "role": "assistant"},
448
+ "finish_reason": finish_reason,
449
+ }
450
+ ],
451
+ "usage": None,
452
+ }
453
+
454
+ @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
455
+ async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
456
+ client_ip = req.client.host
457
+ # Redact user messages only for logging purposes
458
+ redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
459
+
460
+ logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
461
+
462
+ try:
463
+ # Validate that the requested model is available
464
+ if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
465
+ logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
466
+ raise HTTPException(status_code=400, detail="Requested model is not available.")
467
+
468
+ # Process the request with actual message content, but don't log it
469
+ async_generator = Blackbox.create_async_generator(
470
+ model=request.model,
471
+ messages=[{"role": msg.role, "content": msg.content} for msg in request.messages], # Actual message content used here
472
+ image=None,
473
+ image_name=None,
474
+ webSearchMode=request.webSearchMode
475
+ )
476
+
477
+ if request.stream:
478
+ async def generate():
479
+ try:
480
+ async for chunk in async_generator:
481
+ if isinstance(chunk, ImageResponse):
482
+ image_markdown = f"![image]({chunk.url})"
483
+ response_chunk = create_response(image_markdown, request.model)
484
+ else:
485
+ response_chunk = create_response(chunk, request.model)
486
+
487
+ yield f"data: {json.dumps(response_chunk)}\n\n"
488
+
489
+ yield "data: [DONE]\n\n"
490
+ except HTTPException as he:
491
+ error_response = {"error": he.detail}
492
+ yield f"data: {json.dumps(error_response)}\n\n"
493
+ except Exception as e:
494
+ logger.exception(f"Error during streaming response generation from IP: {client_ip}.")
495
+ error_response = {"error": str(e)}
496
+ yield f"data: {json.dumps(error_response)}\n\n"
497
+
498
+ return StreamingResponse(generate(), media_type="text/event-stream")
499
+ else:
500
+ response_content = ""
501
+ async for chunk in async_generator:
502
+ if isinstance(chunk, ImageResponse):
503
+ response_content += f"![image]({chunk.url})\n"
504
+ else:
505
+ response_content += chunk
506
+
507
+ logger.info(f"Completed non-streaming response generation for API key: {api_key} | IP: {client_ip}")
508
+ return {
509
+ "id": f"chatcmpl-{uuid.uuid4()}",
510
+ "object": "chat.completion",
511
+ "created": int(datetime.now().timestamp()),
512
+ "model": request.model,
513
+ "choices": [
514
+ {
515
+ "message": {
516
+ "role": "assistant",
517
+ "content": response_content
518
+ },
519
+ "finish_reason": "stop",
520
+ "index": 0
521
+ }
522
+ ],
523
+ "usage": {
524
+ "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
525
+ "completion_tokens": len(response_content.split()),
526
+ "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
527
+ },
528
+ }
529
+ except ModelNotWorkingException as e:
530
+ logger.warning(f"Model not working: {e} | IP: {client_ip}")
531
+ raise HTTPException(status_code=503, detail=str(e))
532
+ except HTTPException as he:
533
+ logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
534
+ raise he
535
+ except Exception as e:
536
+ logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
537
+ raise HTTPException(status_code=500, detail=str(e))
538
+
539
+ # Re-added endpoints without API key authentication
540
+
541
+ # Endpoint: POST /v1/tokenizer
542
+ class TokenizerRequest(BaseModel):
543
+ text: str
544
+
545
+ @app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
546
+ async def tokenizer(request: TokenizerRequest, req: Request):
547
+ client_ip = req.client.host
548
+ text = request.text
549
+ token_count = len(text.split())
550
+ logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
551
+ return {"text": text, "tokens": token_count}
552
+
553
+ # Endpoint: GET /v1/models
554
+ @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
555
+ async def get_models(req: Request):
556
+ client_ip = req.client.host
557
+ logger.info(f"Fetching available models from IP: {client_ip}")
558
+ return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
559
+
560
+ # Endpoint: GET /v1/models/{model}/status
561
+ @app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
562
+ async def model_status(model: str, req: Request):
563
+ client_ip = req.client.host
564
+ logger.info(f"Model status requested for '{model}' from IP: {client_ip}")
565
+ if model in Blackbox.models:
566
+ return {"model": model, "status": "available"}
567
+ elif model in Blackbox.model_aliases and Blackbox.model_aliases[model] in Blackbox.models:
568
+ actual_model = Blackbox.model_aliases[model]
569
+ return {"model": actual_model, "status": "available via alias"}
570
+ else:
571
+ logger.warning(f"Model not found: {model} from IP: {client_ip}")
572
+ raise HTTPException(status_code=404, detail="Model not found")
573
+
574
+ # Endpoint: GET /v1/health
575
+ @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
576
+ async def health_check(req: Request):
577
+ client_ip = req.client.host
578
+ logger.info(f"Health check requested from IP: {client_ip}")
579
+ return {"status": "ok"}
580
+
581
+ # Endpoint: GET /v1/chat/completions (GET method)
582
+ @app.get("/v1/chat/completions")
583
+ async def chat_completions_get(req: Request):
584
+ client_ip = req.client.host
585
+ logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
586
+ return RedirectResponse(url='about:blank')
587
+
588
+ # Custom exception handler to match OpenAI's error format
589
+ @app.exception_handler(HTTPException)
590
+ async def http_exception_handler(request: Request, exc: HTTPException):
591
+ client_ip = request.client.host
592
+ logger.error(f"HTTPException: {exc.detail} | Path: {request.url.path} | IP: {client_ip}")
593
+ return JSONResponse(
594
+ status_code=exc.status_code,
595
+ content={
596
+ "error": {
597
+ "message": exc.detail,
598
+ "type": "invalid_request_error",
599
+ "param": None,
600
+ "code": None
601
+ }
602
+ },
603
+ )
604
+
605
+ if __name__ == "__main__":
606
+ import uvicorn
607
+ uvicorn.run(app, host="0.0.0.0", port=8000)