Build Tools for an Agent

@toolset.tool
async def search_context(
    ctx: RunContext[AgentDeps],
    query: str,
    scope: Literal["edges", "nodes", "episodes"] = "edges",
    node_labels: list[EntityLabel] | None = None,
    edge_types: list[EdgeTypeName] | None = None,
    report_type: str | None = None,
    product: str | None = None,
    near: list[str] | None = None,
    limit: int = 10,
) -> str:
    """Search the graph for the items most relevant to a query. Returns a ranked sample, not a complete set.
    Use scope="edges" for facts and relationships, scope="nodes" for entities, and scope="episodes" for the source reports. Use report_type and product to restrict the search to reports with that metadata. Use near with one or more handles, and scope="edges" or scope="nodes", to add results from the graph neighborhood of known nodes. The results can also include matches from outside the neighborhood."""
    meta_filters = [
        EpisodeMetadataFilter(comparison_operator="=", property_name=name, property_value=value)
        for name, value in (("report_type", report_type), ("product", product))
        if value
    ]
    results = await ctx.deps.zep.graph.search(
        graph_id=ctx.deps.graph_id,          # pinned
        query=query,
        scope=scope,
        limit=min(limit, SEARCH_LIMIT_MAX),  # pinned maximum
        reranker="cross_encoder",            # pinned
        bfs_origin_node_uuids=[ctx.deps.registry.resolve(h) for h in near] if near else None,
        search_filters=SearchFilters(
            node_labels=node_labels,
            edge_types=edge_types,
            episode_metadata_filters=MetadataFilterGroup(type="and", filters=meta_filters)
            if meta_filters else None,
        ),
    )
    return format_ranked_sample(ctx.deps, results)

@toolset.tool
async def list_nodes(
    ctx: RunContext[AgentDeps],
    label: EntityLabel,
    order_by: Literal["degree"] | None = None,
    limit: int = LIST_LIMIT_MAX,
) -> str:
    """List the nodes that have one entity type. Returns every match up to the limit, and says if the list is complete.
    Use this when the question needs a complete set, such as all open quality issues or all products. Set order_by="degree" to list the most connected nodes first."""
    page_size = min(limit, LIST_LIMIT_MAX)
    nodes = await ctx.deps.zep.graph.node.get_by_graph_id(
        ctx.deps.graph_id,
        filters=SearchFilters(node_labels=[label]),
        order_by=order_by,
        limit=page_size + 1,  # one extra node shows if the list is complete
    )
    return format_list(ctx.deps, nodes, page_size)

@toolset.tool
async def get_neighborhood(
    ctx: RunContext[AgentDeps],
    handle: str,
    edge_types: list[EdgeTypeName] | None = None,
    direction: Literal["out", "in", "both"] = "both",
    limit: int = 30,
) -> str:
    """Get the edges and neighbor nodes of one node. Returns every match up to the limit, and says if the list is complete.
    Use this to follow relationships from a node you already have, such as the components of a product, the supplier of a component, or the manager of an employee. Use edge_types to keep only some relationships. Use direction="out" for edges that start at the node, such as the manager of an employee (REPORTS_TO). Use direction="in" for edges that end at the node, such as the members of a team (MEMBER_OF)."""
    page_size = min(limit, NEIGHBOR_LIMIT_MAX)
    neighbors = await ctx.deps.zep.graph.node.get_neighbors(
        ctx.deps.registry.resolve(handle, prefix="n"),
        filters=SearchFilters(edge_types=edge_types) if edge_types else None,
        direction=direction,
        limit=page_size + 1,
    )
    return format_list(ctx.deps, neighbors, page_size)

@toolset.tool
async def get_details(ctx: RunContext[AgentDeps], handle: str) -> str:
    """Get the full attributes of a node (n handle) or the full text and metadata of an episode (p handle).
    Use this when a search result is relevant but too short to answer from, for example to read a full report."""
    if handle.startswith("p"):
        episode = await ctx.deps.zep.graph.episode.get(ctx.deps.registry.resolve(handle, prefix="p"))
        return format_episode(ctx.deps, episode, max_chars=EPISODE_TEXT_MAX_CHARS)
    node = await ctx.deps.zep.graph.node.get(ctx.deps.registry.resolve(handle, prefix="n"))
    return format_node_details(ctx.deps, node)

@toolset.tool
async def get_employees(
    ctx: RunContext[AgentDeps],
    team: str | None = None,
    product: str | None = None,
) -> str:
    """Find employees by team, by the product they work on, or both. Returns every match, with each person's title, team, and manager.
    Use this for questions about who works on something or who is on a team. Names must match the graph, such as "Reliability Engineering" or "Aster 410"."""
    if team is None and product is None:
        return "pass team, product, or both."
    groups = []
    if team:
        team_node = await node_by_name(ctx.deps, "Team", team)
        groups.append(await members(ctx.deps, team_node, "MEMBER_OF"))   # edges end at the team
    if product:
        product_node = await node_by_name(ctx.deps, "Product", product)
        groups.append(await members(ctx.deps, product_node, "WORKS_ON"))  # edges end at the product
    employees = intersect(groups)
    return await format_employees(ctx.deps, employees)  # one REPORTS_TO lookup (direction="out") for each employee

@toolset.tool
async def search_products(
    ctx: RunContext[AgentDeps],
    query: str | None = None,
    category: str | None = None,
) -> str:
    """Get products with their category, lifecycle status, and every regulatory filing (region, status, and key dates). Returns every product that matches, and says if the product list is complete.
    Use this for questions about which products are cleared, launched, or in a category. Read each filing status: only a cleared 510(k) or a valid CE certificate clears a product in a region."""
    products = await ctx.deps.zep.graph.node.get_by_graph_id(
        ctx.deps.graph_id, filters=SearchFilters(node_labels=["Product"]), limit=LIST_LIMIT_MAX + 1
    )
    truncated = len(products) > LIST_LIMIT_MAX
    products = filter_products(products[:LIST_LIMIT_MAX], query=query, category=category)
    return await format_products_with_filings(ctx.deps, products, truncated)  # FILED_FOR neighbors for each product