Update main.py

main.py

@@ -7,6 +7,7 @@ import random
 import logging
 import time
 from contextlib import asynccontextmanager
+import json
 
 # --- Production-Ready Configuration ---
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
@@ -26,11 +27,56 @@ 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 Function ---
+# --- Helper Functions ---
+
 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):
@@ -97,7 +143,8 @@ 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(
-                    rp_resp.aiter_raw(),
+                    # Use the new async generator to modify the stream
+                    modified_aiter_raw(rp_resp.aiter_raw()),
                     status_code=rp_resp.status_code,
                     headers=rp_resp.headers,
                     background=BackgroundTask(rp_resp.aclose),