|
|
import asyncio |
|
|
import json |
|
|
import random |
|
|
import re |
|
|
import string |
|
|
import time |
|
|
import uuid |
|
|
from datetime import datetime, timezone, timedelta |
|
|
from typing import Any, Dict, List, Optional |
|
|
|
|
|
import boto3 |
|
|
import httpx |
|
|
import tiktoken |
|
|
import platform |
|
|
import hashlib |
|
|
from fastapi import HTTPException |
|
|
|
|
|
from api.config import ( |
|
|
MODEL_MAPPING, |
|
|
get_headers_api_chat, |
|
|
get_headers_chat, |
|
|
BASE_URL, |
|
|
AGENT_MODE, |
|
|
TRENDING_AGENT_MODE, |
|
|
MODEL_PREFIXES |
|
|
) |
|
|
from api.logger import setup_logger |
|
|
from api.models import ChatRequest |
|
|
from api.validate import getHid |
|
|
|
|
|
logger = setup_logger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
R2_ACCESS_KEY_ID = "df9c9eb87e850a8eb27afd3968077b42" |
|
|
R2_SECRET_ACCESS_KEY = "14b08b0855263bb63d2618da3a6537e1b0446d89d51da03a568620b1e5342ea8" |
|
|
R2_ENDPOINT_URL = "https://f2f92ac53fae792c4155f6e93a514989.r2.cloudflarestorage.com" |
|
|
R2_BUCKET_NAME = "snapzion" |
|
|
R2_REPLACED_URLS_KEY = "snapzion.txt" |
|
|
|
|
|
s3 = boto3.client( |
|
|
"s3", |
|
|
endpoint_url=R2_ENDPOINT_URL, |
|
|
aws_access_key_id=R2_ACCESS_KEY_ID, |
|
|
aws_secret_access_key=R2_SECRET_ACCESS_KEY, |
|
|
) |
|
|
|
|
|
BLOCKED_MESSAGE = ( |
|
|
"Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai " |
|
|
"and for API requests replace https://www.blackbox.ai with https://api.blackbox.ai" |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_random_name_email_customer(): |
|
|
first_names = ["Aliace", "B21ob", "Car232ol", "Daavid", "Evewwlyn", "Fraank", "Grssace", "Hefctor", "Ivgy", "Jackdie"] |
|
|
last_names = ["Smilth", "Johnkson", "Dajvis", "Mihller", "Thomgpson", "Garwcia", "Broawn", "Wilfson", "Maartin", "Clarak"] |
|
|
|
|
|
random_name = f"{random.choice(first_names)} {random.choice(last_names)}" |
|
|
email_username = ''.join(random.choices(string.ascii_lowercase + string.digits, k=8)) |
|
|
random_email = f"[email protected]" |
|
|
suffix_length = len("Rldf7IKdNhdhiw") |
|
|
suffix_chars = string.ascii_letters + string.digits |
|
|
random_suffix = ''.join(random.choice(suffix_chars) for _ in range(suffix_length)) |
|
|
random_customer_id = f"cus_{random_suffix}" |
|
|
|
|
|
return random_name, random_email, random_customer_id |
|
|
|
|
|
def generate_session(email: str, id_length: int = 21, days_ahead: int = 365) -> dict: |
|
|
numeric_id = ''.join(random.choice('0123456789') for _ in range(id_length)) |
|
|
future_date = datetime.now(timezone.utc) + timedelta(days=days_ahead) |
|
|
expiry = future_date.isoformat(timespec='milliseconds').replace('+00:00', 'Z') |
|
|
|
|
|
chars = string.ascii_letters + string.digits + "-" |
|
|
random_img_id = ''.join(random.choice(chars) for _ in range(48)) |
|
|
image_url = f"https://lh3.googleusercontent.com/a/ACg8oc{random_img_id}=s96-c" |
|
|
|
|
|
return { |
|
|
"user": { |
|
|
"name": "SNAPZION", |
|
|
"email": email, |
|
|
"image": image_url, |
|
|
"id": numeric_id |
|
|
}, |
|
|
"expires": expiry, |
|
|
"isNewUser": False |
|
|
} |
|
|
|
|
|
def generate_session_data() -> dict: |
|
|
_, email, _ = get_random_name_email_customer() |
|
|
session_data = generate_session(email) |
|
|
logger.info(f"Using generated session with email {email}") |
|
|
return session_data |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def generate_system_fingerprint() -> str: |
|
|
raw = f"{platform.node()}-{time.time()}-{uuid.uuid4()}" |
|
|
short_hash = hashlib.md5(raw.encode()).hexdigest()[:12] |
|
|
return f"fp_{short_hash}" |
|
|
|
|
|
def get_last_user_prompt(messages: List[Any]) -> str: |
|
|
for msg in reversed(messages): |
|
|
if msg.role == "user": |
|
|
if isinstance(msg.content, str): |
|
|
return msg.content.strip() |
|
|
elif isinstance(msg.content, list): |
|
|
for item in msg.content: |
|
|
if item.get("type") == "text": |
|
|
return item.get("text", "").strip() |
|
|
return "" |
|
|
|
|
|
def upload_replaced_urls_to_r2(urls: List[str], alt_text: str = "") -> None: |
|
|
if not urls: |
|
|
logger.info("No replaced or final URLs to store.") |
|
|
return |
|
|
existing = "" |
|
|
try: |
|
|
resp = s3.get_object(Bucket=R2_BUCKET_NAME, Key=R2_REPLACED_URLS_KEY) |
|
|
existing = resp["Body"].read().decode() |
|
|
except s3.exceptions.NoSuchKey: |
|
|
pass |
|
|
except Exception as e: |
|
|
logger.error(f"Error reading {R2_REPLACED_URLS_KEY}: {e}") |
|
|
|
|
|
markdown = "\n".join(f"" for u in urls) |
|
|
content = f"{existing}\n{markdown}" if existing.strip() else markdown |
|
|
|
|
|
try: |
|
|
s3.put_object( |
|
|
Bucket=R2_BUCKET_NAME, |
|
|
Key=R2_REPLACED_URLS_KEY, |
|
|
Body=content.encode(), |
|
|
ContentType="text/plain", |
|
|
) |
|
|
logger.info(f"Appended {len(urls)} URLs to {R2_REPLACED_URLS_KEY}.") |
|
|
except Exception as e: |
|
|
logger.error(f"Error writing {R2_REPLACED_URLS_KEY}: {e}") |
|
|
|
|
|
def calculate_tokens(text: str, model: str) -> int: |
|
|
try: |
|
|
enc = tiktoken.encoding_for_model(model) |
|
|
return len(enc.encode(text)) |
|
|
except Exception: |
|
|
return len(text.split()) |
|
|
|
|
|
def create_chat_completion_data( |
|
|
content: str, |
|
|
model: str, |
|
|
timestamp: int, |
|
|
request_id: str, |
|
|
system_fingerprint: str, |
|
|
prompt_tokens: int = 0, |
|
|
completion_tokens: int = 0, |
|
|
finish_reason: Optional[str] = None, |
|
|
) -> Dict[str, Any]: |
|
|
usage = None |
|
|
if finish_reason == "stop": |
|
|
usage = { |
|
|
"prompt_tokens": prompt_tokens, |
|
|
"completion_tokens": completion_tokens, |
|
|
"total_tokens": prompt_tokens + completion_tokens, |
|
|
} |
|
|
return { |
|
|
"id": request_id, |
|
|
"object": "chat.completion.chunk", |
|
|
"created": timestamp, |
|
|
"model": model, |
|
|
"system_fingerprint": system_fingerprint, |
|
|
"choices": [{"index": 0, "delta": {"content": content, "role": "assistant"}, "finish_reason": finish_reason}], |
|
|
"usage": usage, |
|
|
} |
|
|
|
|
|
def message_to_dict(message, model_prefix: Optional[str] = None): |
|
|
""" |
|
|
Convert a ChatRequest message to a dict for the request payload. |
|
|
Supports up to three images with type-based structure and sends multiple formats. |
|
|
Prepends model_prefix to text content if specified. |
|
|
""" |
|
|
content = "" |
|
|
images_data = [] |
|
|
image_urls = [] |
|
|
|
|
|
if isinstance(message.content, list): |
|
|
for item in message.content: |
|
|
if item.get("type") == "text": |
|
|
content = item.get("text", "").strip() |
|
|
elif item.get("type") == "image_url" and len(images_data) < 3: |
|
|
url = item["image_url"].get("url", "") |
|
|
if url: |
|
|
path = f"MultipleFiles/{uuid.uuid4().hex}.jpg" |
|
|
images_data.append({"filePath": path, "contents": url}) |
|
|
image_urls.append({"image_url": {"url": url}}) |
|
|
elif isinstance(message.content, str): |
|
|
content = message.content.strip() |
|
|
|
|
|
if model_prefix and content: |
|
|
content = f"{model_prefix} {content}" |
|
|
|
|
|
base = {"role": message.role, "content": content} |
|
|
if images_data: |
|
|
base["data"] = { |
|
|
"imageBase64": images_data[0]["contents"], |
|
|
"fileText": "", |
|
|
"title": "snapshot", |
|
|
"imagesData": images_data |
|
|
} |
|
|
|
|
|
for extra in image_urls[1:]: |
|
|
base.setdefault("content", []).append(extra) |
|
|
|
|
|
return base |
|
|
|
|
|
def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str: |
|
|
if model_prefix and content.startswith(model_prefix): |
|
|
return content[len(model_prefix):].strip() |
|
|
return content |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def process_streaming_response(request: ChatRequest): |
|
|
system_fp = generate_system_fingerprint() |
|
|
request_id = f"chatcmpl-{uuid.uuid4()}" |
|
|
logger.info(f"Processing (stream) {request_id} - Model: {request.model}") |
|
|
|
|
|
agent_mode = AGENT_MODE.get(request.model, {}) |
|
|
trending_mode = TRENDING_AGENT_MODE.get(request.model, {}) |
|
|
prefix = MODEL_PREFIXES.get(request.model, "") |
|
|
|
|
|
headers_api = get_headers_api_chat(BASE_URL) |
|
|
if request.model == "o1-preview": |
|
|
await asyncio.sleep(random.randint(1, 60)) |
|
|
|
|
|
h = await getHid() |
|
|
if not h: |
|
|
raise HTTPException(status_code=500, detail="Missing h-value.") |
|
|
|
|
|
msgs = [message_to_dict(m, prefix) for m in request.messages] |
|
|
|
|
|
json_data = { |
|
|
"agentMode": agent_mode, |
|
|
"clickedAnswer2": False, |
|
|
"clickedAnswer3": False, |
|
|
"reasoningMode": False, |
|
|
"clickedForceWebSearch": False, |
|
|
"codeInterpreterMode": False, |
|
|
"codeModelMode": True, |
|
|
"githubToken": "", |
|
|
"deepSearchMode": False, |
|
|
"domains": None, |
|
|
"id": request_id, |
|
|
"imageGenerationMode": False, |
|
|
"isChromeExt": False, |
|
|
"isMicMode": False, |
|
|
"isPremium": True, |
|
|
"isMemoryEnabled": False, |
|
|
"maxTokens": request.max_tokens, |
|
|
"messages": msgs, |
|
|
"mobileClient": False, |
|
|
"playgroundTemperature": request.temperature, |
|
|
"playgroundTopP": request.top_p, |
|
|
"previewToken": None, |
|
|
"trendingAgentMode": trending_mode, |
|
|
"userId": None, |
|
|
"userSelectedModel": MODEL_MAPPING.get(request.model, request.model), |
|
|
"userSystemPrompt": None, |
|
|
"validated": h, |
|
|
"visitFromDelta": False, |
|
|
"webSearchModePrompt": False, |
|
|
"vscodeClient": False, |
|
|
"designerMode": False, |
|
|
"workspaceId": "", |
|
|
"beastMode": False, |
|
|
"customProfile": {"name": "", "occupation": "", "traits": [], "additionalInfo": "", "enableNewChats": False}, |
|
|
"webSearchModeOption": {"autoMode": False, "webMode": False, "offlineMode": True}, |
|
|
"session": generate_session_data(), |
|
|
} |
|
|
|
|
|
prompt_tokens = sum( |
|
|
calculate_tokens(m.get("content", ""), request.model) + |
|
|
sum(calculate_tokens(img["contents"], request.model) for img in m.get("data", {}).get("imagesData", [])) |
|
|
for m in msgs |
|
|
) |
|
|
|
|
|
completion_tokens = 0 |
|
|
final_links: List[str] = [] |
|
|
|
|
|
async with httpx.AsyncClient() as client: |
|
|
try: |
|
|
async with client.stream("POST", f"{BASE_URL}/api/chat", headers=headers_api, json=json_data, timeout=100) as resp: |
|
|
resp.raise_for_status() |
|
|
async for chunk in resp.aiter_text(): |
|
|
if not chunk: |
|
|
continue |
|
|
if chunk.startswith("$@$v=undefined-rv1$@$"): |
|
|
chunk = chunk[21:] |
|
|
if BLOCKED_MESSAGE in chunk: |
|
|
chunk = chunk.replace(BLOCKED_MESSAGE, "").strip() |
|
|
if not chunk: |
|
|
continue |
|
|
chunk = chunk.replace("https://storage.googleapis.com", "https://cdn.snapzion.com") |
|
|
links = re.findall(r"(https://cdn\.snapzion\.com[^\s\)]+)", chunk) |
|
|
final_links.extend(links) |
|
|
clean = strip_model_prefix(chunk, prefix) |
|
|
completion_tokens += calculate_tokens(clean, request.model) |
|
|
ts = int(datetime.now().timestamp()) |
|
|
yield "data: " + json.dumps( |
|
|
create_chat_completion_data(clean, request.model, ts, request_id, system_fp, prompt_tokens, completion_tokens) |
|
|
) + "\n\n" |
|
|
|
|
|
ts = int(datetime.now().timestamp()) |
|
|
yield "data: " + json.dumps( |
|
|
create_chat_completion_data("", request.model, ts, request_id, system_fp, prompt_tokens, completion_tokens, "stop") |
|
|
) + "\n\n" |
|
|
yield "data: [DONE]\n\n" |
|
|
|
|
|
except Exception as e: |
|
|
logger.error(f"Stream error {request_id}: {e}") |
|
|
|
|
|
upload_replaced_urls_to_r2(final_links, alt_text=get_last_user_prompt(request.messages)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def process_non_streaming_response(request: ChatRequest): |
|
|
system_fp = generate_system_fingerprint() |
|
|
request_id = f"chatcmpl-{uuid.uuid4()}" |
|
|
logger.info(f"Processing (non-stream) {request_id} - Model: {request.model}") |
|
|
|
|
|
agent_mode = AGENT_MODE.get(request.model, {}) |
|
|
trending_mode = TRENDING_AGENT_MODE.get(request.model, {}) |
|
|
prefix = MODEL_PREFIXES.get(request.model, "") |
|
|
|
|
|
headers_api = get_headers_api_chat(BASE_URL) |
|
|
if request.model == "o1-preview": |
|
|
await asyncio.sleep(random.randint(20, 60)) |
|
|
|
|
|
h = await getHid() |
|
|
if not h: |
|
|
raise HTTPException(status_code=500, detail="Missing h-value.") |
|
|
|
|
|
msgs = [message_to_dict(m, prefix) for m in request.messages] |
|
|
|
|
|
json_data = { |
|
|
"agentMode": agent_mode, |
|
|
"clickedAnswer2": False, |
|
|
"clickedAnswer3": False, |
|
|
"reasoningMode": False, |
|
|
"clickedForceWebSearch": False, |
|
|
"codeInterpreterMode": False, |
|
|
"codeModelMode": True, |
|
|
"githubToken": "", |
|
|
"deepSearchMode": False, |
|
|
"domains": None, |
|
|
"id": request_id, |
|
|
"imageGenerationMode": False, |
|
|
"isChromeExt": False, |
|
|
"isMicMode": False, |
|
|
"isPremium": True, |
|
|
"isMemoryEnabled": False, |
|
|
"maxTokens": request.max_tokens, |
|
|
"messages": msgs, |
|
|
"mobileClient": False, |
|
|
"playgroundTemperature": request.temperature, |
|
|
"playgroundTopP": request.top_p, |
|
|
"previewToken": None, |
|
|
"trendingAgentMode": trending_mode, |
|
|
"userId": None, |
|
|
"userSelectedModel": MODEL_MAPPING.get(request.model, request.model), |
|
|
"userSystemPrompt": None, |
|
|
"validated": h, |
|
|
"visitFromDelta": False, |
|
|
"webSearchModePrompt": False, |
|
|
"vscodeClient": False, |
|
|
"designerMode": False, |
|
|
"workspaceId": "", |
|
|
"beastMode": False, |
|
|
"customProfile": {"name": "", "occupation": "", "traits": [], "additionalInfo": "", "enableNewChats": False}, |
|
|
"webSearchModeOption": {"autoMode": False, "webMode": False, "offlineMode": True}, |
|
|
"session": generate_session_data(), |
|
|
} |
|
|
|
|
|
prompt_tokens = sum( |
|
|
calculate_tokens(m.get("content", ""), request.model) + |
|
|
sum(calculate_tokens(img["contents"], request.model) for img in m.get("data", {}).get("imagesData", [])) |
|
|
for m in msgs |
|
|
) |
|
|
|
|
|
full_resp = "" |
|
|
final_links: List[str] = [] |
|
|
|
|
|
async with httpx.AsyncClient() as client: |
|
|
try: |
|
|
resp = await client.post(f"{BASE_URL}/api/chat", headers=headers_api, json=json_data) |
|
|
resp.raise_for_status() |
|
|
full_resp = resp.text |
|
|
except Exception as e: |
|
|
return { |
|
|
"id": request_id, |
|
|
"object": "chat.completion", |
|
|
"created": int(datetime.now().timestamp()), |
|
|
"model": request.model, |
|
|
"system_fingerprint": system_fp, |
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": str(e)}, "finish_reason": "error"}], |
|
|
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": 0, "total_tokens": prompt_tokens}, |
|
|
} |
|
|
|
|
|
full_resp = full_resp.replace(BLOCKED_MESSAGE, "").strip() |
|
|
full_resp = full_resp.replace("https://storage.googleapis.com", "https://cdn.snapzion.com") |
|
|
final_links.extend(re.findall(r"(https://cdn\.snapzion\.com[^\s\)]+)", full_resp)) |
|
|
clean = strip_model_prefix(full_resp, prefix) |
|
|
completion_tokens = calculate_tokens(clean, request.model) |
|
|
|
|
|
upload_replaced_urls_to_r2(final_links, alt_text=get_last_user_prompt(request.messages)) |
|
|
|
|
|
return { |
|
|
"id": request_id, |
|
|
"object": "chat.completion", |
|
|
"created": int(datetime.now().timestamp()), |
|
|
"model": request.model, |
|
|
"system_fingerprint": system_fp, |
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": clean}, "finish_reason": "stop"}], |
|
|
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens}, |
|
|
} |
|
|
|
