parallel-function-calls.py

PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
LOCATION = "us-central1"  # @param {type:"string"}

import vertexai
vertexai.init(project=PROJECT_ID, location=LOCATION)

from vertexai.generative_models import (
    FunctionDeclaration,
    GenerationConfig,
    GenerationResponse,
    GenerativeModel,
    Part,
    Tool,
)

list_tables_func = FunctionDeclaration(
    name="list_tables",
    description="List tables in a dataset that will help answer the user's question",
    parameters={
        "type": "object",
        "properties": {
            "dataset_id": {
                "type": "string",
                "description": "Dataset ID to fetch tables from.",
            }
        },
        "required": [
            "dataset_id",
        ],
    },
)

tool = Tool(
    function_declarations=[
        list_tables_func,
    ],
)

model = GenerativeModel(
    "gemini-1.5-pro",
    generation_config=GenerationConfig(temperature=0),
    tools=[tool],
)
chat = model.start_chat()


prompt = "List all of the tables from the datasets named: solar and cloudnext"

response = chat.send_message(prompt)

print("FUNCTION CALLS\n===========\n")
print(response.candidates[0].function_calls)

api_response = ["solar1, solar2, solar3", "cloudnext1, cloudnext2, cloudnext3"]

response = chat.send_message(
    [
        Part.from_function_response(
            name="list_tables",
            response={
                "content": api_response[0],
            },
        ),
        Part.from_function_response(
            name="list_tables",
            response={
                "content": api_response[1],
            },
        ),
    ],
)

print("\nLLM RESPONSE\n===========\n")
print(response.candidates[0].function_calls)
print(response.text)