make websocket concurrent

This commit is contained in:
Matteo Rosati
2026-01-21 10:45:14 +01:00
parent a8461b05c9
commit 08d37cdd94
2 changed files with 3 additions and 3 deletions

2
app.py
View File

@@ -91,7 +91,7 @@ async def websocket_endpoint(websocket: WebSocket):
while True: while True:
data = await websocket.receive_text() 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(chunk)
await websocket.send_text("<<END>>") await websocket.send_text("<<END>>")

4
lib.py
View File

@@ -18,7 +18,7 @@ CORPUS: str = "projects/520464122471/locations/europe-west3/ragCorpora/230584300
GEMINI_MODEL: str = "gemini-3-pro-preview" GEMINI_MODEL: str = "gemini-3-pro-preview"
def generate(prompt: str): async def generate(prompt: str):
"""Generate content using Gemini model with RAG retrieval. """Generate content using Gemini model with RAG retrieval.
This function creates a streaming response from the Gemini model, 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, model=GEMINI_MODEL,
contents=contents, contents=contents,
config=generate_content_config, config=generate_content_config,