Custom tools

import asyncio
from typing import Any, Dict

from google.adk.tools.tool_context import ToolContext
from veadk import Agent, Runner


def calculator(
    a: float, b: float, operation: str, tool_context: ToolContext
) -> Dict[str, Any]:
    """A simple calculator that performs a basic arithmetic operation.

    Args:
        a (float): The first operand.
        b (float): The second operand.
        operation (str): One of "add", "subtract", "multiply", "divide".

    Returns:
        Dict[str, Any]: On success, contains "result", "operation", and
        "status" == "success". On failure (unsupported operation or division
        by zero), contains "status" == "error" and a "message".
    """
    if operation == "add":
        return {"result": a + b, "operation": "+", "status": "success"}
    if operation == "subtract":
        return {"result": a - b, "operation": "-", "status": "success"}
    if operation == "multiply":
        return {"result": a * b, "operation": "*", "status": "success"}
    if operation == "divide":
        if b == 0:
            return {"status": "error", "message": "Divisor cannot be zero."}
        return {"result": a / b, "operation": "/", "status": "success"}
    return {"status": "error", "message": f"Unsupported operation: {operation}"}


agent = Agent(
    name="computing_agent",
    model_name="doubao-seed-2-1-pro-260628",
    instruction="Use the `calculator` tool to perform the calculation the user asks for.",
    tools=[calculator],
)
runner = Runner(agent=agent)

response = asyncio.run(runner.run("Add 2 and 3"))
print(response)

import asyncio

from google.adk.tools.tool_context import ToolContext
from veadk import Agent, Runner


def message_checker(user_message: str, tool_context: ToolContext) -> str:
    """Check a user message and return it normalized.

    Args:
        user_message (str): The user message to check.

    Returns:
        str: The checked message.
    """
    call_count = tool_context.state.get("message_checker_calls", 0) + 1
    tool_context.state["message_checker_calls"] = call_count

    return f"Checked message: {user_message.upper()} (call {call_count})"


agent = Agent(
    name="context_agent",
    model_name="doubao-seed-2-1-pro-260628",
    instruction="Use message_checker to check the user message, then show the checked result.",
    tools=[message_checker],
)
runner = Runner(agent=agent)

response = asyncio.run(runner.run("Hello world!"))
print(response)

import asyncio
from typing import Any

from google.adk.events import Event
from google.adk.tools.long_running_tool import LongRunningFunctionTool
from google.genai.types import Content, FunctionCall, FunctionResponse, Part
from veadk import Agent, Runner

APP_NAME = "long_running_tool_app"
USER_ID = "long_running_tool_user"
SESSION_ID = "long_running_tool_session"


def big_data_processing(data_url: str) -> dict[str, Any]:
    """Start processing big data located at a URL.

    Args:
        data_url (str): The URL of the big data to process.

    Returns:
        dict[str, Any]: The initial task state, with "status" == "pending",
        the "data-url", and a "task-id" to track the job.
    """
    # Kick off the long-running job and return its handle immediately.
    return {
        "status": "pending",
        "data-url": data_url,
        "task-id": "big-data-processing-1",
    }


long_running_tool = LongRunningFunctionTool(func=big_data_processing)

agent = Agent(
    name="long_running_tool_agent",
    model_name="doubao-seed-2-1-pro-260628",
    instruction="Use long_running_tool to process big data.",
    tools=[long_running_tool],
)
runner = Runner(agent=agent, app_name=APP_NAME)


def get_long_running_call(event: Event) -> FunctionCall | None:
    """Return the long-running function call carried by an event, if any."""
    if not event.long_running_tool_ids or not event.content or not event.content.parts:
        return None
    for part in event.content.parts:
        if (
            part.function_call
            and part.function_call.id in event.long_running_tool_ids
        ):
            return part.function_call
    return None


def get_function_response(
    event: Event, function_call_id: str
) -> FunctionResponse | None:
    """Return the function response matching a given call id, if any."""
    if not event.content or not event.content.parts:
        return None
    for part in event.content.parts:
        if (
            part.function_response
            and part.function_response.id == function_call_id
        ):
            return part.function_response
    return None


async def main():
    # Create the session before running.
    session = await runner.short_term_memory.create_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
    )

    query = "Process the big data from https://example.com/data.csv"
    content = Content(role="user", parts=[Part(text=query)])

    print("Running agent...")
    long_running_call = None
    long_running_response = None
    async for event in runner.run_async(
        session_id=session.id, user_id=USER_ID, new_message=content
    ):
        if long_running_call is None:
            long_running_call = get_long_running_call(event)
        elif long_running_response is None:
            long_running_response = get_function_response(event, long_running_call.id)
        if event.content and event.content.parts:
            if text := "".join(part.text or "" for part in event.content.parts):
                print(f"[{event.author}]: {text}")

    # The task finished in the background; send the final result back to the agent.
    if long_running_response is not None:
        updated = long_running_response.model_copy(deep=True)
        updated.response = {"status": "finish"}
        async for event in runner.run_async(
            session_id=session.id,
            user_id=USER_ID,
            new_message=Content(
                role="user", parts=[Part(function_response=updated)]
            ),
        ):
            if event.content and event.content.parts:
                if text := "".join(part.text or "" for part in event.content.parts):
                    print(f"[{event.author}]: {text}")


if __name__ == "__main__":
    asyncio.run(main())