68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
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="")
|