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,
|
Content,
|
||||||
FunctionCall,
|
FunctionCall,
|
||||||
FunctionResponse,
|
FunctionResponse,
|
||||||
|
HttpOptions,
|
||||||
GenerateContentConfig,
|
GenerateContentConfig,
|
||||||
Part,
|
Part,
|
||||||
)
|
)
|
||||||
@@ -678,6 +679,7 @@ class GoogleLLMService(LLMService):
|
|||||||
system_instruction: Optional[str] = None,
|
system_instruction: Optional[str] = None,
|
||||||
tools: Optional[List[Dict[str, Any]]] = None,
|
tools: Optional[List[Dict[str, Any]]] = None,
|
||||||
tool_config: Optional[Dict[str, Any]] = None,
|
tool_config: Optional[Dict[str, Any]] = None,
|
||||||
|
http_options: Optional[HttpOptions] = None,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
):
|
):
|
||||||
"""Initialize the Google LLM service.
|
"""Initialize the Google LLM service.
|
||||||
@@ -689,6 +691,7 @@ class GoogleLLMService(LLMService):
|
|||||||
system_instruction: System instruction/prompt for the model.
|
system_instruction: System instruction/prompt for the model.
|
||||||
tools: List of available tools/functions.
|
tools: List of available tools/functions.
|
||||||
tool_config: Configuration for tool usage.
|
tool_config: Configuration for tool usage.
|
||||||
|
http_options: HTTP options for the client.
|
||||||
**kwargs: Additional arguments passed to parent class.
|
**kwargs: Additional arguments passed to parent class.
|
||||||
"""
|
"""
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
@@ -698,7 +701,8 @@ class GoogleLLMService(LLMService):
|
|||||||
self.set_model_name(model)
|
self.set_model_name(model)
|
||||||
self._api_key = api_key
|
self._api_key = api_key
|
||||||
self._system_instruction = system_instruction
|
self._system_instruction = system_instruction
|
||||||
self._create_client(api_key)
|
self._http_options = http_options
|
||||||
|
self._create_client(api_key, http_options)
|
||||||
self._settings = {
|
self._settings = {
|
||||||
"max_tokens": params.max_tokens,
|
"max_tokens": params.max_tokens,
|
||||||
"temperature": params.temperature,
|
"temperature": params.temperature,
|
||||||
@@ -717,6 +721,9 @@ class GoogleLLMService(LLMService):
|
|||||||
"""
|
"""
|
||||||
return True
|
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:
|
def needs_mcp_alternate_schema(self) -> bool:
|
||||||
"""Check if this LLM service requires alternate MCP schema.
|
"""Check if this LLM service requires alternate MCP schema.
|
||||||
|
|
||||||
@@ -728,9 +735,6 @@ class GoogleLLMService(LLMService):
|
|||||||
"""
|
"""
|
||||||
return True
|
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]):
|
def _maybe_unset_thinking_budget(self, generation_params: Dict[str, Any]):
|
||||||
try:
|
try:
|
||||||
# There's no way to introspect on model capabilities, so
|
# There's no way to introspect on model capabilities, so
|
||||||
|
|||||||
Reference in New Issue
Block a user