Update api/utils.py

api/utils.py

@@ -37,8 +37,6 @@ 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(
@@ -48,7 +46,6 @@ s3 = boto3.client(
     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"
@@ -58,15 +55,12 @@ BLOCKED_MESSAGE = (
 #   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"
+    random_email = f"{email_username}@gmail.com"
     suffix_length = len("Rldf7IKdNhdhiw")
     suffix_chars = string.ascii_letters + string.digits
     random_suffix = ''.join(random.choice(suffix_chars) for _ in range(suffix_length))
@@ -75,9 +69,6 @@ def get_random_name_email_customer():
     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')
@@ -88,7 +79,7 @@ def generate_session(email: str, id_length: int = 21, days_ahead: int = 365) ->
 
     return {
         "user": {
-            "name": "BLACKBOX.AI",
+            "name": "SNAPZION",
             "email": email,
             "image": image_url,
             "id": numeric_id
@@ -98,9 +89,6 @@ def generate_session(email: str, id_length: int = 21, days_ahead: int = 365) ->
     }
 
 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}")
@@ -110,8 +98,8 @@ def generate_session_data() -> dict:
 #         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]
+    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:
@@ -127,37 +115,36 @@ def get_last_user_prompt(messages: List[Any]) -> str:
 
 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.")
+        logger.info("No replaced or final URLs to store.")
         return
-    existing_data = ""
+    existing = ""
     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.")
+        resp = s3.get_object(Bucket=R2_BUCKET_NAME, Key=R2_REPLACED_URLS_KEY)
+        existing = resp["Body"].read().decode()
     except s3.exceptions.NoSuchKey:
-        logger.info("snapzion.txt does not exist yet. Will create a new one.")
+        pass
     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
+        logger.error(f"Error reading {R2_REPLACED_URLS_KEY}: {e}")
+
+    markdown = "\n".join(f"![{alt_text}]({u})" 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=updated_content.encode("utf-8"),
+            Body=content.encode(),
             ContentType="text/plain",
         )
-        logger.info(f"Appended {len(urls)} new URLs to snapzion.txt in R2 (in Markdown format).")
+        logger.info(f"Appended {len(urls)} URLs to {R2_REPLACED_URLS_KEY}.")
     except Exception as e:
-        logger.error(f"Failed to upload replaced URLs to R2: {e}")
+        logger.error(f"Error writing {R2_REPLACED_URLS_KEY}: {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.")
+        enc = tiktoken.encoding_for_model(model)
+        return len(enc.encode(text))
+    except Exception:
         return len(text.split())
 
 def create_chat_completion_data(
@@ -188,13 +175,47 @@ def create_chat_completion_data(
     }
 
 def message_to_dict(message, model_prefix: Optional[str] = None):
-    # ... existing implementation ...
-    # unchanged from your original code for handling content/images ...
-    pass
+    """
+    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
+        }
+        # embed any extra URLs in content list
+        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):
-        logger.debug(f"Stripping prefix '{model_prefix}' from content.")
         return content[len(model_prefix):].strip()
     return content
 
@@ -202,27 +223,23 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
 #    STREAMING RESPONSE HANDLER
 # ---------------------------------------------
 async def process_streaming_response(request: ChatRequest):
-    system_fingerprint = generate_system_fingerprint()
+    system_fp = generate_system_fingerprint()
     request_id = f"chatcmpl-{uuid.uuid4()}"
-    logger.info(f"Processing request (stream) {request_id} - Model: {request.model}")
+    logger.info(f"Processing (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)
+    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":
-        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)
+        await asyncio.sleep(random.randint(1, 60))
 
-    h_value = await getHid()
-    if not h_value:
-        logger.error("No h-value for validation.")
+    h = await getHid()
+    if not h:
         raise HTTPException(status_code=500, detail="Missing h-value.")
 
-    messages = [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages]
+    msgs = [message_to_dict(m, prefix) for m in request.messages]
 
     json_data = {
         "agentMode": agent_mode,
@@ -242,16 +259,16 @@ async def process_streaming_response(request: ChatRequest):
         "isPremium": True,
         "isMemoryEnabled": False,
         "maxTokens": request.max_tokens,
-        "messages": messages,
+        "messages": msgs,
         "mobileClient": False,
         "playgroundTemperature": request.temperature,
         "playgroundTopP": request.top_p,
         "previewToken": None,
-        "trendingAgentMode": trending_agent_mode,
+        "trendingAgentMode": trending_mode,
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
-        "validated": h_value,
+        "validated": h,
         "visitFromDelta": False,
         "webSearchModePrompt": False,
         "vscodeClient": False,
@@ -264,19 +281,19 @@ async def process_streaming_response(request: ChatRequest):
     }
 
     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
+        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_snapzion_links: List[str] = []
+    final_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():
+            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$@$"):
@@ -285,59 +302,48 @@ async def process_streaming_response(request: ChatRequest):
                         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())
+                    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(cleaned, request.model, timestamp, request_id, system_fingerprint, prompt_tokens, completion_tokens)
+                        create_chat_completion_data(clean, request.model, ts, request_id, system_fp, prompt_tokens, completion_tokens)
                     ) + "\n\n"
 
-                # send final stop
-                timestamp = int(datetime.now().timestamp())
+                ts = int(datetime.now().timestamp())
                 yield "data: " + json.dumps(
-                    create_chat_completion_data("", request.model, timestamp, request_id, system_fingerprint, prompt_tokens, completion_tokens, "stop")
+                    create_chat_completion_data("", request.model, ts, request_id, system_fp, 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))
+            logger.error(f"Stream error {request_id}: {e}")
 
+    upload_replaced_urls_to_r2(final_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()
+    system_fp = generate_system_fingerprint()
     request_id = f"chatcmpl-{uuid.uuid4()}"
-    logger.info(f"Processing request (non-stream) {request_id} - Model: {request.model}")
+    logger.info(f"Processing (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)
+    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":
-        delay_seconds = random.randint(20, 60)
-        logger.info(f"Delay {delay_seconds}s for 'o1-preview' (Request: {request_id})")
-        await asyncio.sleep(delay_seconds)
+        await asyncio.sleep(random.randint(20, 60))
 
-    h_value = await getHid()
-    if not h_value:
-        logger.error("Failed to retrieve h-value.")
+    h = await getHid()
+    if not h:
         raise HTTPException(status_code=500, detail="Missing h-value.")
 
-    messages = [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages]
+    msgs = [message_to_dict(m, prefix) for m in request.messages]
 
     json_data = {
         "agentMode": agent_mode,
@@ -357,16 +363,16 @@ async def process_non_streaming_response(request: ChatRequest):
         "isPremium": True,
         "isMemoryEnabled": False,
         "maxTokens": request.max_tokens,
-        "messages": messages,
+        "messages": msgs,
         "mobileClient": False,
         "playgroundTemperature": request.temperature,
         "playgroundTopP": request.top_p,
         "previewToken": None,
-        "trendingAgentMode": trending_agent_mode,
+        "trendingAgentMode": trending_mode,
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
-        "validated": h_value,
+        "validated": h,
         "visitFromDelta": False,
         "webSearchModePrompt": False,
         "vscodeClient": False,
@@ -379,49 +385,44 @@ async def process_non_streaming_response(request: ChatRequest):
     }
 
     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
+        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_response = ""
-    final_snapzion_links: List[str] = []
+    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_chat, json=json_data)
+            resp = await client.post(f"{BASE_URL}/api/chat", headers=headers_api, json=json_data)
             resp.raise_for_status()
-            full_response = resp.text
+            full_resp = 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,
+                "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},
             }
 
-    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")
+    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)
 
-    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))
+    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_fingerprint,
-        "choices": [{"index": 0, "message": {"role": "assistant", "content": cleaned}, "finish_reason": "stop"}],
+        "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},
     }