Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, Response
 from openai import AsyncOpenAI
 
 app = FastAPI()
@@ -37,15 +37,45 @@ async def generate_ai_response(prompt: str):
         yield f"Error: {str(err)}"
         raise HTTPException(status_code=500, detail="AI generation failed")
 
+class CustomStreamingResponse(Response):
+    def __init__(self, content, token, media_type="text/event-stream", status_code=200):
+        super().__init__(content=content, media_type=media_type, status_code=status_code)
+        self.token = token  # Store token as an attribute
+
+    async def __call__(self, scope, receive, send):
+        # Send the token as part of the initial response
+        await send({
+            "type": "http.response.start",
+            "status": self.status_code,
+            "headers": [
+                (b"content-type", self.media_type.encode()),
+                # Include token in a custom field (not headers)
+                (b"x-token-value", self.token.encode())
+            ]
+        })
+        # Stream the content
+        async for chunk in self.body_iterator:
+            await send({
+                "type": "http.response.body",
+                "body": chunk.encode() if isinstance(chunk, str) else chunk,
+                "more_body": True
+            })
+        await send({
+            "type": "http.response.body",
+            "body": b"",
+            "more_body": False
+        })
+
 @app.post("/generate")
 async def generate_response(prompt: str):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     
-    return StreamingResponse(
-        response_with_token(),
-        media_type="text/event-stream",
-        headers={"X-Token": token}  # Optionally include token in headers
+    global token  # Access the global token
+    return CustomStreamingResponse(
+        content=generate_ai_response(prompt),
+        token=token,
+        media_type="text/event-stream"
     )
 
 def get_app():