Compare commits
5 Commits
20919298c8
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
865f4ee1b6 | ||
|
|
3a7f329499 | ||
|
|
f14c56ddd3 | ||
|
|
da726dc777 | ||
|
|
ead140a2f4 |
42
lib.py
42
lib.py
@@ -4,6 +4,8 @@ This module provides functionality to generate content using Google's Gemini mod
|
|||||||
with Vertex AI RAG (Retrieval-Augmented Generation) support.
|
with Vertex AI RAG (Retrieval-Augmented Generation) support.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
from llm_config import generate_content_config
|
||||||
|
import logging
|
||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
@@ -13,43 +15,58 @@ from google.genai import types
|
|||||||
from google.oauth2 import service_account
|
from google.oauth2 import service_account
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
from llm_config import generate_content_config
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
# Load environment variables from .env file
|
# Load environment variables from .env file
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
|
|
||||||
def get_credentials():
|
def get_credentials():
|
||||||
"""Get Google Cloud credentials from environment.
|
"""Get Google Cloud credentials and project ID from environment.
|
||||||
|
|
||||||
Supports two methods:
|
Supports two methods:
|
||||||
1. GOOGLE_CREDENTIALS_JSON: Direct JSON content as string (production)
|
1. GOOGLE_CREDENTIALS_JSON: Direct JSON content as string (production)
|
||||||
2. GOOGLE_APPLICATION_CREDENTIALS: Path to JSON file (local development)
|
2. GOOGLE_APPLICATION_CREDENTIALS: Path to JSON file (local development)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
service_account.Credentials: The loaded credentials
|
tuple: (credentials, project_id) where credentials is a
|
||||||
|
service_account.Credentials object and project_id is the Google Cloud project
|
||||||
"""
|
"""
|
||||||
|
# OAuth scopes required for Vertex AI API
|
||||||
|
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||||
|
|
||||||
# Try to load credentials from JSON content directly
|
# Try to load credentials from JSON content directly
|
||||||
credentials_json = os.getenv("GOOGLE_CREDENTIALS_JSON")
|
credentials_json = os.getenv("GOOGLE_CREDENTIALS_JSON")
|
||||||
|
|
||||||
if credentials_json:
|
if credentials_json:
|
||||||
try:
|
try:
|
||||||
credentials_info = json.loads(credentials_json)
|
credentials_info = json.loads(credentials_json)
|
||||||
return service_account.Credentials.from_service_account_info(
|
project_id = credentials_info.get("project_id")
|
||||||
credentials_info
|
|
||||||
|
credentials = service_account.Credentials.from_service_account_info(
|
||||||
|
credentials_info, scopes=SCOPES
|
||||||
)
|
)
|
||||||
|
return credentials, project_id
|
||||||
except json.JSONDecodeError as e:
|
except json.JSONDecodeError as e:
|
||||||
raise ValueError(f"Invalid JSON in GOOGLE_CREDENTIALS_JSON: {e}")
|
raise ValueError(f"Invalid JSON in GOOGLE_CREDENTIALS_JSON: {e}")
|
||||||
|
|
||||||
# Fall back to file-based credentials (standard behavior)
|
# Fall back to file-based credentials (standard behavior)
|
||||||
credentials_path = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
|
credentials_path = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
|
||||||
|
|
||||||
if credentials_path and os.path.exists(credentials_path):
|
if credentials_path and os.path.exists(credentials_path):
|
||||||
return service_account.Credentials.from_service_account_file(
|
with open(credentials_path) as f:
|
||||||
credentials_path
|
credentials_info = json.load(f)
|
||||||
|
project_id = credentials_info.get("project_id")
|
||||||
|
|
||||||
|
credentials = service_account.Credentials.from_service_account_file(
|
||||||
|
credentials_path, scopes=SCOPES
|
||||||
)
|
)
|
||||||
|
return credentials, project_id
|
||||||
|
|
||||||
# If neither is provided, return None to let the client use default credentials
|
# If neither is provided, return None to let the client use default credentials
|
||||||
# (useful when running on Google Cloud with service account attached)
|
# (useful when running on Google Cloud with service account attached)
|
||||||
return None
|
return None, None
|
||||||
|
|
||||||
|
|
||||||
# Gemini model name
|
# Gemini model name
|
||||||
@@ -78,8 +95,13 @@ async def generate(prompt: str):
|
|||||||
def run_streaming():
|
def run_streaming():
|
||||||
"""Run the synchronous streaming in a separate thread."""
|
"""Run the synchronous streaming in a separate thread."""
|
||||||
try:
|
try:
|
||||||
credentials = get_credentials()
|
credentials, project_id = get_credentials()
|
||||||
client = genai.Client(vertexai=True, credentials=credentials)
|
|
||||||
|
client = genai.Client(
|
||||||
|
vertexai=True,
|
||||||
|
credentials=credentials,
|
||||||
|
project=project_id
|
||||||
|
)
|
||||||
|
|
||||||
contents = [
|
contents = [
|
||||||
types.Content(role="user", parts=[
|
types.Content(role="user", parts=[
|
||||||
|
|||||||
@@ -34,6 +34,6 @@ generate_content_config = types.GenerateContentConfig(
|
|||||||
],
|
],
|
||||||
tools=tools,
|
tools=tools,
|
||||||
thinking_config=types.ThinkingConfig(
|
thinking_config=types.ThinkingConfig(
|
||||||
thinking_level="MEDIUM",
|
thinking_level="HIGH",
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user