rename main -> lib

2026-01-21 10:17:29 +01:00
parent 5906bfb0ab
commit 993c5cf8aa
2 changed files with 1 additions and 1 deletions
--- a/lib.py
+++ b/lib.py
@@ -0,0 +1,90 @@
+"""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"
+
+
+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",
+        ),
+    )
+
+    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="")