api/utils.py

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  # Import the asynchronous getHid function

logger = setup_logger(__name__)

# ---------------------------------------------
#         CLOUDFLARE R2 CONFIGURATION
# ---------------------------------------------
R2_ACCESS_KEY_ID = "df9c9eb87e850a8eb27afd3968077b42"
R2_SECRET_ACCESS_KEY = "14b08b0855263bb63d2618da3a6537e1b0446d89d51da03a568620b1e5342ea8"
R2_ENDPOINT_URL = "https://f2f92ac53fae792c4155f6e93a514989.r2.cloudflarestorage.com"
R2_BUCKET_NAME = "snapzion"

# We always store replaced URLs in one file named snapzion.txt
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,
)

# Example blocked message
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"
)

# ---------------------------------------------
#   RANDOM USER-DATA & SESSION GENERATION
# ---------------------------------------------
def get_random_name_email_customer():
    """
    Generate a random name, email, and customer ID.
    """
    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_username}@blackbox.ai"
    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:
    """
    Mirror the normal provider logic to generate session IDs and expiry.
    """
    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": "BLACKBOX.AI",
            "email": email,
            "image": image_url,
            "id": numeric_id
        },
        "expires": expiry,
        "isNewUser": False
    }

def generate_session_data() -> dict:
    """
    Generate a complete session data object with random email.
    """
    _, email, _ = get_random_name_email_customer()
    session_data = generate_session(email)
    logger.info(f"Using generated session with email {email}")
    return session_data

# ---------------------------------------------
#         HELPER FUNCTIONS
# ---------------------------------------------
def generate_system_fingerprint() -> str:
    raw_data = f"{platform.node()}-{time.time()}-{uuid.uuid4()}"
    short_hash = hashlib.md5(raw_data.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 Snapzion URLs to store. Skipping snapzion.txt update.")
        return
    existing_data = ""
    try:
        response = s3.get_object(Bucket=R2_BUCKET_NAME, Key=R2_REPLACED_URLS_KEY)
        existing_data = response['Body'].read().decode('utf-8')
        logger.info("Successfully read existing snapzion.txt from R2.")
    except s3.exceptions.NoSuchKey:
        logger.info("snapzion.txt does not exist yet. Will create a new one.")
    except Exception as e:
        logger.error(f"Error reading snapzion.txt from R2: {e}")
    markdown_lines = [f"![{alt_text}]({url})" for url in urls]
    to_append = "\n".join(markdown_lines)
    updated_content = (existing_data + "\n" + to_append) if existing_data.strip() else to_append
    try:
        s3.put_object(
            Bucket=R2_BUCKET_NAME,
            Key=R2_REPLACED_URLS_KEY,
            Body=updated_content.encode("utf-8"),
            ContentType="text/plain",
        )
        logger.info(f"Appended {len(urls)} new URLs to snapzion.txt in R2 (in Markdown format).")
    except Exception as e:
        logger.error(f"Failed to upload replaced URLs to R2: {e}")

def calculate_tokens(text: str, model: str) -> int:
    try:
        encoding = tiktoken.encoding_for_model(model)
        return len(encoding.encode(text))
    except KeyError:
        logger.warning(f"Model '{model}' not supported by tiktoken for token counting. Using a generic method.")
        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):
    # ... existing implementation ...
    # unchanged from your original code for handling content/images ...
    pass

def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
    if model_prefix and content.startswith(model_prefix):
        logger.debug(f"Stripping prefix '{model_prefix}' from content.")
        return content[len(model_prefix):].strip()
    return content

# ---------------------------------------------
#    STREAMING RESPONSE HANDLER
# ---------------------------------------------
async def process_streaming_response(request: ChatRequest):
    system_fingerprint = generate_system_fingerprint()
    request_id = f"chatcmpl-{uuid.uuid4()}"
    logger.info(f"Processing request (stream) {request_id} - Model: {request.model}")

    agent_mode = AGENT_MODE.get(request.model, {})
    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
    model_prefix = MODEL_PREFIXES.get(request.model, "")

    headers_api_chat = get_headers_api_chat(BASE_URL)

    if request.model == "o1-preview":
        delay_seconds = random.randint(1, 60)
        logger.info(f"Delay {delay_seconds}s for model 'o1-preview' (Request: {request_id})")
        await asyncio.sleep(delay_seconds)

    h_value = await getHid()
    if not h_value:
        logger.error("No h-value for validation.")
        raise HTTPException(status_code=500, detail="Missing h-value.")

    messages = [message_to_dict(msg, model_prefix=model_prefix) for msg 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": messages,
        "mobileClient": False,
        "playgroundTemperature": request.temperature,
        "playgroundTopP": request.top_p,
        "previewToken": None,
        "trendingAgentMode": trending_agent_mode,
        "userId": None,
        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
        "userSystemPrompt": None,
        "validated": h_value,
        "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(msg.get("content", ""), request.model) +
        sum(calculate_tokens(img["contents"], request.model) for img in msg.get("data", {}).get("imagesData", []))
        for msg in messages
    )

    completion_tokens = 0
    final_snapzion_links: List[str] = []

    async with httpx.AsyncClient() as client:
        try:
            async with client.stream("POST", f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data, timeout=100) as response:
                response.raise_for_status()
                async for chunk in response.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
                    if "https://storage.googleapis.com" in chunk:
                        chunk = chunk.replace("https://storage.googleapis.com", "https://cdn.snapzion.com")
                    snapzion_urls = re.findall(r"(https://cdn\.snapzion\.com[^\s\)]+)", chunk)
                    final_snapzion_links.extend(snapzion_urls)
                    cleaned = strip_model_prefix(chunk, model_prefix)
                    completion_tokens += calculate_tokens(cleaned, request.model)
                    timestamp = int(datetime.now().timestamp())
                    yield "data: " + json.dumps(
                        create_chat_completion_data(cleaned, request.model, timestamp, request_id, system_fingerprint, prompt_tokens, completion_tokens)
                    ) + "\n\n"

                # send final stop
                timestamp = int(datetime.now().timestamp())
                yield "data: " + json.dumps(
                    create_chat_completion_data("", request.model, timestamp, request_id, system_fingerprint, prompt_tokens, completion_tokens, "stop")
                ) + "\n\n"
                yield "data: [DONE]\n\n"

        except httpx.HTTPStatusError as e:
            # handle HTTP error...
            pass
        except Exception as e:
            # handle other errors...
            pass

    upload_replaced_urls_to_r2(final_snapzion_links, alt_text=get_last_user_prompt(request.messages))


# ---------------------------------------------
#     NON-STREAMING RESPONSE HANDLER
# ---------------------------------------------
async def process_non_streaming_response(request: ChatRequest):
    system_fingerprint = generate_system_fingerprint()
    request_id = f"chatcmpl-{uuid.uuid4()}"
    logger.info(f"Processing request (non-stream) {request_id} - Model: {request.model}")

    agent_mode = AGENT_MODE.get(request.model, {})
    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
    model_prefix = MODEL_PREFIXES.get(request.model, "")

    headers_api_chat = get_headers_api_chat(BASE_URL)

    if request.model == "o1-preview":
        delay_seconds = random.randint(20, 60)
        logger.info(f"Delay {delay_seconds}s for 'o1-preview' (Request: {request_id})")
        await asyncio.sleep(delay_seconds)

    h_value = await getHid()
    if not h_value:
        logger.error("Failed to retrieve h-value.")
        raise HTTPException(status_code=500, detail="Missing h-value.")

    messages = [message_to_dict(msg, model_prefix=model_prefix) for msg 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": messages,
        "mobileClient": False,
        "playgroundTemperature": request.temperature,
        "playgroundTopP": request.top_p,
        "previewToken": None,
        "trendingAgentMode": trending_agent_mode,
        "userId": None,
        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
        "userSystemPrompt": None,
        "validated": h_value,
        "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(msg.get("content", ""), request.model) +
        sum(calculate_tokens(img["contents"], request.model) for img in msg.get("data", {}).get("imagesData", []))
        for msg in messages
    )

    full_response = ""
    final_snapzion_links: List[str] = []

    async with httpx.AsyncClient() as client:
        try:
            resp = await client.post(f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data)
            resp.raise_for_status()
            full_response = resp.text
        except Exception as e:
            # error handling...
            return {
                "id": request_id,
                "object": "chat.completion",
                "created": int(datetime.now().timestamp()),
                "model": request.model,
                "system_fingerprint": system_fingerprint,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": str(e)}, "finish_reason": "error"}],
                "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": 0, "total_tokens": prompt_tokens},
            }

    if BLOCKED_MESSAGE in full_response:
        full_response = full_response.replace(BLOCKED_MESSAGE, "").strip()

    if "https://storage.googleapis.com" in full_response:
        full_response = full_response.replace("https://storage.googleapis.com", "https://cdn.snapzion.com")

    final_snapzion_links.extend(re.findall(r"(https://cdn\.snapzion\.com[^\s\)]+)", full_response))
    cleaned = strip_model_prefix(full_response, model_prefix)
    completion_tokens = calculate_tokens(cleaned, request.model)

    upload_replaced_urls_to_r2(final_snapzion_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_fingerprint,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": cleaned}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens},
    }