91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
"""Google Gemini API integration for Akern-Genai project.
|
|
|
|
This module provides functionality to generate content using Google's Gemini model
|
|
with Vertex AI RAG (Retrieval-Augmented Generation) support.
|
|
"""
|
|
|
|
from google import genai
|
|
from google.genai import types
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables from .env file
|
|
load_dotenv()
|
|
|
|
# Vertex AI RAG Corpus resource path
|
|
CORPUS: str = "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"
|
|
|
|
# Gemini model name
|
|
GEMINI_MODEL: str = "gemini-3-pro-preview"
|
|
|
|
|
|
async def generate(prompt: str):
|
|
"""Generate content using Gemini model with RAG retrieval.
|
|
|
|
This function creates a streaming response from the Gemini model,
|
|
augmented with content from the configured RAG corpus.
|
|
|
|
Args:
|
|
prompt: The user's input prompt to generate content for.
|
|
|
|
Yields:
|
|
str: Text chunks from the generated response.
|
|
"""
|
|
client = genai.Client(
|
|
vertexai=True,
|
|
)
|
|
|
|
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",
|
|
),
|
|
)
|
|
|
|
async for chunk in client.models.generate_content_stream(
|
|
model=GEMINI_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="")
|