make websocket concurrent

app.py

@@ -91,7 +91,7 @@ async def websocket_endpoint(websocket: WebSocket):
     while True:
         data = await websocket.receive_text()
 
-        for chunk in generate(data):
+        async for chunk in generate(data):
             await websocket.send_text(chunk)
 
         await websocket.send_text("<<END>>")

lib.py

@@ -18,7 +18,7 @@ CORPUS: str = "projects/520464122471/locations/europe-west3/ragCorpora/230584300
 GEMINI_MODEL: str = "gemini-3-pro-preview"
 
 
-def generate(prompt: str):
+async def generate(prompt: str):
     """Generate content using Gemini model with RAG retrieval.
 
     This function creates a streaming response from the Gemini model,
@@ -70,7 +70,7 @@ def generate(prompt: str):
         ),
     )
 
-    for chunk in client.models.generate_content_stream(
+    async for chunk in client.models.generate_content_stream(
         model=GEMINI_MODEL,
         contents=contents,
         config=generate_content_config,