abdullahalioo commited on
Commit
471a43e
·
verified ·
1 Parent(s): 81b6332

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -36
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import os
2
  from fastapi import FastAPI, HTTPException
3
- from fastapi.responses import StreamingResponse, Response
4
  from openai import AsyncOpenAI
5
 
6
  app = FastAPI()
@@ -37,46 +37,20 @@ async def generate_ai_response(prompt: str):
37
  yield f"Error: {str(err)}"
38
  raise HTTPException(status_code=500, detail="AI generation failed")
39
 
40
- class CustomStreamingResponse(Response):
41
- def __init__(self, content, token, media_type="text/event-stream", status_code=200):
42
- super().__init__(content=content, media_type=media_type, status_code=status_code)
43
- self.token = token # Store token as an attribute
44
-
45
- async def __call__(self, scope, receive, send):
46
- # Send the token as part of the initial response
47
- await send({
48
- "type": "http.response.start",
49
- "status": self.status_code,
50
- "headers": [
51
- (b"content-type", self.media_type.encode()),
52
- # Include token in a custom field (not headers)
53
- (b"x-token-value", self.token.encode())
54
- ]
55
- })
56
- # Stream the content
57
- async for chunk in self.body_iterator:
58
- await send({
59
- "type": "http.response.body",
60
- "body": chunk.encode() if isinstance(chunk, str) else chunk,
61
- "more_body": True
62
- })
63
- await send({
64
- "type": "http.response.body",
65
- "body": b"",
66
- "more_body": False
67
- })
68
-
69
  @app.post("/generate")
70
  async def generate_response(prompt: str):
71
  if not prompt:
72
  raise HTTPException(status_code=400, detail="Prompt cannot be empty")
73
 
74
- global token # Access the global token
75
- return CustomStreamingResponse(
76
- content=generate_ai_response(prompt),
77
- token=token,
78
- media_type="text/event-stream"
79
  )
80
 
81
  def get_app():
82
- return app
 
 
 
 
 
 
1
  import os
2
  from fastapi import FastAPI, HTTPException
3
+ from fastapi.responses import StreamingResponse
4
  from openai import AsyncOpenAI
5
 
6
  app = FastAPI()
 
37
  yield f"Error: {str(err)}"
38
  raise HTTPException(status_code=500, detail="AI generation failed")
39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
  @app.post("/generate")
41
  async def generate_response(prompt: str):
42
  if not prompt:
43
  raise HTTPException(status_code=400, detail="Prompt cannot be empty")
44
 
45
+ return StreamingResponse(
46
+ response_with_token(),
47
+ media_type="text/event-stream",
 
 
48
  )
49
 
50
  def get_app():
51
+ return app
52
+
53
+ @app.get("/gettoken")
54
+ async def get_token():
55
+ global token
56
+ return {"token": token}