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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user