Ashrafb commited on
Commit
089a07b
1 Parent(s): 10501a1

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +9 -4
main.py CHANGED
@@ -32,15 +32,20 @@ async def respond(message, history: list[tuple[str, str]]):
32
 
33
  messages.append({"role": "user", "content": message})
34
 
35
- async for message in client.chat_completion(
36
  messages,
37
  max_tokens=MAX_TOKENS,
38
  stream=True,
39
  temperature=TEMPERATURE,
40
  top_p=TOP_P,
41
- ):
42
- token = message.choices[0].delta.content
43
- yield token
 
 
 
 
 
44
 
45
  @app.post("/generate/")
46
  async def generate(request: Request):
 
32
 
33
  messages.append({"role": "user", "content": message})
34
 
35
+ response = await client.chat_completion(
36
  messages,
37
  max_tokens=MAX_TOKENS,
38
  stream=True,
39
  temperature=TEMPERATURE,
40
  top_p=TOP_P,
41
+ )
42
+
43
+ if hasattr(response, "__aiter__"): # Check if the response supports async iteration
44
+ async for message in response:
45
+ yield message.choices[0].delta.content
46
+ else:
47
+ for message in response: # Handle regular iteration
48
+ yield message.choices[0].delta.content
49
 
50
  @app.post("/generate/")
51
  async def generate(request: Request):