Update api/utils.py

api/utils.py

@@ -4,7 +4,7 @@ import uuid
 import asyncio
 import random
 import string
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
 
 import httpx
 from fastapi import HTTPException
@@ -23,11 +23,57 @@ from api.logger import setup_logger
 from api.validate import getHid  # Import the asynchronous getHid function
 import tiktoken
 
+# ---------------- NEW IMPORTS FOR CLOUDFLARE R2 UPLOAD ----------------
+import boto3
+import re
+
 logger = setup_logger(__name__)
 
 # Define the 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"
 
+# ---------------- R2 CONFIG EXAMPLE ----------------
+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_PREFIX = "nai"  # or however you want to name your uploads
+
+# Initialize your R2 client
+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,
+)
+
+# Function to upload replaced URLs into R2 as a text file
+def upload_replaced_urls_to_r2(urls: List[str], request_id: str) -> None:
+    """
+    Given a list of replaced URLs, upload them as a .txt file to your R2 bucket.
+    The file name includes the request_id so it's unique for each request.
+    """
+    if not urls:
+        logger.info(f"No replaced URLs found for request {request_id}. Skipping upload.")
+        return
+
+    # Join all replaced URLs with a newline
+    content_body = "\n".join(urls)
+
+    # Create an object key, for example: replaced-urls/chatcmpl-<uuid>.txt
+    object_key = f"{R2_REPLACED_URLS_KEY_PREFIX}/{request_id}.txt"
+
+    try:
+        s3.put_object(
+            Bucket=R2_BUCKET_NAME,
+            Key=object_key,
+            Body=content_body.encode("utf-8"),
+        )
+        logger.info(f"Uploaded replaced URLs to R2: {object_key}")
+    except Exception as e:
+        logger.error(f"Failed to upload replaced URLs to R2: {e}")
+
+
 # Function to calculate tokens using tiktoken
 def calculate_tokens(text: str, model: str) -> int:
     try:
@@ -100,7 +146,8 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Process streaming response with headers from config.py
+
+# --------------------- PROCESS STREAMING RESPONSE ---------------------
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -168,6 +215,10 @@ async def process_streaming_response(request: ChatRequest):
             prompt_tokens += calculate_tokens(message['data']['imageBase64'], request.model)
 
     completion_tokens = 0
+
+    # We'll keep track of any replaced URLs in this list
+    replaced_urls = []
+
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -181,23 +232,39 @@ async def process_streaming_response(request: ChatRequest):
                 async for chunk in response.aiter_text():
                     timestamp = int(datetime.now().timestamp())
                     if chunk:
-                        content = chunk
-                        if content.startswith("$@$v=undefined-rv1$@$"):
-                            content = content[21:]
+                        # First handle any prefix from the server, if present
+                        if chunk.startswith("$@$v=undefined-rv1$@$"):
+                            chunk = chunk[21:]
+
                         # Remove the blocked message if present
-                        if BLOCKED_MESSAGE in content:
+                        if BLOCKED_MESSAGE in chunk:
                             logger.info(
                                 f"Blocked message detected in response for Request ID {request_id}."
                             )
-                            content = content.replace(BLOCKED_MESSAGE, '').strip()
-                            if not content:
+                            chunk = chunk.replace(BLOCKED_MESSAGE, '').strip()
+                            if not chunk:
                                 continue  # Skip if content is empty after removal
-                        cleaned_content = strip_model_prefix(content, model_prefix)
+
+                        # ---------------- REPLACE STORAGE URLS ----------------
+                        # We'll search for "https://storage.googleapis.com" and replace
+                        # with "https://cdn.snapzion.com"
+                        # Also collect them for uploading to R2
+                        found_urls = re.findall(r"https://storage\.googleapis\.com\S*", chunk)
+                        if found_urls:
+                            replaced_urls.extend(found_urls)
+                            chunk = chunk.replace("https://storage.googleapis.com", 
+                                                  "https://cdn.snapzion.com")
+
+                        cleaned_content = strip_model_prefix(chunk, model_prefix)
                         completion_tokens += calculate_tokens(cleaned_content, request.model)
+
                         yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp, request_id))}\n\n"
 
+                # At the end of the stream, finalize the response
+                # yield final piece with finish_reason="stop"
                 yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, request_id, prompt_tokens, completion_tokens, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
+
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred for Request ID {request_id}: {e}")
             error_message = f"HTTP error occurred: {e}"
@@ -209,7 +276,7 @@ async def process_streaming_response(request: ChatRequest):
             
             yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp, request_id, prompt_tokens, completion_tokens, 'error'))}\n\n"
             yield "data: [DONE]\n\n"
-            # raise HTTPException(status_code=e.response.status_code, detail=error_message)
+
         except httpx.RequestError as e:
             logger.error(
                 f"Error occurred during request for Request ID {request_id}: {e}"
@@ -217,15 +284,18 @@ async def process_streaming_response(request: ChatRequest):
             error_message = f"Request error occurred: {e}"
             yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp, request_id, prompt_tokens, completion_tokens, 'error'))}\n\n"
             yield "data: [DONE]\n\n"
-            # raise HTTPException(status_code=500, detail=error_message)
+
         except Exception as e:
             logger.error(f"An unexpected error occurred for Request ID {request_id}: {e}")
             error_message = f"An unexpected error occurred: {e}"
             yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp, request_id, prompt_tokens, completion_tokens, 'error'))}\n\n"
             yield "data: [DONE]\n\n"
-            # raise HTTPException(status_code=500, detail=error_message)
 
-# Process non-streaming response with headers from config.py
+    # Once the entire stream is done, upload any replaced URLs to R2
+    upload_replaced_urls_to_r2(replaced_urls, request_id)
+
+
+# --------------------- PROCESS NON-STREAMING RESPONSE ---------------------
 async def process_non_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -253,6 +323,7 @@ async def process_non_streaming_response(request: ChatRequest):
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
+    # Hard-coded for demonstration
     h_value = "00f37b34-a166-4efb-bce5-1312d87f2f94"
     if not h_value:
         logger.error("Failed to retrieve h-value for validation.")
@@ -298,6 +369,9 @@ async def process_non_streaming_response(request: ChatRequest):
             prompt_tokens += calculate_tokens(message['data']['imageBase64'], request.model)
 
     full_response = ""
+    # We'll keep track of any replaced URLs here
+    replaced_urls = []
+
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -381,10 +455,11 @@ async def process_non_streaming_response(request: ChatRequest):
                 },
             }
 
+    # Clean up server prefix if needed
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
-    # Remove the blocked message if present
+    # Remove blocked message
     if BLOCKED_MESSAGE in full_response:
         logger.info(
             f"Blocked message detected in response for Request ID {request_id}."
@@ -395,9 +470,21 @@ async def process_non_streaming_response(request: ChatRequest):
                 status_code=500, detail="Blocked message detected in response."
             )
 
+    # ---------------- REPLACE STORAGE URLS ----------------
+    # We'll search for "https://storage.googleapis.com" and replace
+    # with "https://cdn.snapzion.com"
+    found_urls = re.findall(r"https://storage\.googleapis\.com\S*", full_response)
+    if found_urls:
+        replaced_urls.extend(found_urls)
+        full_response = full_response.replace("https://storage.googleapis.com", "https://cdn.snapzion.com")
+
+    # Strip model prefix if present
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)
     completion_tokens = calculate_tokens(cleaned_full_response, request.model)
 
+    # Upload replaced URLs to R2
+    upload_replaced_urls_to_r2(replaced_urls, request_id)
+
     return {
         "id": request_id,
         "object": "chat.completion",
@@ -415,4 +502,4 @@ async def process_non_streaming_response(request: ChatRequest):
             "completion_tokens": completion_tokens,
             "total_tokens": prompt_tokens + completion_tokens,
         },
-    }
+    }