Update main.py

main.py

@@ -7,7 +7,6 @@ import random
 import logging
 import time
 from contextlib import asynccontextmanager
-import json
 
 # --- Production-Ready Configuration ---
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
@@ -17,7 +16,7 @@ logging.basicConfig(
 )
 
 TARGET_URL = os.getenv("TARGET_URL", "https://api.gmi-serving.com")
-MAX_RETRIES = int(os.getenv("MAX_RETRIES", "10"))
+MAX_RETRIES = int(os.getenv("MAX_RETRIES", "12"))
 DEFAULT_RETRY_CODES = "429,500,502,503,504"
 RETRY_CODES_STR = os.getenv("RETRY_CODES", DEFAULT_RETRY_CODES)
 try:
@@ -27,56 +26,11 @@ except ValueError:
     logging.error(f"Invalid RETRY_CODES format: '{RETRY_CODES_STR}'. Falling back to default: {DEFAULT_RETRY_CODES}")
     RETRY_STATUS_CODES = {int(code.strip()) for code in DEFAULT_RETRY_CODES.split(',')}
 
-# --- Helper Functions ---
-
+# --- Helper Function ---
 def generate_random_ip():
     """Generates a random, valid-looking IPv4 address."""
     return ".".join(str(random.randint(1, 254)) for _ in range(4))
 
-async def modified_aiter_raw(original_aiter):
-    """
-    An async generator that intercepts and modifies the streaming data chunks.
-    It adds a prefix to the 'id' and includes a 'provider' field.
-    """
-    buffer = ""
-    async for chunk in original_aiter:
-        buffer += chunk.decode('utf-8')
-        while '\n' in buffer:
-            line, buffer = buffer.split('\n', 1)
-            if line.startswith('data:'):
-                try:
-                    # Strip the "data: " prefix to get the JSON string
-                    json_str = line[len('data: '):].strip()
-                    
-                    # Process only if it's not the SSE termination message
-                    if json_str and json_str != '[DONE]':
-                        data = json.loads(json_str)
-                        
-                        # Add 'NAI-' prefix to the id
-                        if 'id' in data:
-                            data['id'] = f"NAI-{data['id']}"
-                        
-                        # Add the provider field
-                        data['provider'] = 'TypeGPT'
-                        
-                        # Reconstruct the SSE data line
-                        modified_line = f"data: {json.dumps(data)}"
-                        yield (modified_line + '\n').encode('utf-8')
-                    else:
-                        # Pass through messages like 'data: [DONE]'
-                        yield (line + '\n').encode('utf-8')
-                except json.JSONDecodeError:
-                    # If it's not valid JSON, pass it through as is
-                    yield (line + '\n').encode('utf-8')
-            else:
-                # Pass through non-data lines (e.g., empty lines, comments)
-                yield (line + '\n').encode('utf-8')
-    
-    # Yield any remaining data in the buffer
-    if buffer:
-        yield buffer.encode('utf-8')
-
-
 # --- HTTPX Client Lifecycle Management ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -143,8 +97,7 @@ async def reverse_proxy_handler(request: Request):
                 log_func(f"Request finished: {request.method} {request.url.path} status_code={rp_resp.status_code} latency={duration_ms:.2f}ms")
                 
                 return StreamingResponse(
-                    # Use the new async generator to modify the stream
-                    modified_aiter_raw(rp_resp.aiter_raw()),
+                    rp_resp.aiter_raw(),
                     status_code=rp_resp.status_code,
                     headers=rp_resp.headers,
                     background=BackgroundTask(rp_resp.aclose),