admin/Akern-Genai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix the async mess

Browse Source
This commit is contained in:
Matteo Rosati
2026-01-22 09:44:17 +01:00
parent 31ee3fed8c
commit 1ed452f1d9
2 changed files with 19 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
app.py
View File

@@ -8,7 +8,14 @@ import os
import logging import logging
from typing import Annotated 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.security import HTTPBasic, HTTPBasicCredentials
from fastapi.templating import Jinja2Templates from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles from fastapi.staticfiles import StaticFiles
@@ -67,7 +74,7 @@ templates = Jinja2Templates(directory=os.path.join(STATIC_DIR, TEMPLATES_DIR))
@app.get("/") @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. """Render the main index page.
Args: Args:
@@ -91,7 +98,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>>")

17
lib.py
View File

@@ -12,13 +12,15 @@ from dotenv import load_dotenv
load_dotenv() load_dotenv()
# Vertex AI RAG Corpus resource path # 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 name
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,
@@ -42,8 +44,7 @@ def generate(prompt: str):
types.Tool( types.Tool(
retrieval=types.Retrieval( retrieval=types.Retrieval(
vertex_rag_store=types.VertexRagStore( vertex_rag_store=types.VertexRagStore(
rag_resources=[ rag_resources=[types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
) )
) )
) )
@@ -54,16 +55,14 @@ def generate(prompt: str):
top_p=0.95, top_p=0.95,
max_output_tokens=65535, max_output_tokens=65535,
safety_settings=[ safety_settings=[
types.SafetySetting( types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
types.SafetySetting( types.SafetySetting(
category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF" category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"
), ),
types.SafetySetting( types.SafetySetting(
category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF" category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"
), ),
types.SafetySetting( types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
], ],
tools=tools, tools=tools,
thinking_config=types.ThinkingConfig( thinking_config=types.ThinkingConfig(
@@ -76,6 +75,8 @@ def generate(prompt: str):
contents=contents, contents=contents,
config=generate_content_config, config=generate_content_config,
): ):
# DEBUG: Log chunk type to confirm generator behavior
print(f"[DEBUG] Chunk type: {type(chunk)}")
if ( if (
not chunk.candidates not chunk.candidates
or not chunk.candidates[0].content or not chunk.candidates[0].content