Refactor GoogleVertexLLMService to use GoogleLLMService as a base class

This commit is contained in:
Mark Backman
2025-11-05 09:33:02 -05:00
parent bbc7d3e2fb
commit 80d127aaa4
3 changed files with 75 additions and 37 deletions

View File

@@ -31,6 +31,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
you cancel a task with `PipelineTask.cancel(reason="cancellation your
reason")`.
### Changed
- Updated the `GoogleVertexLLMService` to use the `GoogleLLMService` as a base
class instead of the `OpenAILLMService`.
### Fixed
- Fixed an issue where the `SmallWebRTCRequest` dataclass in runner would scrub

View File

@@ -715,7 +715,6 @@ class GoogleLLMService(LLMService):
self._system_instruction = system_instruction
self._http_options = http_options
self._create_client(api_key, http_options)
self._settings = {
"max_tokens": params.max_tokens,
"temperature": params.temperature,
@@ -726,6 +725,9 @@ class GoogleLLMService(LLMService):
self._tools = tools
self._tool_config = tool_config
# Initialize the API client. Subclasses can override this if needed.
self.create_client()
def can_generate_metrics(self) -> bool:
"""Check if the service can generate usage metrics.
@@ -734,8 +736,9 @@ class GoogleLLMService(LLMService):
"""
return True
def _create_client(self, api_key: str, http_options: Optional[HttpOptions] = None):
self._client = genai.Client(api_key=api_key, http_options=http_options)
def create_client(self):
"""Create the Gemini client instance. Subclasses can override this."""
self._client = genai.Client(api_key=self._api_key, http_options=self._http_options)
async def run_inference(self, context: LLMContext | OpenAILLMContext) -> Optional[str]:
"""Run a one-shot, out-of-band (i.e. out-of-pipeline) inference with the given LLM context.

View File

@@ -6,8 +6,8 @@
"""Google Vertex AI LLM service implementation.
This module provides integration with Google's AI models via Vertex AI while
maintaining OpenAI API compatibility through Google's OpenAI-compatible endpoint.
This module provides integration with Google's AI models via Vertex AI,
extending the GoogleLLMService with Vertex AI authentication.
"""
import json
@@ -20,12 +20,14 @@ from typing import Optional
from loguru import logger
from pipecat.services.openai.llm import OpenAILLMService
from pipecat.services.google.llm import GoogleLLMService
try:
from google.auth import default
from google.auth.exceptions import GoogleAuthError
from google.auth.transport.requests import Request
from google.genai import Client
from google.genai.types import HttpOptions
from google.oauth2 import service_account
except ModuleNotFoundError as e:
@@ -36,19 +38,19 @@ except ModuleNotFoundError as e:
raise Exception(f"Missing module: {e}")
class GoogleVertexLLMService(OpenAILLMService):
"""Google Vertex AI LLM service with OpenAI API compatibility.
class GoogleVertexLLMService(GoogleLLMService):
"""Google Vertex AI LLM service extending GoogleLLMService.
Provides access to Google's AI models via Vertex AI while maintaining
OpenAI API compatibility. Handles authentication using Google service
account credentials and constructs appropriate endpoint URLs for
different GCP regions and projects.
Provides access to Google's AI models via Vertex AI while using the same
Google AI client and message format as GoogleLLMService. Handles authentication
using Google service account credentials and configures the client for
Vertex AI endpoints.
Reference:
https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/call-vertex-using-openai-library
https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference
"""
class InputParams(OpenAILLMService.InputParams):
class InputParams(GoogleLLMService.InputParams):
"""Input parameters specific to Vertex AI.
Parameters:
@@ -100,6 +102,11 @@ class GoogleVertexLLMService(OpenAILLMService):
model: str = "google/gemini-2.0-flash-001",
location: Optional[str] = None,
project_id: Optional[str] = None,
params: Optional[GoogleLLMService.InputParams] = None,
system_instruction: Optional[str] = None,
tools: Optional[list] = None,
tool_config: Optional[dict] = None,
http_options: Optional[HttpOptions] = None,
**kwargs,
):
"""Initializes the VertexLLMService.
@@ -110,11 +117,26 @@ class GoogleVertexLLMService(OpenAILLMService):
model: Model identifier (e.g., "google/gemini-2.0-flash-001").
location: GCP region for Vertex AI endpoint (e.g., "us-east4").
project_id: Google Cloud project ID.
**kwargs: Additional arguments passed to OpenAILLMService.
params: Input parameters for the model.
system_instruction: System instruction/prompt for the model.
tools: List of available tools/functions.
tool_config: Configuration for tool usage.
http_options: HTTP options for the client.
**kwargs: Additional arguments passed to GoogleLLMService.
"""
# Check if user incorrectly passed api_key, which is used by parent
# class but not here.
if "api_key" in kwargs:
logger.error(
"GoogleVertexLLMService does not accept 'api_key' parameter. "
"Use 'credentials' or 'credentials_path' instead for Vertex AI authentication."
)
raise ValueError(
"Invalid parameter 'api_key'. Use 'credentials' or 'credentials_path' for Vertex AI authentication."
)
# Handle deprecated InputParams fields
if "params" in kwargs and isinstance(kwargs["params"], GoogleVertexLLMService.InputParams):
params = kwargs["params"]
if params and isinstance(params, GoogleVertexLLMService.InputParams):
# Extract location and project_id from params if not provided
# directly, for backward compatibility
if project_id is None:
@@ -122,13 +144,12 @@ class GoogleVertexLLMService(OpenAILLMService):
if location is None:
location = params.location
# Convert to base InputParams
params = OpenAILLMService.InputParams(
params = GoogleLLMService.InputParams(
**params.model_dump(exclude={"location", "project_id"}, exclude_unset=True)
)
kwargs["params"] = params
# Validate project_id and location parameters
# NOTE: once we remove Vertex-spcific InputParams class, we can update
# NOTE: once we remove Vertex-specific InputParams class, we can update
# __init__() signature as follows:
# - location: str = "us-east4",
# - project_id: str,
@@ -143,29 +164,38 @@ class GoogleVertexLLMService(OpenAILLMService):
logger.warning("location is not provided. Defaulting to 'us-east4'.")
location = "us-east4" # Default location if not provided
base_url = self._get_base_url(location, project_id)
self._api_key = self._get_api_token(credentials, credentials_path)
# These need to be set before calling super().__init__() because
# super().__init__() invokes _create_client(), which needs these.
self._credentials = self._get_credentials(credentials, credentials_path)
self._project_id = project_id
self._location = location
# Call parent constructor with dummy api_key
# (api_key is required by parent class, but not actually used with Vertex)
super().__init__(
api_key=self._api_key,
base_url=base_url,
api_key="dummy",
model=model,
params=params,
system_instruction=system_instruction,
tools=tools,
tool_config=tool_config,
http_options=http_options,
**kwargs,
)
@staticmethod
def _get_base_url(location: str, project_id: str) -> str:
"""Construct the base URL for Vertex AI API."""
# Determine the correct API host based on location
if location == "global":
api_host = "aiplatform.googleapis.com"
else:
api_host = f"{location}-aiplatform.googleapis.com"
return f"https://{api_host}/v1/projects/{project_id}/locations/{location}/endpoints/openapi"
def create_client(self):
"""Create the Gemini client instance configured for Vertex AI."""
self._client = Client(
vertexai=True,
credentials=self._credentials,
project=self._project_id,
location=self._location,
http_options=self._http_options,
)
@staticmethod
def _get_api_token(credentials: Optional[str], credentials_path: Optional[str]) -> str:
"""Retrieve an authentication token using Google service account credentials.
def _get_credentials(credentials: Optional[str], credentials_path: Optional[str]):
"""Retrieve Credentials using Google service account credentials.
Supports multiple authentication methods:
1. Direct JSON credentials string
@@ -177,7 +207,7 @@ class GoogleVertexLLMService(OpenAILLMService):
credentials_path: Path to the service account JSON file.
Returns:
OAuth token for API authentication.
Google credentials object for API authentication.
Raises:
ValueError: If no valid credentials are provided or found.
@@ -209,4 +239,4 @@ class GoogleVertexLLMService(OpenAILLMService):
creds.refresh(Request()) # Ensure token is up-to-date, lifetime is 1 hour.
return creds.token
return creds