from google.genai import types

# Vertex AI RAG Corpus resource path
CORPUS: str = (
    "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"
)

tools = [
    types.Tool(
        retrieval=types.Retrieval(
            vertex_rag_store=types.VertexRagStore(
                rag_resources=[
                    types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
            )
        )
    )
]

generate_content_config = types.GenerateContentConfig(
    temperature=1,
    top_p=0.95,
    max_output_tokens=65535,
    safety_settings=[
        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"),
    ],
    tools=tools,
    thinking_config=types.ThinkingConfig(
        thinking_level="HIGH",
    ),
)