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, List
+from typing import Any, Dict, Optional
 
 import httpx
 from fastapi import HTTPException
@@ -23,57 +23,11 @@ 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:
@@ -146,8 +100,7 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-
-# --------------------- PROCESS STREAMING RESPONSE ---------------------
+# Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -170,7 +123,7 @@ async def process_streaming_response(request: ChatRequest):
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
-    h_value = "00f37b34-a166-4efb-bce5-1312d87f2f94"
+    h_value = await getHid()
     if not h_value:
         logger.error("Failed to retrieve h-value for validation.")
         raise HTTPException(
@@ -215,10 +168,6 @@ 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(
@@ -232,39 +181,23 @@ async def process_streaming_response(request: ChatRequest):
                 async for chunk in response.aiter_text():
                     timestamp = int(datetime.now().timestamp())
                     if chunk:
-                        # First handle any prefix from the server, if present
-                        if chunk.startswith("$@$v=undefined-rv1$@$"):
-                            chunk = chunk[21:]
-
+                        content = chunk
+                        if content.startswith("$@$v=undefined-rv1$@$"):
+                            content = content[21:]
                         # Remove the blocked message if present
-                        if BLOCKED_MESSAGE in chunk:
+                        if BLOCKED_MESSAGE in content:
                             logger.info(
                                 f"Blocked message detected in response for Request ID {request_id}."
                             )
-                            chunk = chunk.replace(BLOCKED_MESSAGE, '').strip()
-                            if not chunk:
+                            content = content.replace(BLOCKED_MESSAGE, '').strip()
+                            if not content:
                                 continue  # Skip if content is empty after removal
-
-                        # ---------------- 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)
+                        cleaned_content = strip_model_prefix(content, 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}"
@@ -276,7 +209,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}"
@@ -284,18 +217,15 @@ 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)
 
-    # 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 ---------------------
+# Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -323,7 +253,6 @@ 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.")
@@ -369,9 +298,6 @@ 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(
@@ -455,11 +381,10 @@ 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 blocked message
+    # Remove the blocked message if present
     if BLOCKED_MESSAGE in full_response:
         logger.info(
             f"Blocked message detected in response for Request ID {request_id}."
@@ -470,21 +395,9 @@ 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",
@@ -503,3 +416,4 @@ async def process_non_streaming_response(request: ChatRequest):
             "total_tokens": prompt_tokens + completion_tokens,
         },
     }
+