diff --git a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py index 98f6f26b..26944818 100644 --- a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py @@ -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 diff --git a/hindsight-api/hindsight_api/engine/llm_wrapper.py b/hindsight-api/hindsight_api/engine/llm_wrapper.py index 9452d797..81ab8a2b 100644 --- a/hindsight-api/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api/hindsight_api/engine/llm_wrapper.py @@ -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": diff --git a/hindsight-api/hindsight_api/engine/vertexai_token_refresher.py b/hindsight-api/hindsight_api/engine/vertexai_token_refresher.py deleted file mode 100644 index 946e62a8..00000000 --- a/hindsight-api/hindsight_api/engine/vertexai_token_refresher.py +++ /dev/null @@ -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" - ) diff --git a/hindsight-api/tests/test_vertexai_provider.py b/hindsight-api/tests/test_vertexai_provider.py index 2afe2333..97d05858 100644 --- a/hindsight-api/tests/test_vertexai_provider.py +++ b/hindsight-api/tests/test_vertexai_provider.py @@ -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()