Files
einstein/functions/call_function.py
T

63 lines
1.8 KiB
Python

from typing import Callable
from google.genai import types
from functions.get_file_content import (get_file_content,
schema_get_file_content)
from functions.get_files_info import get_files_info, schema_get_files_info
from functions.run_python_file import run_python_file, schema_run_python_file
from functions.write_file import schema_write_file, write_file
available_functions = types.Tool(
function_declarations=[
schema_get_file_content,
schema_get_files_info,
schema_run_python_file,
schema_write_file,
]
)
function_map: dict[str, Callable[..., str]] = {
"get_file_content": get_file_content,
"get_files_info": get_files_info,
"run_python_file": run_python_file,
"write_file": write_file,
}
def call_function(
function_call: types.FunctionCall, verbose: bool = False
) -> types.Content:
if verbose:
print(f"Calling function: {function_call.name}({function_call.args})")
else:
print(f"Calling function: {function_call.name}")
function_name = function_call.name or ""
function = function_map.get(function_name)
if function is None:
return types.Content(
role="tool",
parts=[
types.Part.from_function_response(
name=function_name,
response={"error": f"Unknown function: {function_name}"},
)
],
)
args = dict(function_call.args) if function_call.args else {}
args["working_directory"] = "./calculator"
function_result = function(**args)
return types.Content(
role="tool",
parts=[
types.Part.from_function_response(
name=function_name,
response={"result": function_result},
)
],
)