fix docstrings, update env dist

main.py

@@ -1,19 +1,39 @@
+"""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()
 
-CORPUS = "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"
+# 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,
     )
 
-    model = "gemini-3-pro-preview"
     contents = [
         types.Content(role="user", parts=[types.Part.from_text(text=prompt)]),
     ]
@@ -21,7 +41,8 @@ def generate(prompt: str):
         types.Tool(
             retrieval=types.Retrieval(
                 vertex_rag_store=types.VertexRagStore(
-                    rag_resources=[types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
+                    rag_resources=[
+                        types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
                 )
             )
         )
@@ -32,14 +53,16 @@ def generate(prompt: str):
         top_p=0.95,
         max_output_tokens=65535,
         safety_settings=[
-            types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
+            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"),
+            types.SafetySetting(
+                category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
         ],
         tools=tools,
         thinking_config=types.ThinkingConfig(
@@ -48,7 +71,7 @@ def generate(prompt: str):
     )
 
     for chunk in client.models.generate_content_stream(
-        model=model,
+        model=GEMINI_MODEL,
         contents=contents,
         config=generate_content_config,
     ):