From 1ed452f1d9a6044c219dda6a5750e134ab9c787d Mon Sep 17 00:00:00 2001
From: Matteo Rosati <rosati.matteo@gmail.com>
Date: Thu, 22 Jan 2026 09:44:17 +0100
Subject: [PATCH] fix the async mess

---
 app.py | 13 ++++++++++---
 lib.py | 17 +++++++++--------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/app.py b/app.py
index b7f065b..3e3e5b5 100644
--- a/app.py
+++ b/app.py
@@ -8,7 +8,14 @@ import os
 import logging
 from typing import Annotated
 
-from fastapi import FastAPI, Request, WebSocket, Depends, HTTPException, status
+from fastapi import (
+    FastAPI,
+    Request,
+    WebSocket,
+    Depends,
+    HTTPException,
+    status,
+)
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.templating import Jinja2Templates
 from fastapi.staticfiles import StaticFiles
@@ -67,7 +74,7 @@ templates = Jinja2Templates(directory=os.path.join(STATIC_DIR, TEMPLATES_DIR))
 
 
 @app.get("/")
-async def home(request: Request, username: Annotated[str, Depends(verify_credentials)]):
+def home(request: Request, username: Annotated[str, Depends(verify_credentials)]):
     """Render the main index page.
 
     Args:
@@ -91,7 +98,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>>")
diff --git a/lib.py b/lib.py
index 06ebaf1..bb807cd 100644
--- a/lib.py
+++ b/lib.py
@@ -12,13 +12,15 @@ from dotenv import load_dotenv
 load_dotenv()
 
 # Vertex AI RAG Corpus resource path
-CORPUS: str = "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"
+CORPUS: str = (
+    "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"
+)
 
 # Gemini model name
 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,
@@ -42,8 +44,7 @@ def generate(prompt: str):
         types.Tool(
             retrieval=types.Retrieval(
                 vertex_rag_store=types.VertexRagStore(
-                    rag_resources=[
-                        types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
+                    rag_resources=[types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
                 )
             )
         )
@@ -54,16 +55,14 @@ def generate(prompt: str):
         top_p=0.95,
         max_output_tokens=65535,
         safety_settings=[
-            types.SafetySetting(
-                category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
+            types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
             types.SafetySetting(
                 category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"
             ),
             types.SafetySetting(
                 category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"
             ),
-            types.SafetySetting(
-                category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
+            types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
         ],
         tools=tools,
         thinking_config=types.ThinkingConfig(
@@ -76,6 +75,8 @@ def generate(prompt: str):
         contents=contents,
         config=generate_content_config,
     ):
+        # DEBUG: Log chunk type to confirm generator behavior
+        print(f"[DEBUG] Chunk type: {type(chunk)}")
         if (
             not chunk.candidates
             or not chunk.candidates[0].content