diff --git a/app.py b/app.py index b7f065b..fd0c23d 100644 --- a/app.py +++ b/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("<>") diff --git a/lib.py b/lib.py index eb2ab9a..75a4e80 100644 --- a/lib.py +++ b/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,