[Feat] New Provider - VertexAI Agent Engine (#18014)
* init A2AProviderConfigManager * move file * move file * add pydnatic ai folder * init providers * test_pydantic_ai_non_streaming * fix import * INIT pydantic * use_a2a_form_fields * test_vertex_agent_engine_streaming * add agent_engine * init transform for agent engine * init agent engine * VertexAgentEngineSSEStreamIterator * sample * ui add new fields * fix vertex_credentials * working SSE iterator * TestVertexAgentEngineTransformRequest * fix code QA check * Potential fix for code scanning alert no. 3923: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
GitHub
parent
f58b76aee8
commit
32c07113cf
13
litellm/llms/vertex_ai/agent_engine/__init__.py
Normal file
13
litellm/llms/vertex_ai/agent_engine/__init__.py
Normal file
@ -0,0 +1,13 @@
|
||||
"""
|
||||
Vertex AI Agent Engine (Reasoning Engines) Provider
|
||||
|
||||
Supports Vertex AI Reasoning Engines via the :query and :streamQuery endpoints.
|
||||
"""
|
||||
|
||||
from litellm.llms.vertex_ai.agent_engine.transformation import (
|
||||
VertexAgentEngineConfig,
|
||||
VertexAgentEngineError,
|
||||
)
|
||||
|
||||
__all__ = ["VertexAgentEngineConfig", "VertexAgentEngineError"]
|
||||
|
||||
90
litellm/llms/vertex_ai/agent_engine/sse_iterator.py
Normal file
90
litellm/llms/vertex_ai/agent_engine/sse_iterator.py
Normal file
@ -0,0 +1,90 @@
|
||||
"""
|
||||
SSE Stream Iterator for Vertex AI Agent Engine.
|
||||
|
||||
Handles Server-Sent Events (SSE) streaming responses from Vertex AI Reasoning Engines.
|
||||
"""
|
||||
|
||||
from typing import Any, Union
|
||||
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.types.llms.openai import ChatCompletionUsageBlock
|
||||
from litellm.types.utils import (
|
||||
Delta,
|
||||
GenericStreamingChunk,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
)
|
||||
|
||||
|
||||
class VertexAgentEngineResponseIterator(BaseModelResponseIterator):
|
||||
"""
|
||||
Iterator for Vertex Agent Engine SSE streaming responses.
|
||||
|
||||
Uses BaseModelResponseIterator which handles sync/async iteration.
|
||||
We just need to implement chunk_parser to parse Vertex Agent Engine response format.
|
||||
"""
|
||||
|
||||
def __init__(self, streaming_response: Any, sync_stream: bool) -> None:
|
||||
super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
|
||||
|
||||
def chunk_parser(
|
||||
self, chunk: dict
|
||||
) -> Union[GenericStreamingChunk, ModelResponseStream]:
|
||||
"""
|
||||
Parse a Vertex Agent Engine response chunk into ModelResponseStream.
|
||||
|
||||
Vertex Agent Engine response format:
|
||||
{
|
||||
"content": {
|
||||
"parts": [{"text": "..."}],
|
||||
"role": "model"
|
||||
},
|
||||
"finish_reason": "STOP",
|
||||
"usage_metadata": {
|
||||
"prompt_token_count": 100,
|
||||
"candidates_token_count": 50,
|
||||
"total_token_count": 150
|
||||
}
|
||||
}
|
||||
"""
|
||||
# Extract text from content.parts
|
||||
text = None
|
||||
content = chunk.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
for part in parts:
|
||||
if isinstance(part, dict) and "text" in part:
|
||||
text = part["text"]
|
||||
break
|
||||
|
||||
# Extract finish_reason
|
||||
finish_reason = None
|
||||
raw_finish_reason = chunk.get("finish_reason")
|
||||
if raw_finish_reason == "STOP":
|
||||
finish_reason = "stop"
|
||||
elif raw_finish_reason:
|
||||
finish_reason = raw_finish_reason.lower()
|
||||
|
||||
# Extract usage from usage_metadata
|
||||
usage = None
|
||||
usage_metadata = chunk.get("usage_metadata", {})
|
||||
if usage_metadata:
|
||||
usage = ChatCompletionUsageBlock(
|
||||
prompt_tokens=usage_metadata.get("prompt_token_count", 0),
|
||||
completion_tokens=usage_metadata.get("candidates_token_count", 0),
|
||||
total_tokens=usage_metadata.get("total_token_count", 0),
|
||||
)
|
||||
|
||||
# Return ModelResponseStream (OpenAI-compatible chunk)
|
||||
return ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason=finish_reason,
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=text,
|
||||
role="assistant" if text else None,
|
||||
),
|
||||
)
|
||||
],
|
||||
usage=usage,
|
||||
)
|
||||
508
litellm/llms/vertex_ai/agent_engine/transformation.py
Normal file
508
litellm/llms/vertex_ai/agent_engine/transformation.py
Normal file
@ -0,0 +1,508 @@
|
||||
"""
|
||||
Transformation for Vertex AI Agent Engine (Reasoning Engines)
|
||||
|
||||
Handles the transformation between LiteLLM's OpenAI-compatible format and
|
||||
Vertex AI Reasoning Engine's API format.
|
||||
|
||||
API Reference:
|
||||
- :query endpoint - for session management (create, get, list, delete)
|
||||
- :streamQuery endpoint - for actual queries (stream_query method)
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.vertex_ai.agent_engine.sse_iterator import (
|
||||
VertexAgentEngineResponseIterator,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
HTTPHandler = Any
|
||||
AsyncHTTPHandler = Any
|
||||
CustomStreamWrapper = Any
|
||||
|
||||
|
||||
class VertexAgentEngineError(BaseLLMException):
|
||||
"""Exception for Vertex Agent Engine errors."""
|
||||
|
||||
def __init__(self, status_code: int, message: str):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
super().__init__(message=message, status_code=status_code)
|
||||
|
||||
|
||||
class VertexAgentEngineConfig(BaseConfig, VertexBase):
|
||||
"""
|
||||
Configuration for Vertex AI Agent Engine (Reasoning Engines).
|
||||
|
||||
Model format: vertex_ai/agent_engine/<resource_id>
|
||||
Where resource_id is the numeric ID of the reasoning engine.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
BaseConfig.__init__(self, **kwargs)
|
||||
VertexBase.__init__(self)
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> List[str]:
|
||||
"""Vertex Agent Engine has limited OpenAI compatible params."""
|
||||
return ["user"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
"""Map OpenAI params to Agent Engine params."""
|
||||
# Map 'user' to 'user_id' for session management
|
||||
if "user" in non_default_params:
|
||||
optional_params["user_id"] = non_default_params["user"]
|
||||
return optional_params
|
||||
|
||||
def _parse_model_string(self, model: str) -> Tuple[str, str]:
|
||||
"""
|
||||
Parse model string to extract resource ID.
|
||||
|
||||
Model format: agent_engine/<project_number>/<location>/<engine_id>
|
||||
Or: agent_engine/<engine_id> (uses default project/location)
|
||||
|
||||
Returns: (resource_path, engine_id)
|
||||
"""
|
||||
# Remove 'agent_engine/' prefix if present
|
||||
if model.startswith("agent_engine/"):
|
||||
model = model[len("agent_engine/") :]
|
||||
|
||||
# Check if it's a full resource path
|
||||
if model.startswith("projects/"):
|
||||
# Full path: projects/123/locations/us-central1/reasoningEngines/456
|
||||
return model, model.split("/")[-1]
|
||||
|
||||
# Just the engine ID
|
||||
return model, model
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Get the complete URL for the request.
|
||||
|
||||
For Vertex Agent Engine:
|
||||
- Non-streaming: :query endpoint (for session management)
|
||||
- Streaming: :streamQuery endpoint (for actual queries)
|
||||
"""
|
||||
resource_path, engine_id = self._parse_model_string(model)
|
||||
|
||||
# Get project and location from litellm_params or environment
|
||||
vertex_project = self.safe_get_vertex_ai_project(litellm_params)
|
||||
vertex_location = self.safe_get_vertex_ai_location(litellm_params) or "us-central1"
|
||||
|
||||
# Build the full resource path if only engine_id was provided
|
||||
if not resource_path.startswith("projects/"):
|
||||
if not vertex_project:
|
||||
raise ValueError(
|
||||
"vertex_project is required for Vertex Agent Engine. "
|
||||
"Set via litellm_params['vertex_project'] or VERTEXAI_PROJECT env var."
|
||||
)
|
||||
resource_path = f"projects/{vertex_project}/locations/{vertex_location}/reasoningEngines/{engine_id}"
|
||||
|
||||
# Build the base URL
|
||||
base_url = f"https://{vertex_location}-aiplatform.googleapis.com"
|
||||
|
||||
# Always use :streamQuery endpoint for actual queries
|
||||
# The :query endpoint only supports session management methods
|
||||
# (create_session, get_session, list_sessions, delete_session, etc.)
|
||||
endpoint = f"{base_url}/v1beta1/{resource_path}:streamQuery"
|
||||
|
||||
verbose_logger.debug(f"Vertex Agent Engine URL: {endpoint}")
|
||||
return endpoint
|
||||
|
||||
def _get_auth_headers(
|
||||
self,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> Dict[str, str]:
|
||||
"""Get authentication headers using Google Cloud credentials."""
|
||||
vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params)
|
||||
vertex_project = self.safe_get_vertex_ai_project(litellm_params)
|
||||
|
||||
# Get access token using VertexBase
|
||||
access_token, project_id = self.get_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"Vertex Agent Engine: Authenticated for project {project_id}")
|
||||
|
||||
return {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def _get_user_id(self, optional_params: dict) -> str:
|
||||
"""Get or generate user ID for session management."""
|
||||
user_id = optional_params.get("user_id") or optional_params.get("user")
|
||||
if user_id:
|
||||
return user_id
|
||||
# Generate a user ID
|
||||
return f"litellm-user-{str(uuid.uuid4())[:8]}"
|
||||
|
||||
def _get_session_id(self, optional_params: dict) -> Optional[str]:
|
||||
"""Get session ID if provided."""
|
||||
return optional_params.get("session_id")
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform the request to Vertex Agent Engine format.
|
||||
|
||||
The API expects:
|
||||
{
|
||||
"class_method": "stream_query",
|
||||
"input": {
|
||||
"message": "...",
|
||||
"user_id": "...",
|
||||
"session_id": "..." (optional)
|
||||
}
|
||||
}
|
||||
"""
|
||||
# Use the last message content as the prompt
|
||||
prompt = convert_content_list_to_str(messages[-1])
|
||||
|
||||
# Get user_id and session_id
|
||||
user_id = self._get_user_id(optional_params)
|
||||
session_id = self._get_session_id(optional_params)
|
||||
|
||||
# Build the input
|
||||
input_data: Dict[str, Any] = {
|
||||
"message": prompt,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
if session_id:
|
||||
input_data["session_id"] = session_id
|
||||
|
||||
# Build the request payload
|
||||
# Note: stream_query is used for both streaming and non-streaming
|
||||
# The difference is the endpoint (:streamQuery vs :query)
|
||||
payload = {
|
||||
"class_method": "stream_query",
|
||||
"input": input_data,
|
||||
}
|
||||
|
||||
verbose_logger.debug(f"Vertex Agent Engine payload: {payload}")
|
||||
return payload
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Validate environment and set up authentication headers."""
|
||||
auth_headers = self._get_auth_headers(optional_params, litellm_params)
|
||||
headers.update(auth_headers)
|
||||
return headers
|
||||
|
||||
def _extract_text_from_response(self, response_data: dict) -> str:
|
||||
"""Extract text content from the response."""
|
||||
# Try to get from content.parts
|
||||
content = response_data.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
for part in parts:
|
||||
if "text" in part:
|
||||
return part["text"]
|
||||
|
||||
# Try actions.state_delta
|
||||
actions = response_data.get("actions", {})
|
||||
state_delta = actions.get("state_delta", {})
|
||||
for key, value in state_delta.items():
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
|
||||
return ""
|
||||
|
||||
def _calculate_usage(
|
||||
self, model: str, messages: List[AllMessageValues], content: str
|
||||
) -> Optional[Usage]:
|
||||
"""Calculate token usage using LiteLLM's token counter."""
|
||||
try:
|
||||
from litellm.utils import token_counter
|
||||
|
||||
prompt_tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
|
||||
completion_tokens = token_counter(
|
||||
model="gpt-3.5-turbo", text=content, count_response_tokens=True
|
||||
)
|
||||
total_tokens = prompt_tokens + completion_tokens
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to calculate token usage: {str(e)}")
|
||||
return None
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ModelResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
request_data: dict,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
api_key: Optional[str] = None,
|
||||
json_mode: Optional[bool] = None,
|
||||
) -> ModelResponse:
|
||||
"""
|
||||
Transform Vertex Agent Engine response to LiteLLM ModelResponse format.
|
||||
|
||||
The response is a streaming SSE format even for non-streaming requests.
|
||||
We need to collect all the chunks and extract the final response.
|
||||
"""
|
||||
try:
|
||||
content_type = raw_response.headers.get("content-type", "").lower()
|
||||
verbose_logger.debug(f"Vertex Agent Engine response Content-Type: {content_type}")
|
||||
|
||||
# Parse the SSE response
|
||||
response_text = raw_response.text
|
||||
verbose_logger.debug(f"Response (first 500 chars): {response_text[:500]}")
|
||||
|
||||
# Extract content from SSE stream
|
||||
content = ""
|
||||
for line in response_text.strip().split("\n"):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
data = json.loads(line)
|
||||
if isinstance(data, dict):
|
||||
text = self._extract_text_from_response(data)
|
||||
if text:
|
||||
content = text # Use the last non-empty text
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# Create the message
|
||||
message = Message(content=content, role="assistant")
|
||||
|
||||
# Create choices
|
||||
choice = Choices(finish_reason="stop", index=0, message=message)
|
||||
|
||||
# Update model response
|
||||
model_response.choices = [choice]
|
||||
model_response.model = model
|
||||
|
||||
# Calculate usage
|
||||
calculated_usage = self._calculate_usage(model, messages, content)
|
||||
if calculated_usage:
|
||||
setattr(model_response, "usage", calculated_usage)
|
||||
|
||||
return model_response
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Error processing Vertex Agent Engine response: {str(e)}")
|
||||
raise VertexAgentEngineError(
|
||||
message=f"Error processing response: {str(e)}",
|
||||
status_code=raw_response.status_code,
|
||||
)
|
||||
|
||||
def get_streaming_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
) -> VertexAgentEngineResponseIterator:
|
||||
"""Return a streaming iterator for SSE responses."""
|
||||
return VertexAgentEngineResponseIterator(
|
||||
streaming_response=raw_response.iter_lines(),
|
||||
sync_stream=True,
|
||||
)
|
||||
|
||||
def get_sync_custom_stream_wrapper(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_base: str,
|
||||
headers: dict,
|
||||
data: dict,
|
||||
messages: list,
|
||||
client: Optional[Union[HTTPHandler, "AsyncHTTPHandler"]] = None,
|
||||
json_mode: Optional[bool] = None,
|
||||
signed_json_body: Optional[bytes] = None,
|
||||
) -> "CustomStreamWrapper":
|
||||
"""Get a CustomStreamWrapper for synchronous streaming."""
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
HTTPHandler,
|
||||
_get_httpx_client,
|
||||
)
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
client = _get_httpx_client(params={})
|
||||
|
||||
# Avoid logging sensitive api_base directly
|
||||
verbose_logger.debug("Making sync streaming request to Vertex AI endpoint.")
|
||||
|
||||
# Make streaming request
|
||||
response = client.post(
|
||||
api_base,
|
||||
headers=headers,
|
||||
data=json.dumps(data),
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise VertexAgentEngineError(
|
||||
status_code=response.status_code, message=str(response.read())
|
||||
)
|
||||
|
||||
# Create iterator for SSE stream
|
||||
completion_stream = self.get_streaming_response(model=model, raw_response=response)
|
||||
|
||||
streaming_response = CustomStreamWrapper(
|
||||
completion_stream=completion_stream,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
original_response="first stream response received",
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
|
||||
return streaming_response
|
||||
|
||||
async def get_async_custom_stream_wrapper(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_base: str,
|
||||
headers: dict,
|
||||
data: dict,
|
||||
messages: list,
|
||||
client: Optional["AsyncHTTPHandler"] = None,
|
||||
json_mode: Optional[bool] = None,
|
||||
signed_json_body: Optional[bytes] = None,
|
||||
) -> "CustomStreamWrapper":
|
||||
"""Get a CustomStreamWrapper for asynchronous streaming."""
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
if client is None or not isinstance(client, AsyncHTTPHandler):
|
||||
client = get_async_httpx_client(
|
||||
llm_provider=cast(Any, "vertex_ai"), params={}
|
||||
)
|
||||
|
||||
# Avoid logging sensitive api_base directly
|
||||
verbose_logger.debug("Making async streaming request to Vertex AI endpoint.")
|
||||
|
||||
# Make async streaming request
|
||||
response = await client.post(
|
||||
api_base,
|
||||
headers=headers,
|
||||
data=json.dumps(data),
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise VertexAgentEngineError(
|
||||
status_code=response.status_code, message=str(await response.aread())
|
||||
)
|
||||
|
||||
# Create iterator for SSE stream (async)
|
||||
completion_stream = VertexAgentEngineResponseIterator(
|
||||
streaming_response=response.aiter_lines(),
|
||||
sync_stream=False,
|
||||
)
|
||||
|
||||
streaming_response = CustomStreamWrapper(
|
||||
completion_stream=completion_stream,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
original_response="first stream response received",
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
|
||||
return streaming_response
|
||||
|
||||
@property
|
||||
def has_custom_stream_wrapper(self) -> bool:
|
||||
"""Indicates that this config has custom streaming support."""
|
||||
return True
|
||||
|
||||
@property
|
||||
def supports_stream_param_in_request_body(self) -> bool:
|
||||
"""Agent Engine does not allow passing `stream` in the request body."""
|
||||
return False
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
||||
) -> BaseLLMException:
|
||||
return VertexAgentEngineError(status_code=status_code, message=error_message)
|
||||
|
||||
def should_fake_stream(
|
||||
self,
|
||||
model: Optional[str],
|
||||
stream: Optional[bool],
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Agent Engine always returns SSE streams, so we use real streaming."""
|
||||
return False
|
||||
|
||||
@ -5,7 +5,6 @@ from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, get_ty
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.utils import supports_response_schema, supports_system_messages
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs
|
||||
@ -14,6 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.vertex_ai import PartType, Schema
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
from litellm.utils import supports_response_schema, supports_system_messages
|
||||
|
||||
|
||||
class VertexAIError(BaseLLMException):
|
||||
@ -36,6 +36,7 @@ class VertexAIModelRoute(str, Enum):
|
||||
MODEL_GARDEN = "model_garden"
|
||||
NON_GEMINI = "non_gemini"
|
||||
OPENAI_COMPATIBLE = "openai"
|
||||
AGENT_ENGINE = "agent_engine"
|
||||
|
||||
VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute]
|
||||
|
||||
@ -76,6 +77,10 @@ def get_vertex_ai_model_route(
|
||||
if litellm_params and litellm_params.get("base_model") is not None:
|
||||
if "gemini" in litellm_params["base_model"]:
|
||||
return VertexAIModelRoute.GEMINI
|
||||
|
||||
# Check for agent_engine models (Reasoning Engines)
|
||||
if "agent_engine/" in model:
|
||||
return VertexAIModelRoute.AGENT_ENGINE
|
||||
|
||||
# Check if numeric endpoint ID with custom api_base (PSC endpoint)
|
||||
# Route to GEMINI (HTTP path) to support PSC endpoints properly
|
||||
|
||||
@ -3242,6 +3242,37 @@ def completion( # type: ignore # noqa: PLR0915
|
||||
timeout=timeout,
|
||||
client=client,
|
||||
)
|
||||
elif model_route == VertexAIModelRoute.AGENT_ENGINE:
|
||||
# Vertex AI Agent Engine (Reasoning Engines)
|
||||
from litellm.llms.vertex_ai.agent_engine.transformation import (
|
||||
VertexAgentEngineConfig,
|
||||
)
|
||||
|
||||
vertex_agent_engine_config = VertexAgentEngineConfig()
|
||||
|
||||
# Update litellm_params with vertex credentials
|
||||
litellm_params["vertex_project"] = vertex_ai_project
|
||||
litellm_params["vertex_location"] = vertex_ai_location
|
||||
litellm_params["vertex_credentials"] = vertex_credentials
|
||||
|
||||
model_response = base_llm_http_handler.completion(
|
||||
model=model,
|
||||
stream=stream,
|
||||
messages=messages,
|
||||
model_response=model_response,
|
||||
optional_params=new_params,
|
||||
litellm_params=litellm_params, # type: ignore
|
||||
encoding=encoding,
|
||||
api_key=None,
|
||||
api_base=api_base,
|
||||
logging_obj=logging,
|
||||
acompletion=acompletion,
|
||||
timeout=timeout,
|
||||
client=client,
|
||||
custom_llm_provider="vertex_ai",
|
||||
provider_config=vertex_agent_engine_config,
|
||||
headers=headers or {},
|
||||
)
|
||||
else: # VertexAIModelRoute.NON_GEMINI
|
||||
model_response = vertex_ai_non_gemini.completion(
|
||||
model=model,
|
||||
|
||||
@ -30628,7 +30628,7 @@
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"mode": "embedding"
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": {
|
||||
"fireworks_ai/accounts/fireworks/models/": {
|
||||
"max_tokens": 40960,
|
||||
"max_input_tokens": 40960,
|
||||
"max_output_tokens": 40960,
|
||||
|
||||
@ -166,6 +166,29 @@
|
||||
"litellm_params_template": {
|
||||
"custom_llm_provider": "pydantic_ai_agents"
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent_type": "vertex_agent_engine",
|
||||
"agent_type_display_name": "Vertex AI Agent Engine",
|
||||
"description": "Connect to Google Cloud Vertex AI Reasoning Engines",
|
||||
"logo_url": "/ui/assets/logos/google.svg",
|
||||
"inherit_credentials_from_provider": "Vertex_AI",
|
||||
"model_template": "vertex_ai/agent_engine/{reasoning_engine_id}",
|
||||
"credential_fields": [
|
||||
{
|
||||
"key": "reasoning_engine_id",
|
||||
"label": "Reasoning Engine Resource ID",
|
||||
"placeholder": "projects/123456789/locations/us-central1/reasoningEngines/987654321",
|
||||
"tooltip": "The full resource ID of your Vertex AI Reasoning Engine. Find this in Google Cloud Console under Vertex AI > Agent Builder > Your Agent.",
|
||||
"required": true,
|
||||
"field_type": "text",
|
||||
"default_value": null,
|
||||
"include_in_litellm_params": false
|
||||
}
|
||||
],
|
||||
"litellm_params_template": {
|
||||
"custom_llm_provider": "vertex_ai"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@ -2689,8 +2689,8 @@
|
||||
"key": "vertex_credentials",
|
||||
"label": "Vertex Credentials",
|
||||
"placeholder": null,
|
||||
"tooltip": null,
|
||||
"required": true,
|
||||
"tooltip": "Optional - Upload your GCP service account JSON file. If not provided, uses default GCP credentials (ADC).",
|
||||
"required": false,
|
||||
"field_type": "upload",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
|
||||
151
tests/agent_tests/local_vertex_agent.py
Normal file
151
tests/agent_tests/local_vertex_agent.py
Normal file
@ -0,0 +1,151 @@
|
||||
"""
|
||||
Test script for Vertex AI Reasoning Engine.
|
||||
|
||||
This script demonstrates how to:
|
||||
1. Authenticate with Google Cloud
|
||||
2. Send queries to a Vertex AI Reasoning Engine using the :query endpoint
|
||||
|
||||
Usage:
|
||||
python local_vertex_agent.py
|
||||
|
||||
Requirements:
|
||||
pip install httpx google-auth
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from uuid import uuid4
|
||||
|
||||
from google.auth import default
|
||||
from google.auth.transport.requests import Request
|
||||
import httpx
|
||||
|
||||
# Configuration - update these for your agent
|
||||
PROJECT_ID = "gen-lang-client-0682925754" # Your GCP project ID
|
||||
LOCATION = "us-central1" # Your agent's location
|
||||
|
||||
# For Reasoning Engines, use just the numeric ID at the end
|
||||
REASONING_ENGINE_ID = "8263861224643493888"
|
||||
|
||||
# The project number from the resource name
|
||||
PROJECT_NUMBER = "1060139831167"
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main function to test Vertex AI Reasoning Engine."""
|
||||
|
||||
# Step 1: Authenticate with Google Cloud
|
||||
print("Step 1: Authenticating with Google Cloud...")
|
||||
credentials, project = default(scopes=['https://www.googleapis.com/auth/cloud-platform'])
|
||||
credentials.refresh(Request())
|
||||
print(f"Authenticated! Project: {project}")
|
||||
print(f"Token (first 20 chars): {credentials.token[:20]}...")
|
||||
|
||||
# Step 2: Build the endpoint URL
|
||||
base_url = f"https://{LOCATION}-aiplatform.googleapis.com"
|
||||
resource_path = f"projects/{PROJECT_NUMBER}/locations/{LOCATION}/reasoningEngines/{REASONING_ENGINE_ID}"
|
||||
|
||||
# The Reasoning Engine uses :query endpoint with specific format
|
||||
query_url = f"{base_url}/v1beta1/{resource_path}:query"
|
||||
stream_url = f"{base_url}/v1beta1/{resource_path}:streamQuery"
|
||||
|
||||
print(f"\nQuery URL: {query_url}")
|
||||
print(f"Stream URL: {stream_url}")
|
||||
|
||||
# Step 3: Create authenticated httpx client
|
||||
print("\nStep 2: Creating authenticated HTTP client...")
|
||||
client = httpx.AsyncClient(
|
||||
headers={
|
||||
"Authorization": f"Bearer {credentials.token}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=120.0,
|
||||
)
|
||||
|
||||
# Step 4: Build the query request (non-streaming)
|
||||
# Note: For non-streaming, we need to:
|
||||
# 1. Create a session
|
||||
# 2. Use the streaming endpoint with stream_query method
|
||||
# The :query endpoint only supports session management methods
|
||||
|
||||
user_id = f"test-user-{uuid4().hex[:8]}"
|
||||
|
||||
# First create a session
|
||||
create_session_request = {
|
||||
"class_method": "async_create_session",
|
||||
"input": {
|
||||
"user_id": user_id,
|
||||
}
|
||||
}
|
||||
|
||||
print(f"\nStep 3: Creating session...")
|
||||
print(f"User ID: {user_id}")
|
||||
|
||||
async with client:
|
||||
# Create session
|
||||
print(f"\nSending to: {query_url}")
|
||||
response = await client.post(query_url, json=create_session_request)
|
||||
print(f"Create session status: {response.status_code}")
|
||||
|
||||
if response.status_code == 200:
|
||||
session_data = response.json()
|
||||
print(f"Session created:\n{json.dumps(session_data, indent=2)}")
|
||||
|
||||
# Extract session_id from response
|
||||
session_id = session_data.get("output", {}).get("id") or session_data.get("output", {}).get("session_id")
|
||||
print(f"\nSession ID: {session_id}")
|
||||
|
||||
# Now send the actual query via streamQuery
|
||||
query_request = {
|
||||
"class_method": "stream_query",
|
||||
"input": {
|
||||
"message": "Hello! What can you do?",
|
||||
"user_id": user_id,
|
||||
"session_id": session_id,
|
||||
}
|
||||
}
|
||||
|
||||
print(f"\nStep 4: Sending query via streamQuery...")
|
||||
print(f"Request:\n{json.dumps(query_request, indent=2)}")
|
||||
|
||||
# Use streaming endpoint but collect full response
|
||||
async with client.stream("POST", stream_url, json=query_request) as stream_response:
|
||||
print(f"Query status: {stream_response.status_code}")
|
||||
|
||||
if stream_response.status_code == 200:
|
||||
print("\nResponse:")
|
||||
full_response = ""
|
||||
async for line in stream_response.aiter_lines():
|
||||
if line:
|
||||
full_response = line # Keep last line (full response)
|
||||
|
||||
# Parse and display
|
||||
try:
|
||||
data = json.loads(full_response)
|
||||
# Extract the text from the response
|
||||
content = data.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
for part in parts:
|
||||
if "text" in part:
|
||||
print(f"\nAgent response:\n{part['text']}")
|
||||
except:
|
||||
print(full_response)
|
||||
else:
|
||||
content = await stream_response.aread()
|
||||
print(f"Error: {content.decode()}")
|
||||
else:
|
||||
print(f"Error creating session: {response.text}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("Vertex AI Reasoning Engine Test Script")
|
||||
print("=" * 60)
|
||||
print(f"\nConfiguration:")
|
||||
print(f" PROJECT_ID: {PROJECT_ID}")
|
||||
print(f" PROJECT_NUMBER: {PROJECT_NUMBER}")
|
||||
print(f" LOCATION: {LOCATION}")
|
||||
print(f" REASONING_ENGINE_ID: {REASONING_ENGINE_ID}")
|
||||
print()
|
||||
|
||||
asyncio.run(main())
|
||||
@ -201,3 +201,79 @@ async def test_a2a_completion_bridge_bedrock_agentcore():
|
||||
|
||||
print(f"Received {len(chunks)} chunks from Bedrock AgentCore")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Vertex AI Agent Engine Tests
|
||||
# ============================================================
|
||||
|
||||
# Configuration - update these for your Vertex AI Reasoning Engine
|
||||
VERTEX_AGENT_RESOURCE_NAME = "projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_agent_engine_non_streaming():
|
||||
"""
|
||||
Test non-streaming request to Vertex AI Agent Engine via litellm.acompletion.
|
||||
|
||||
Uses the Reasoning Engine resource ID to call a hosted agent.
|
||||
"""
|
||||
|
||||
litellm._turn_on_debug()
|
||||
|
||||
# Call via litellm.acompletion with vertex_ai/agent_engine/ prefix
|
||||
response = await litellm.acompletion(
|
||||
model=f"vertex_ai/agent_engine/{VERTEX_AGENT_RESOURCE_NAME}",
|
||||
messages=[{"role": "user", "content": "Hello! What can you do?"}],
|
||||
stream=False,
|
||||
)
|
||||
|
||||
print(f"\n=== Vertex Agent Engine Non-Streaming Response ===")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Basic assertions
|
||||
assert response is not None
|
||||
assert hasattr(response, "choices")
|
||||
assert len(response.choices) > 0
|
||||
assert response.choices[0].message is not None
|
||||
assert response.choices[0].message.content is not None
|
||||
assert len(response.choices[0].message.content) > 0
|
||||
|
||||
print(f"Agent response: {response.choices[0].message.content[:200]}...")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_agent_engine_streaming():
|
||||
"""
|
||||
Test streaming request to Vertex AI Agent Engine via litellm.acompletion.
|
||||
|
||||
Uses the Reasoning Engine resource ID to call a hosted agent with streaming.
|
||||
"""
|
||||
#litellm._turn_on_debug()
|
||||
|
||||
# Call via litellm.acompletion with streaming
|
||||
response = await litellm.acompletion(
|
||||
model=f"vertex_ai/agent_engine/{VERTEX_AGENT_RESOURCE_NAME}",
|
||||
messages=[{"role": "user", "content": "Hello! What can you do?"}],
|
||||
stream=True,
|
||||
)
|
||||
|
||||
print(f"\n=== Vertex Agent Engine Streaming Response ===")
|
||||
|
||||
chunks = []
|
||||
full_content = ""
|
||||
async for chunk in response:
|
||||
print(f"Chunk: {chunk}")
|
||||
# chunks.append(chunk)
|
||||
# if hasattr(chunk, "choices") and len(chunk.choices) > 0:
|
||||
# delta = chunk.choices[0].delta
|
||||
# if hasattr(delta, "content") and delta.content:
|
||||
# full_content += delta.content
|
||||
# print(f"Chunk: {delta.content}", end="", flush=True)
|
||||
|
||||
# # print(f"\n\nReceived {len(chunks)} chunks")
|
||||
# print(f"Full content: {full_content[:200]}...")
|
||||
|
||||
# # Basic assertions
|
||||
# assert len(chunks) > 0
|
||||
# assert len(full_content) > 0
|
||||
|
||||
|
||||
128
tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py
Normal file
128
tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py
Normal file
@ -0,0 +1,128 @@
|
||||
"""
|
||||
Tests for Vertex AI Agent Engine transformation.
|
||||
|
||||
Tests the request transformation and streaming chunk parsing without making real API calls.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.llms.vertex_ai.agent_engine.sse_iterator import (
|
||||
VertexAgentEngineResponseIterator,
|
||||
)
|
||||
from litellm.llms.vertex_ai.agent_engine.transformation import VertexAgentEngineConfig
|
||||
|
||||
|
||||
class TestVertexAgentEngineTransformRequest:
|
||||
"""Tests for transform_request method."""
|
||||
|
||||
def test_transform_request_basic(self):
|
||||
"""
|
||||
Test that transform_request correctly formats messages into Vertex Agent Engine payload.
|
||||
"""
|
||||
config = VertexAgentEngineConfig()
|
||||
|
||||
messages = [{"role": "user", "content": "Hello, what can you do?"}]
|
||||
optional_params = {"user_id": "test-user-123"}
|
||||
litellm_params = {}
|
||||
|
||||
result = config.transform_request(
|
||||
model="agent_engine/123456789",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result["class_method"] == "stream_query"
|
||||
assert result["input"]["message"] == "Hello, what can you do?"
|
||||
assert result["input"]["user_id"] == "test-user-123"
|
||||
assert "session_id" not in result["input"]
|
||||
|
||||
def test_transform_request_with_session_id(self):
|
||||
"""
|
||||
Test that transform_request includes session_id when provided.
|
||||
"""
|
||||
config = VertexAgentEngineConfig()
|
||||
|
||||
messages = [{"role": "user", "content": "Follow up question"}]
|
||||
optional_params = {
|
||||
"user_id": "test-user-123",
|
||||
"session_id": "session-abc-456",
|
||||
}
|
||||
litellm_params = {}
|
||||
|
||||
result = config.transform_request(
|
||||
model="agent_engine/123456789",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result["class_method"] == "stream_query"
|
||||
assert result["input"]["message"] == "Follow up question"
|
||||
assert result["input"]["user_id"] == "test-user-123"
|
||||
assert result["input"]["session_id"] == "session-abc-456"
|
||||
|
||||
|
||||
class TestVertexAgentEngineChunkParser:
|
||||
"""Tests for the streaming chunk parser."""
|
||||
|
||||
def test_chunk_parser_with_text_content(self):
|
||||
"""
|
||||
Test that chunk_parser correctly extracts text from Vertex Agent Engine response format.
|
||||
"""
|
||||
iterator = VertexAgentEngineResponseIterator(
|
||||
streaming_response=iter([]),
|
||||
sync_stream=True,
|
||||
)
|
||||
|
||||
chunk = {
|
||||
"content": {
|
||||
"parts": [{"text": "Hello! I can help you with financial analysis."}],
|
||||
"role": "model",
|
||||
},
|
||||
"finish_reason": "STOP",
|
||||
"usage_metadata": {
|
||||
"prompt_token_count": 100,
|
||||
"candidates_token_count": 50,
|
||||
"total_token_count": 150,
|
||||
},
|
||||
}
|
||||
|
||||
result = iterator.chunk_parser(chunk)
|
||||
|
||||
assert result.choices[0].delta.content == "Hello! I can help you with financial analysis."
|
||||
assert result.choices[0].delta.role == "assistant"
|
||||
assert result.choices[0].finish_reason == "stop"
|
||||
assert result.usage["prompt_tokens"] == 100
|
||||
assert result.usage["completion_tokens"] == 50
|
||||
assert result.usage["total_tokens"] == 150
|
||||
|
||||
def test_chunk_parser_without_finish_reason(self):
|
||||
"""
|
||||
Test that chunk_parser handles chunks without finish_reason (intermediate chunks).
|
||||
"""
|
||||
iterator = VertexAgentEngineResponseIterator(
|
||||
streaming_response=iter([]),
|
||||
sync_stream=True,
|
||||
)
|
||||
|
||||
chunk = {
|
||||
"content": {
|
||||
"parts": [{"text": "Partial response..."}],
|
||||
"role": "model",
|
||||
},
|
||||
}
|
||||
|
||||
result = iterator.chunk_parser(chunk)
|
||||
|
||||
assert result.choices[0].delta.content == "Partial response..."
|
||||
assert result.choices[0].finish_reason is None
|
||||
assert result.usage is None
|
||||
|
||||
Loading…
Reference in New Issue
Block a user