Agent with Neo4j Logging¶

from __future__ import annotations

from typing import Dict

from google.adk.models.lite_llm import LiteLlm
from google.adk.tools.tool_context import ToolContext

from opensage.agents.opensage_agent import OpenSageAgent
from opensage.toolbox.general.orchestration_tools import (
    call_subagent,
    create_subagent,
    list_subagents,
)


def add_numbers(a: float, b: float) -> float:
    """Add two numbers together."""
    return a + b


def multiply_numbers(a: float, b: float) -> float:
    """Multiply two numbers together."""
    return a * b


def calculate_area_and_perimeter(
    length: float, width: float, tool_context: ToolContext
) -> Dict[str, float]:
    """Calculate area and perimeter of a rectangle."""
    return {"area": length * width, "perimeter": 2 * (length + width)}


geometry_calculator = OpenSageAgent(
    name="geometry_calculator",
    description="Calculates geometric properties like area and perimeter of shapes.",
    model=LiteLlm(model="openai/gpt-5.5"),
    instruction="You specialize in calculating geometric properties.",
    tools=[calculate_area_and_perimeter],
)


def mk_agent(opensage_session_id: str):
    return OpenSageAgent(
        name="calculation_orchestrator",
        description="Main agent that coordinates calculations with Neo4j history logging.",
        model=LiteLlm(model="openai/gpt-5.5"),
        instruction="""
        You are a calculation orchestrator. Delegate geometric calculations to
        the `geometry_calculator` sub-agent via the `call_subagent` tool.
        Formulate the final answer as a single number inside
        <final_answer>...</final_answer> tags.
        """,
        subagents=[geometry_calculator],
        tools=[
            call_subagent,
            create_subagent,
            list_subagents,
            add_numbers,
            multiply_numbers,
        ],
    )