wehub-resource-sync
113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
import json
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
from mlflow.types.chat import (
|
|
ChatTool,
|
|
Function,
|
|
FunctionParams,
|
|
FunctionToolDefinition,
|
|
ParamProperty,
|
|
ToolCall,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from google import genai
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def convert_gemini_func_to_mlflow_chat_tool(
|
|
function_def: "genai.types.FunctionDeclaration",
|
|
) -> ChatTool:
|
|
"""
|
|
Convert Gemini function definition into MLflow's standard format (OpenAI compatible).
|
|
Ref: https://ai.google.dev/gemini-api/docs/function-calling
|
|
|
|
Args:
|
|
function_def: A genai.types.FunctionDeclaration or genai.protos.FunctionDeclaration object
|
|
representing a function definition.
|
|
|
|
Returns:
|
|
ChatTool: MLflow's standard tool definition object.
|
|
"""
|
|
return ChatTool(
|
|
type="function",
|
|
function=FunctionToolDefinition(
|
|
name=function_def.name,
|
|
description=function_def.description,
|
|
parameters=_convert_gemini_function_param_to_mlflow_function_param(
|
|
function_def.parameters
|
|
),
|
|
),
|
|
)
|
|
|
|
|
|
def convert_gemini_func_call_to_mlflow_tool_call(
|
|
func_call: "genai.types.FunctionCall",
|
|
) -> ToolCall:
|
|
"""
|
|
Convert Gemini function call into MLflow's standard format (OpenAI compatible).
|
|
Ref: https://ai.google.dev/gemini-api/docs/function-calling
|
|
|
|
Args:
|
|
func_call: A genai.types.FunctionCall or genai.protos.FunctionCall object
|
|
representing a single func call.
|
|
|
|
Returns:
|
|
ToolCall: MLflow's standard tool call object.
|
|
"""
|
|
# original args object is not json serializable
|
|
args = func_call.args or {}
|
|
|
|
return ToolCall(
|
|
# Gemini does not have func call id
|
|
id=func_call.name,
|
|
type="function",
|
|
function=Function(name=func_call.name, arguments=json.dumps(dict(args))),
|
|
)
|
|
|
|
|
|
def _convert_gemini_param_property_to_mlflow_param_property(param_property) -> ParamProperty:
|
|
"""
|
|
Convert Gemini parameter property definition into MLflow's standard format (OpenAI compatible).
|
|
Ref: https://ai.google.dev/gemini-api/docs/function-calling
|
|
|
|
Args:
|
|
param_property: A genai.types.Schema or genai.protos.Schema object
|
|
representing a parameter property.
|
|
|
|
Returns:
|
|
ParamProperty: MLflow's standard param property object.
|
|
"""
|
|
type_name = param_property.type
|
|
type_name = type_name.name.lower() if hasattr(type_name, "name") else type_name.lower()
|
|
return ParamProperty(
|
|
description=param_property.description,
|
|
enum=param_property.enum,
|
|
type=type_name,
|
|
)
|
|
|
|
|
|
def _convert_gemini_function_param_to_mlflow_function_param(
|
|
function_params: "genai.types.Schema",
|
|
) -> FunctionParams:
|
|
"""
|
|
Convert Gemini function parameter definition into MLflow's standard format (OpenAI compatible).
|
|
Ref: https://ai.google.dev/gemini-api/docs/function-calling
|
|
|
|
Args:
|
|
function_params: A genai.types.Schema or genai.protos.Schema object
|
|
representing function parameters.
|
|
|
|
Returns:
|
|
FunctionParams: MLflow's standard function parameter object.
|
|
"""
|
|
return FunctionParams(
|
|
properties={
|
|
k: _convert_gemini_param_property_to_mlflow_param_property(v)
|
|
for k, v in function_params.properties.items()
|
|
},
|
|
required=function_params.required,
|
|
)
|