from google import genai
from google.genai import types
from dotenv import load_dotenv


load_dotenv()

CORPUS = "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"


def generate(prompt: str):
    client = genai.Client(
        vertexai=True,
    )

    model = "gemini-3-pro-preview"
    contents = [
        types.Content(role="user", parts=[types.Part.from_text(text=prompt)]),
    ]
    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",
        ),
    )

    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        if (
            not chunk.candidates
            or not chunk.candidates[0].content
            or not chunk.candidates[0].content.parts
        ):
            continue

        yield chunk.text


if __name__ == "__main__":
    for chunk in generate("Come si calcola il rapporto sodio potassio?"):
        print(chunk, end="")