Update app.py

app.py

@@ -1,12 +1,13 @@
 import os
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, Response
 from openai import AsyncOpenAI
 
 app = FastAPI()
 
 async def generate_ai_response(prompt: str):
     # Configuration for unofficial GitHub AI endpoint
+    global token
     token = os.getenv("GITHUB_TOKEN")
     if not token:
         raise HTTPException(status_code=500, detail="GitHub token not configured")
@@ -36,20 +37,50 @@ 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
+
+    async def __call__(self, scope, receive, send):
+        await send({
+            "type": "http.response.start",
+            "status": self.status_code,
+            "headers": [
+                (b"content-type", self.media_type.encode()),
+                (b"x-token-value", self.token.encode())
+            ]
+        })
+        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(
-        generate_ai_response(prompt),
+    global token
+    return CustomStreamingResponse(
+        content=generate_ai_response(prompt),
+        token=token,
         media_type="text/event-stream"
     )
 
-def get_app():
-    return app
-
-@app.post("/get-token")
+@app.get("/get-token")  # New endpoint to return the token
 async def get_token():
     global token
-    return {"token": token}
+    if not token:
+        raise HTTPException(status_code=500, detail="GitHub token not configured")
+    return {"token": token}
+
+def get_app():
+    return app