Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, Response
+from fastapi.responses import StreamingResponse
 from openai import AsyncOpenAI
 
 app = FastAPI()
@@ -37,46 +37,20 @@ 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")
     
-    global token  # Access the global token
-    return CustomStreamingResponse(
-        content=generate_ai_response(prompt),
-        token=token,
-        media_type="text/event-stream"
+    return StreamingResponse(
+        response_with_token(),
+        media_type="text/event-stream",
     )
 
 def get_app():
-    return app
+    return app
+
+@app.get("/gettoken")
+async def get_token():
+    global token
+    return {"token": token}