Add option to change the base URL for Google Generative AI. (#2113)

* Add option to change the base URL for Google Generative AI.
This would be useful to support private instance or gateway of the API

* fix: add proper type hints for http_options in Google LLM service
This commit is contained in:
otaqwawi
2025-07-04 01:12:35 +07:00
committed by GitHub
parent b184d62634
commit 6045a8ad8c

View File

@@ -67,6 +67,7 @@ try:
Content,
FunctionCall,
FunctionResponse,
HttpOptions,
GenerateContentConfig,
Part,
)
@@ -678,6 +679,7 @@ class GoogleLLMService(LLMService):
system_instruction: Optional[str] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_config: Optional[Dict[str, Any]] = None,
http_options: Optional[HttpOptions] = None,
**kwargs,
):
"""Initialize the Google LLM service.
@@ -689,6 +691,7 @@ class GoogleLLMService(LLMService):
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 parent class.
"""
super().__init__(**kwargs)
@@ -698,7 +701,8 @@ class GoogleLLMService(LLMService):
self.set_model_name(model)
self._api_key = api_key
self._system_instruction = system_instruction
self._create_client(api_key)
self._http_options = http_options
self._create_client(api_key, http_options)
self._settings = {
"max_tokens": params.max_tokens,
"temperature": params.temperature,
@@ -717,6 +721,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 needs_mcp_alternate_schema(self) -> bool:
"""Check if this LLM service requires alternate MCP schema.
@@ -728,9 +735,6 @@ class GoogleLLMService(LLMService):
"""
return True
def _create_client(self, api_key: str):
self._client = genai.Client(api_key=api_key)
def _maybe_unset_thinking_budget(self, generation_params: Dict[str, Any]):
try:
# There's no way to introspect on model capabilities, so