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.
|
||||
"""
|
||||
|
||||
from llm_config import generate_content_config
|
||||
import logging
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
@@ -13,43 +15,58 @@ from google.genai import types
|
||||
from google.oauth2 import service_account
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from llm_config import generate_content_config
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Load environment variables from .env file
|
||||
load_dotenv()
|
||||
|
||||
|
||||
def get_credentials():
|
||||
"""Get Google Cloud credentials from environment.
|
||||
"""Get Google Cloud credentials and project ID from environment.
|
||||
|
||||
Supports two methods:
|
||||
1. GOOGLE_CREDENTIALS_JSON: Direct JSON content as string (production)
|
||||
2. GOOGLE_APPLICATION_CREDENTIALS: Path to JSON file (local development)
|
||||
|
||||
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
|
||||
credentials_json = os.getenv("GOOGLE_CREDENTIALS_JSON")
|
||||
|
||||
if credentials_json:
|
||||
try:
|
||||
credentials_info = json.loads(credentials_json)
|
||||
return service_account.Credentials.from_service_account_info(
|
||||
credentials_info
|
||||
project_id = credentials_info.get("project_id")
|
||||
|
||||
credentials = service_account.Credentials.from_service_account_info(
|
||||
credentials_info, scopes=SCOPES
|
||||
)
|
||||
return credentials, project_id
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid JSON in GOOGLE_CREDENTIALS_JSON: {e}")
|
||||
|
||||
# Fall back to file-based credentials (standard behavior)
|
||||
credentials_path = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
|
||||
|
||||
if credentials_path and os.path.exists(credentials_path):
|
||||
return service_account.Credentials.from_service_account_file(
|
||||
credentials_path
|
||||
with open(credentials_path) as f:
|
||||
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
|
||||
# (useful when running on Google Cloud with service account attached)
|
||||
return None
|
||||
return None, None
|
||||
|
||||
|
||||
# Gemini model name
|
||||
@@ -78,8 +95,13 @@ async def generate(prompt: str):
|
||||
def run_streaming():
|
||||
"""Run the synchronous streaming in a separate thread."""
|
||||
try:
|
||||
credentials = get_credentials()
|
||||
client = genai.Client(vertexai=True, credentials=credentials)
|
||||
credentials, project_id = get_credentials()
|
||||
|
||||
client = genai.Client(
|
||||
vertexai=True,
|
||||
credentials=credentials,
|
||||
project=project_id
|
||||
)
|
||||
|
||||
contents = [
|
||||
types.Content(role="user", parts=[
|
||||
|
||||
@@ -34,6 +34,6 @@ generate_content_config = types.GenerateContentConfig(
|
||||
],
|
||||
tools=tools,
|
||||
thinking_config=types.ThinkingConfig(
|
||||
thinking_level="MEDIUM",
|
||||
thinking_level="HIGH",
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user