Update main.py

main.py

@@ -32,15 +32,20 @@ async def respond(message, history: list[tuple[str, str]]):
 
     messages.append({"role": "user", "content": message})
 
-    async for message in client.chat_completion(
+    response = await client.chat_completion(
         messages,
         max_tokens=MAX_TOKENS,
         stream=True,
         temperature=TEMPERATURE,
         top_p=TOP_P,
-    ):
-        token = message.choices[0].delta.content
-        yield token
+    )
+
+    if hasattr(response, "__aiter__"):  # Check if the response supports async iteration
+        async for message in response:
+            yield message.choices[0].delta.content
+    else:
+        for message in response:  # Handle regular iteration
+            yield message.choices[0].delta.content
 
 @app.post("/generate/")
 async def generate(request: Request):