def retrieve(corpus_name: str, question: str, top_k: int = 3) -> list[dict]:
    resp = rag.retrieval_query(
        text=question, rag_resources=[rag.RagResource(rag_corpus=corpus_name)],
        rag_retrieval_config=rag.RagRetrievalConfig(top_k=top_k, filter=rag.Filter(vector_distance_threshold=0.6)))
    return format_contexts(resp)

rag_tool = types.Tool(retrieval=types.Retrieval(vertex_rag_store=types.VertexRagStore(
    rag_resources=[types.VertexRagStoreRagResource(rag_corpus=corpus_name)],
    rag_retrieval_config=types.RagRetrievalConfig(top_k=4))))
resp = client.models.generate_content(
    model=MODEL, contents=question,
    config=types.GenerateContentConfig(
        tools=[rag_tool], temperature=0,
        system_instruction="Réponds en français, uniquement à partir des documents. "
                           "Si l'information n'y est pas, dis-le."))

def cited_sources(gm) -> list[str]:
    chunks = (gm.grounding_chunks or []) if gm else []
    cited = sorted({i for s in (gm.grounding_supports or []) for i in (s.grounding_chunk_indices or [])}) if gm else []
    picked = [chunks[i] for i in cited if i < len(chunks)]
    uris = [ch.retrieved_context.uri for ch in picked if ch.retrieved_context and ch.retrieved_context.uri]
    return list(dict.fromkeys(u.rsplit("/", 1)[-1] for u in uris))

recherche_documents = VertexAiRagRetrieval(
    name="recherche_documents",
    description="Cherche dans les documents internes de Nimbus Conseil (RH, frais, sécurité, accueil).",
    rag_corpora=[RAG_CORPUS] if RAG_CORPUS else None,
    similarity_top_k=4,
    vector_distance_threshold=0.6,
)

root_agent = Agent(
    name="assistant_documentaire",
    model=Gemini(model=MODEL, retry_options=types.HttpRetryOptions(attempts=5, initial_delay=2)),
    instruction="Tu es l'assistant documentaire de Nimbus Conseil. ... cite le nom du document. "
                "Si l'information n'y est pas, dis-le.",
    tools=[recherche_documents],
)

export GOOGLE_CLOUD_LOCATION=europe-west4
export RAG_CORPUS=projects/<numéro>/locations/europe-west4/ragCorpora/<id>   # affiché par setup
python agent.py "Je suis arrivé lundi, combien de jours puis-je télétravailler cette semaine ?"

$ python rag_engine_lab.py ask "Combien de jours de télétravail par semaine et quelle indemnité ?"
  [0.184] politique-teletravail.md : 'Politique de télétravail — Nimbus Conseil (société fictive)\nVersion 3, …'
  [0.336] politique-teletravail.md : 'Équipement\nNimbus Conseil fournit un ordinateur portable, un écran …'
  [0.377] notes-de-frais.md : 'Procédure notes de frais — Nimbus Conseil (société fictive)\n…'
RÉPONSE : Il est possible de télétravailler jusqu'à 3 jours par semaine, dont au plus 2 jours consécutifs.
Une indemnité forfaitaire de 2,50 € par jour télétravaillé est versée, dans la limite de 55 € par mois.
SOURCES : ['politique-teletravail.md']

$ python rag_engine_lab.py ask "Quel est le budget formation annuel par salarié ?"
RÉPONSE : L'information concernant le budget formation annuel par salarié n'est pas disponible dans les documents fournis.

$ python agent.py "Je suis arrivé lundi, combien de jours puis-je télétravailler cette semaine ?"
Vous pouvez télétravailler 1 jour cette semaine, car le télétravail est limité à 1 jour par semaine
pendant les 4 premières semaines d'un nouvel arrivant (politique-teletravail.md).
[sources citées : ['politique-teletravail.md']]