Switch Vertex AI provider to native genai SDK (#242)
Replace the OpenAI-compatible endpoint approach with the native google-genai SDK for Vertex AI. This eliminates the custom token refresher, TokenInjectingTransport, and async lifecycle complexity while also removing the 8192 output token cap that the OpenAI endpoint enforced. Changes: - vertexai provider now uses genai.Client(vertexai=True) instead of AsyncOpenAI with token-injecting transport - Routes through existing _call_gemini/_call_with_tools_gemini paths - Strips google/ prefix from model names (native SDK uses bare names) - Preserves service account key auth via credentials parameter - Delete vertexai_token_refresher.py (no longer needed) - Strip markdown code fences in consolidator JSON parsing - Rewrite vertexai tests for native SDK integration
This commit is contained in:
Chris Bartholomew
GitHub
parent
c2ac7d0440
commit
49ae55af03
4 changed files with 169 additions and 374 deletions
|
|
@ -865,7 +865,14 @@ Focus on DURABLE knowledge that serves this mission, not ephemeral state.
|
|||
)
|
||||
# Parse JSON response - should be an array
|
||||
if isinstance(result, str):
|
||||
result = json.loads(result)
|
||||
# Strip markdown code fences (some models wrap JSON in ```json ... ```)
|
||||
clean = result.strip()
|
||||
if clean.startswith("```"):
|
||||
clean = clean.split("\n", 1)[1] if "\n" in clean else clean[3:]
|
||||
if clean.endswith("```"):
|
||||
clean = clean[:-3]
|
||||
clean = clean.strip()
|
||||
result = json.loads(clean)
|
||||
# Ensure result is a list
|
||||
if isinstance(result, list):
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -105,9 +105,6 @@ class LLMProvider:
|
|||
self._mock_calls: list[dict] = []
|
||||
self._mock_response: Any = None
|
||||
|
||||
# Vertex AI token refresher
|
||||
self._vertexai_refresher: Any = None
|
||||
|
||||
# Set default base URLs
|
||||
if not self.base_url:
|
||||
if self.provider == "groq":
|
||||
|
|
@ -117,62 +114,48 @@ class LLMProvider:
|
|||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
|
||||
# Handle Vertex AI provider
|
||||
if self.provider == "vertexai":
|
||||
if not VERTEXAI_AVAILABLE:
|
||||
raise ValueError("Vertex AI requires 'google-auth' package. Install with: pip install google-auth")
|
||||
# Vertex AI config — stored for client creation below
|
||||
self._vertexai_project_id: str | None = None
|
||||
self._vertexai_region: str | None = None
|
||||
self._vertexai_credentials: Any = None
|
||||
|
||||
if self.provider == "vertexai":
|
||||
from ..config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
project_id = config.llm_vertexai_project_id
|
||||
if not project_id:
|
||||
self._vertexai_project_id = config.llm_vertexai_project_id
|
||||
if not self._vertexai_project_id:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID is required for Vertex AI provider. "
|
||||
"Set it to your GCP project ID."
|
||||
)
|
||||
|
||||
region = config.llm_vertexai_region or "us-central1"
|
||||
self._vertexai_region = config.llm_vertexai_region or "us-central1"
|
||||
service_account_key = config.llm_vertexai_service_account_key
|
||||
|
||||
# Try ADC first
|
||||
credentials = None
|
||||
auth_method = None
|
||||
|
||||
try:
|
||||
credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
|
||||
auth_method = "ADC"
|
||||
logger.info("Vertex AI: Using Application Default Credentials")
|
||||
except google.auth.exceptions.DefaultCredentialsError:
|
||||
logger.debug("Vertex AI: ADC not available, trying service account")
|
||||
|
||||
# Fall back to service account key file
|
||||
if credentials is None and service_account_key:
|
||||
try:
|
||||
credentials = service_account.Credentials.from_service_account_file(
|
||||
service_account_key,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
# Load explicit service account credentials if provided
|
||||
if service_account_key:
|
||||
if not VERTEXAI_AVAILABLE:
|
||||
raise ValueError(
|
||||
"Vertex AI service account auth requires 'google-auth' package. "
|
||||
"Install with: pip install google-auth"
|
||||
)
|
||||
auth_method = "Service Account"
|
||||
logger.info(f"Vertex AI: Using service account key: {service_account_key}")
|
||||
except Exception as e:
|
||||
logger.error(f"Vertex AI: Failed to load service account key: {e}")
|
||||
|
||||
if credentials is None:
|
||||
raise ValueError(
|
||||
"Vertex AI authentication failed. Either:\n"
|
||||
" 1. Set up ADC: gcloud auth application-default login\n"
|
||||
" 2. Set HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY to path of service account JSON key"
|
||||
self._vertexai_credentials = service_account.Credentials.from_service_account_file(
|
||||
service_account_key,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
logger.info(f"Vertex AI: Using service account key: {service_account_key}")
|
||||
|
||||
# Initialize token refresher
|
||||
from .vertexai_token_refresher import VertexAITokenRefresher
|
||||
# Strip google/ prefix from model name — native SDK uses bare names
|
||||
# e.g. "google/gemini-2.0-flash-lite-001" -> "gemini-2.0-flash-lite-001"
|
||||
if self.model.startswith("google/"):
|
||||
self.model = self.model[len("google/"):]
|
||||
|
||||
self._vertexai_refresher = VertexAITokenRefresher(credentials, project_id, region)
|
||||
self.base_url = self._vertexai_refresher.get_base_url()
|
||||
|
||||
logger.info(f"Vertex AI: project={project_id}, region={region}, auth={auth_method}")
|
||||
logger.info(
|
||||
f"Vertex AI: project={self._vertexai_project_id}, region={self._vertexai_region}, "
|
||||
f"model={self.model}, auth={'service_account' if service_account_key else 'ADC'}"
|
||||
)
|
||||
|
||||
# Validate API key (not needed for ollama, lmstudio, vertexai, or mock)
|
||||
if self.provider not in ("ollama", "lmstudio", "vertexai", "mock") and not self.api_key:
|
||||
|
|
@ -202,30 +185,16 @@ class LLMProvider:
|
|||
anthropic_kwargs["timeout"] = self.timeout
|
||||
self._anthropic_client = AsyncAnthropic(**anthropic_kwargs)
|
||||
elif self.provider == "vertexai":
|
||||
# Custom transport for token injection
|
||||
class TokenInjectingTransport(httpx.AsyncHTTPTransport):
|
||||
def __init__(self, refresher, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._refresher = refresher
|
||||
|
||||
async def handle_async_request(self, request):
|
||||
token = self._refresher.get_token()
|
||||
request.headers["Authorization"] = f"Bearer {token}"
|
||||
return await super().handle_async_request(request)
|
||||
|
||||
transport = TokenInjectingTransport(self._vertexai_refresher)
|
||||
# Native genai SDK with Vertex AI — handles ADC automatically,
|
||||
# or uses explicit service account credentials if provided
|
||||
client_kwargs = {
|
||||
"api_key": "dummy", # Required by AsyncOpenAI but unused (we inject token via transport)
|
||||
"base_url": self.base_url,
|
||||
"max_retries": 0,
|
||||
"http_client": httpx.AsyncClient(transport=transport),
|
||||
"vertexai": True,
|
||||
"project": self._vertexai_project_id,
|
||||
"location": self._vertexai_region,
|
||||
}
|
||||
if self.timeout:
|
||||
client_kwargs["timeout"] = self.timeout
|
||||
self._client = AsyncOpenAI(**client_kwargs)
|
||||
|
||||
# Start background refresh
|
||||
self._vertexai_refresher.start_refresh_task()
|
||||
if self._vertexai_credentials is not None:
|
||||
client_kwargs["credentials"] = self._vertexai_credentials
|
||||
self._gemini_client = genai.Client(**client_kwargs)
|
||||
elif self.provider in ("ollama", "lmstudio"):
|
||||
# Use dummy key if not provided for local
|
||||
api_key = self.api_key or "local"
|
||||
|
|
@ -317,8 +286,8 @@ class LLMProvider:
|
|||
return_usage,
|
||||
)
|
||||
|
||||
# Handle Gemini provider separately
|
||||
if self.provider == "gemini":
|
||||
# Handle Gemini and Vertex AI providers (both use native genai SDK)
|
||||
if self.provider in ("gemini", "vertexai"):
|
||||
return await self._call_gemini(
|
||||
messages,
|
||||
response_format,
|
||||
|
|
@ -682,8 +651,8 @@ class LLMProvider:
|
|||
messages, tools, max_completion_tokens, max_retries, initial_backoff, max_backoff, start_time, scope
|
||||
)
|
||||
|
||||
# Handle Gemini (convert to Gemini tool format)
|
||||
if self.provider == "gemini":
|
||||
# Handle Gemini and Vertex AI (convert to Gemini tool format)
|
||||
if self.provider in ("gemini", "vertexai"):
|
||||
return await self._call_with_tools_gemini(
|
||||
messages, tools, max_retries, initial_backoff, max_backoff, start_time, scope
|
||||
)
|
||||
|
|
@ -1603,10 +1572,8 @@ class LLMProvider:
|
|||
self._mock_calls = []
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (e.g., stop token refresh tasks)."""
|
||||
if self._vertexai_refresher is not None:
|
||||
await self._vertexai_refresher.stop()
|
||||
logger.debug("Vertex AI token refresher stopped")
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def for_memory(cls) -> "LLMProvider":
|
||||
|
|
|
|||
|
|
@ -1,120 +0,0 @@
|
|||
"""Vertex AI token refresher with background refresh and caching."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VertexAITokenRefresher:
|
||||
"""
|
||||
Background token refresher for Vertex AI.
|
||||
|
||||
Refreshes Google Cloud access tokens every 50 minutes to ensure they don't expire (60-min default).
|
||||
Thread-safe token caching for concurrent access from multiple async tasks.
|
||||
"""
|
||||
|
||||
def __init__(self, credentials: Any, project_id: str, region: str):
|
||||
"""
|
||||
Initialize the token refresher.
|
||||
|
||||
Args:
|
||||
credentials: Google Cloud credentials object (from google.auth.default or service_account)
|
||||
project_id: GCP project ID
|
||||
region: GCP region (e.g., "us-central1")
|
||||
"""
|
||||
self._credentials = credentials
|
||||
self._project_id = project_id
|
||||
self._region = region
|
||||
|
||||
# Thread-safe token cache
|
||||
self._token: str | None = None
|
||||
self._token_expiry: datetime | None = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# Background refresh task
|
||||
self._refresh_task: asyncio.Task | None = None
|
||||
self._stop_event = asyncio.Event()
|
||||
|
||||
# Initial token fetch (synchronous, must complete before returning)
|
||||
self._refresh_token_sync()
|
||||
|
||||
def _refresh_token_sync(self) -> None:
|
||||
"""Synchronously refresh the token (thread-safe)."""
|
||||
try:
|
||||
import google.auth.transport.requests
|
||||
|
||||
request = google.auth.transport.requests.Request()
|
||||
self._credentials.refresh(request)
|
||||
|
||||
with self._lock:
|
||||
self._token = self._credentials.token
|
||||
self._token_expiry = self._credentials.expiry
|
||||
|
||||
logger.debug(f"Vertex AI token refreshed, expires at {self._token_expiry}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to refresh Vertex AI token: {e}")
|
||||
raise
|
||||
|
||||
async def _refresh_loop(self) -> None:
|
||||
"""Background refresh loop (runs every 50 minutes)."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
# Wait 50 minutes or until stop event
|
||||
await asyncio.wait_for(self._stop_event.wait(), timeout=50 * 60)
|
||||
# If we get here, stop was signaled
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
# 50 minutes passed, refresh token
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self._refresh_token_sync)
|
||||
except Exception as e:
|
||||
logger.error(f"Background token refresh failed: {e}")
|
||||
# Continue loop - next API call will fail with auth error
|
||||
|
||||
def start_refresh_task(self) -> None:
|
||||
"""Start the background refresh task."""
|
||||
if self._refresh_task is None or self._refresh_task.done():
|
||||
self._refresh_task = asyncio.create_task(self._refresh_loop())
|
||||
logger.info("Vertex AI token refresh task started (refreshes every 50 minutes)")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the background refresh task."""
|
||||
if self._refresh_task is not None and not self._refresh_task.done():
|
||||
self._stop_event.set()
|
||||
try:
|
||||
await asyncio.wait_for(self._refresh_task, timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("Vertex AI token refresh task did not stop within 5 seconds")
|
||||
logger.info("Vertex AI token refresh task stopped")
|
||||
|
||||
def get_token(self) -> str:
|
||||
"""
|
||||
Get current access token (thread-safe).
|
||||
|
||||
Returns:
|
||||
Current Google Cloud access token
|
||||
|
||||
Raises:
|
||||
RuntimeError: If token is not available
|
||||
"""
|
||||
with self._lock:
|
||||
if self._token is None:
|
||||
raise RuntimeError("Vertex AI token not available")
|
||||
return self._token
|
||||
|
||||
def get_base_url(self) -> str:
|
||||
"""
|
||||
Get the Vertex AI OpenAI-compatible endpoint URL.
|
||||
|
||||
Returns:
|
||||
Base URL for Vertex AI OpenAI API
|
||||
"""
|
||||
return (
|
||||
f"https://{self._region}-aiplatform.googleapis.com/v1beta1/"
|
||||
f"projects/{self._project_id}/locations/{self._region}/endpoints/openapi"
|
||||
)
|
||||
|
|
@ -1,10 +1,9 @@
|
|||
"""
|
||||
Test Vertex AI provider integration including token refresh and API calls.
|
||||
Test Vertex AI provider integration using native genai SDK.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -12,110 +11,38 @@ import pytest
|
|||
pytest.importorskip("google.auth")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_refresher_initialization():
|
||||
"""Test token refresher initialization with mocked credentials."""
|
||||
from hindsight_api.engine.vertexai_token_refresher import VertexAITokenRefresher
|
||||
|
||||
# Mock credentials
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "test-token-123"
|
||||
mock_credentials.expiry = None
|
||||
|
||||
with patch("google.auth.transport.requests.Request"):
|
||||
refresher = VertexAITokenRefresher(mock_credentials, "test-project", "us-central1")
|
||||
|
||||
# Verify token was fetched
|
||||
assert refresher.get_token() == "test-token-123"
|
||||
|
||||
# Verify base URL is correctly formatted
|
||||
expected_url = (
|
||||
"https://us-central1-aiplatform.googleapis.com/v1beta1/"
|
||||
"projects/test-project/locations/us-central1/endpoints/openapi"
|
||||
)
|
||||
assert refresher.get_base_url() == expected_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_refresher_background_refresh():
|
||||
"""Test that background refresh task starts and stops correctly."""
|
||||
from hindsight_api.engine.vertexai_token_refresher import VertexAITokenRefresher
|
||||
|
||||
# Mock credentials
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "test-token-123"
|
||||
mock_credentials.expiry = None
|
||||
|
||||
with patch("google.auth.transport.requests.Request"):
|
||||
refresher = VertexAITokenRefresher(mock_credentials, "test-project", "us-central1")
|
||||
|
||||
# Start refresh task
|
||||
refresher.start_refresh_task()
|
||||
assert refresher._refresh_task is not None
|
||||
assert not refresher._refresh_task.done()
|
||||
|
||||
# Stop refresh task
|
||||
await refresher.stop()
|
||||
assert refresher._refresh_task.done()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_refresher_thread_safety():
|
||||
"""Test that token access is thread-safe."""
|
||||
from hindsight_api.engine.vertexai_token_refresher import VertexAITokenRefresher
|
||||
|
||||
# Mock credentials
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "test-token-123"
|
||||
mock_credentials.expiry = None
|
||||
|
||||
with patch("google.auth.transport.requests.Request"):
|
||||
refresher = VertexAITokenRefresher(mock_credentials, "test-project", "us-central1")
|
||||
|
||||
# Access token from multiple tasks concurrently
|
||||
async def get_token_task():
|
||||
return refresher.get_token()
|
||||
|
||||
results = await asyncio.gather(*[get_token_task() for _ in range(10)])
|
||||
|
||||
# All should return the same token
|
||||
assert all(token == "test-token-123" for token in results)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_refresher_no_token_error():
|
||||
"""Test that getting token without refresh raises error."""
|
||||
from hindsight_api.engine.vertexai_token_refresher import VertexAITokenRefresher
|
||||
|
||||
# Mock credentials that fail to refresh
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = None
|
||||
|
||||
with patch("google.auth.transport.requests.Request") as mock_request:
|
||||
mock_request.side_effect = Exception("Refresh failed")
|
||||
|
||||
with pytest.raises(Exception, match="Refresh failed"):
|
||||
VertexAITokenRefresher(mock_credentials, "test-project", "us-central1")
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_missing_dependency():
|
||||
"""Test error when google-auth is not available."""
|
||||
"""Test error when google-auth is not available and service account key is set."""
|
||||
from hindsight_api.engine import llm_wrapper
|
||||
|
||||
# Temporarily disable Vertex AI availability
|
||||
# VERTEXAI_AVAILABLE only matters when a service account key is provided
|
||||
original_available = llm_wrapper.VERTEXAI_AVAILABLE
|
||||
try:
|
||||
llm_wrapper.VERTEXAI_AVAILABLE = False
|
||||
|
||||
with pytest.raises(ValueError, match="google-auth"):
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project",
|
||||
"HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY": "/path/to/key.json",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
clear_config_cache()
|
||||
|
||||
with pytest.raises(ValueError, match="google-auth"):
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
finally:
|
||||
llm_wrapper.VERTEXAI_AVAILABLE = original_available
|
||||
|
||||
|
|
@ -123,7 +50,6 @@ def test_llm_wrapper_vertexai_missing_dependency():
|
|||
def test_llm_wrapper_vertexai_missing_project_id():
|
||||
"""Test error when project ID is not configured."""
|
||||
with patch.dict(os.environ, {"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": ""}, clear=False):
|
||||
# Clear config cache to reload from env
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
|
@ -138,58 +64,52 @@ def test_llm_wrapper_vertexai_missing_project_id():
|
|||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
# Restore config cache
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_wrapper_vertexai_adc_auth():
|
||||
"""Test Vertex AI with ADC authentication (mocked)."""
|
||||
def test_llm_wrapper_vertexai_adc_auth():
|
||||
"""Test Vertex AI with ADC authentication creates native genai client."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "test-token-adc"
|
||||
mock_credentials.expiry = None
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
# Clear config cache to reload from env
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with patch("google.auth.default", return_value=(mock_credentials, "test-project")):
|
||||
with patch("google.auth.transport.requests.Request"):
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
# genai.Client handles ADC internally — just verify it creates the client
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
assert provider.provider == "vertexai"
|
||||
assert provider._vertexai_refresher is not None
|
||||
assert "aiplatform.googleapis.com" in provider.base_url
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
await provider.cleanup()
|
||||
assert provider.provider == "vertexai"
|
||||
assert provider.model == "gemini-2.0-flash-001" # google/ prefix stripped
|
||||
assert provider._gemini_client is not None
|
||||
|
||||
# Verify genai.Client was called with vertexai=True
|
||||
mock_client_cls.assert_called_once_with(
|
||||
vertexai=True,
|
||||
project="test-project",
|
||||
location="us-central1",
|
||||
)
|
||||
|
||||
# Restore config cache
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_wrapper_vertexai_sa_auth():
|
||||
"""Test Vertex AI with service account authentication (mocked)."""
|
||||
def test_llm_wrapper_vertexai_sa_auth():
|
||||
"""Test Vertex AI with service account authentication passes credentials to genai client."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
import google.auth.exceptions
|
||||
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "test-token-sa"
|
||||
mock_credentials.expiry = None
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
|
|
@ -199,69 +119,91 @@ async def test_llm_wrapper_vertexai_sa_auth():
|
|||
},
|
||||
clear=False,
|
||||
):
|
||||
# Clear config cache to reload from env
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
# Mock ADC failure, SA success
|
||||
with patch(
|
||||
"google.auth.default",
|
||||
side_effect=google.auth.exceptions.DefaultCredentialsError("ADC not available"),
|
||||
"google.oauth2.service_account.Credentials.from_service_account_file",
|
||||
return_value=mock_credentials,
|
||||
):
|
||||
with patch(
|
||||
"google.oauth2.service_account.Credentials.from_service_account_file",
|
||||
return_value=mock_credentials,
|
||||
):
|
||||
with patch("google.auth.transport.requests.Request"):
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
assert provider.provider == "vertexai"
|
||||
assert provider._vertexai_refresher is not None
|
||||
|
||||
# Cleanup
|
||||
await provider.cleanup()
|
||||
|
||||
# Restore config cache
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_wrapper_vertexai_auth_failure():
|
||||
"""Test Vertex AI with both ADC and SA auth failing."""
|
||||
import google.auth.exceptions
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
# Clear config cache to reload from env
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
# Mock both ADC and SA failures
|
||||
with patch(
|
||||
"google.auth.default",
|
||||
side_effect=google.auth.exceptions.DefaultCredentialsError("ADC failed"),
|
||||
):
|
||||
with pytest.raises(ValueError, match="authentication failed"):
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
LLMProvider(
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
# Restore config cache
|
||||
assert provider.provider == "vertexai"
|
||||
assert provider._gemini_client is not None
|
||||
|
||||
# Verify credentials were passed to genai.Client
|
||||
mock_client_cls.assert_called_once_with(
|
||||
vertexai=True,
|
||||
project="test-project",
|
||||
location="us-central1",
|
||||
credentials=mock_credentials,
|
||||
)
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_strips_google_prefix():
|
||||
"""Test that google/ prefix is stripped from model name for native SDK."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="google/gemini-2.0-flash-lite-001",
|
||||
)
|
||||
|
||||
assert provider.model == "gemini-2.0-flash-lite-001"
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_llm_wrapper_vertexai_no_prefix_model():
|
||||
"""Test that model without google/ prefix is unchanged."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID": "test-project"},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="vertexai",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="gemini-2.0-flash-001",
|
||||
)
|
||||
|
||||
assert provider.model == "gemini-2.0-flash-001"
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
|
|
@ -299,5 +241,4 @@ async def test_vertexai_integration_actual_api():
|
|||
assert len(response) > 0
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await provider.cleanup()
|
||||
|
|
|
|||
Loading…
Reference in a new issue