agent/agent_tools.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "iconclass",
#     "sentence-transformers",
#     "numpy",
# ]
# ///
"""Improved Iconclass tools for the interactive tool-agent (NEXT_SESSION Step 0).

The old tools in `../grpo-tools/iconclass_tools.py` were built for the dead GRPO
path and have three problems as *agent* tools:
  1. `search_iconclass` is hardwired to `all-MiniLM-L6-v2` indexing each code by
     its LEAF text only -> siblings collapse, ~18% exact-leaf. We swap the backend
     to the fine-tuned **`davanstrien/iconclass-retriever-bge-ft`** (promptless,
     `leaf` doc-mode) -- the retriever-FT component win (exact-leaf 18.2 -> 28.2%).
  2. `lookup_iconclass_code` returns parent/child/related as BARE CODES, so the
     agent can't reason about leaf-vs-sibling. We decode the TEXT of parents,
     children, and *siblings* (the disambiguation set).
  3. `with_timeout` uses `signal.SIGALRM` (main-thread only) -> breaks under
     smolagents' threaded tool execution. Dropped entirely; these calls are
     in-process and fast.

These are PLAIN typed functions with Google-style docstrings (no smolagents
dependency here, so they're trivially testable). `agent_iconclass.py` wraps them
with `smolagents.tool()` to expose them to the agent.

The FT retriever reuses `ret_index.RetrieverIndex`, whose 40K-tree doc embeddings
are cached under `ret_cache/` -- a live query is one encode + matmul (fast).

USAGE
    uv run agent_tools.py        # self-test the tools on a few examples
"""

from __future__ import annotations

import os
import sys
from functools import lru_cache

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from iconclass import init as _ic_init  # noqa: E402

_ic = _ic_init()

# --------------------------------------------------------------------------- FT retriever backend
FT_RETRIEVER = "davanstrien/iconclass-retriever-bge-ft"
_INDEX = None


def _get_index():
    """Lazily build the FT-retriever index (cached doc embeddings -> instant)."""
    global _INDEX
    if _INDEX is None:
        from ret_index import RetrieverIndex

        _INDEX = RetrieverIndex(model_name=FT_RETRIEVER, doc_mode="leaf")
    return _INDEX


# --------------------------------------------------------------------------- text helpers
def _base(code):
    """Strip any (KEY)/(+xx) suffix to get the structural base code."""
    return code.split("(")[0].strip() if isinstance(code, str) else ""


@lru_cache(maxsize=200_000)
def _text(code):
    """Decoded Iconclass text for a code (base code; falls back to the code)."""
    try:
        b = _base(code)
        n = _ic[b]
        return (n() if callable(n) else str(n)) or b
    except Exception:
        return _base(code) or code


# =====================================================================
# search_iconclass — FT bge retriever (the component win)
# =====================================================================
def search_iconclass(query: str, limit: int = 8) -> list[dict]:
    """Search the Iconclass tree for codes matching a text description.

    Backed by a fine-tuned bge retriever specialised for picking the correct
    leaf code from an art-historical description. Use this whenever you can
    describe what is depicted (a figure, object, scene, pose, or theme) but do
    not know the exact code.

    Args:
        query: A concise text description to search for, e.g. "the Crucifixion",
            "head turned to the left", "a kneeling man", "grapes".
        limit: Maximum number of candidate codes to return (1-24, default 8).

    Returns:
        A list of dicts, each with: code (the Iconclass code), description (its
        decoded text), and similarity (retrieval score, higher = closer).
    """
    if not query or not isinstance(query, str) or not query.strip():
        return []
    limit = max(1, min(int(limit), 24))
    hits = _get_index().search(query.strip(), k=limit)
    return [
        {"code": code, "description": leaf, "similarity": round(float(sim), 4)}
        for code, leaf, sim in hits
    ]


# =====================================================================
# lookup_iconclass_code — enriched with decoded sibling/parent/child text
# =====================================================================
def lookup_iconclass_code(code: str) -> dict:
    """Look up an Iconclass code: its meaning, ancestors, children, and siblings.

    Returns the decoded TEXT of related codes so you can disambiguate a leaf from
    its siblings (e.g. "head turned to the left" 31A2421 vs "...to the right"
    31A2422). Use this to verify a candidate against the image and to navigate to
    the correct sibling.

    Args:
        code: The Iconclass code to look up, e.g. "31A2421", "11H(FRANCIS)",
            "25F23". Any (KEY)/(+xx) suffix is handled.

    Returns:
        A dict with: valid (bool), code, description (its text), parents
        (list of {code, text} from root to here), children (list of {code, text}
        one level deeper), siblings (list of {code, text} sharing the same
        parent -- the leaf-vs-sibling disambiguation set). children/siblings are
        capped for brevity; n_children/n_siblings give the true counts.
    """
    b = _base(code)
    if not b:
        return {
            "valid": False,
            "code": code,
            "description": "Empty code provided",
            "parents": [],
            "children": [],
            "siblings": [],
        }
    try:
        node = _ic[b]
    except Exception:
        return {
            "valid": False,
            "code": code,
            "description": f"Code '{code}' not found in Iconclass",
            "parents": [],
            "children": [],
            "siblings": [],
        }

    obj = node.obj
    desc = node() if callable(node) else str(node)
    path = obj.get("p", []) or []  # root -> this code (inclusive)

    parents = [{"code": p, "text": _text(p)} for p in path if p != b]

    children_codes = obj.get("c", []) or []
    children = [{"code": c, "text": _text(c)} for c in children_codes[:10]]

    # siblings = the parent's children minus self (the disambiguation set)
    sibling_codes = []
    if len(path) >= 2:
        parent = path[-2]
        try:
            sibling_codes = [c for c in (_ic[parent].obj.get("c", []) or []) if c != b]
        except Exception:
            sibling_codes = []
    siblings = [{"code": c, "text": _text(c)} for c in sibling_codes[:10]]

    return {
        "valid": True,
        "code": b,
        "description": desc,
        "parents": parents,
        "children": children,
        "n_children": len(children_codes),
        "siblings": siblings,
        "n_siblings": len(sibling_codes),
    }


# =====================================================================
# get_parent_code — navigate up the hierarchy
# =====================================================================
def get_parent_code(code: str) -> dict:
    """Get the immediate parent (one level more general) of an Iconclass code.

    Use this to move up the hierarchy, e.g. to back off from an over-specific
    leaf to a more defensible general code.

    Args:
        code: The Iconclass code to get the parent of, e.g. "25F23".

    Returns:
        A dict with: valid (bool), parent_code (or None at root), and
        parent_description (the parent's decoded text).
    """
    b = _base(code)
    if not b:
        return {"valid": False, "parent_code": None, "parent_description": "Empty code"}
    try:
        path = _ic[b].obj.get("p", []) or []
    except Exception:
        return {
            "valid": False,
            "parent_code": None,
            "parent_description": f"Code '{code}' not found",
        }
    if len(path) <= 1:
        return {
            "valid": True,
            "parent_code": None,
            "parent_description": "At root level",
        }
    parent = path[-2]
    return {"valid": True, "parent_code": parent, "parent_description": _text(parent)}


# =====================================================================
# validate_iconclass_code — cheap existence check
# =====================================================================
def validate_iconclass_code(code: str) -> bool:
    """Check whether an Iconclass code is structurally valid (exists in the tree).

    Use this to confirm a code before committing it.

    Args:
        code: The Iconclass code to validate, e.g. "25F23", "41A".

    Returns:
        True if the code (or a valid prefix of it) exists, False otherwise.
    """
    b = _base(code)
    if not b:
        return False
    # Strict: every ':' segment of the base must exist in the tree at full depth.
    # The iconclass DB is complete, so any valid code exists. The old prefix-
    # fallback (walking down to the 1-char root, e.g. "4" = Society exists) let
    # wholesale hallucinated families like 4C11211…4C1121100 pass as "valid".
    for seg in b.split(":"):
        seg = seg.strip()
        if not seg:
            return False
        try:
            _ = _ic[seg]
        except Exception:
            return False
    return True


# =====================================================================
# define_codes — batch code -> text (token economy)
# =====================================================================
def define_codes(codes: list[str]) -> dict:
    """Decode several Iconclass codes to their text meanings in one call.

    Token-efficient way to get definitions for a whole candidate set at once,
    instead of looking codes up individually.

    Args:
        codes: A list of Iconclass codes to define, e.g. ["31A2421", "25F23"].

    Returns:
        A dict mapping each input code to its decoded text (or the code itself
        if it cannot be decoded).
    """
    if not codes:
        return {}
    out = {}
    for c in codes:
        if isinstance(c, str) and c.strip():
            out[c] = _text(c)
    return out


# =====================================================================
# Self-test
# =====================================================================
if __name__ == "__main__":
    print("=== agent_tools self-test ===\n")

    print("1. search_iconclass — FT retriever (sibling discrimination):")
    for q in [
        "the Crucifixion",
        "head turned to the left",
        "portrait of a woman",
        "grapes",
    ]:
        print(f"\n  query: {q!r}")
        for r in search_iconclass(q, limit=5):
            d = r["description"][:54]
            print(f"    {r['code']:<16} {d:<56} sim={r['similarity']}")

    print("\n2. lookup_iconclass_code('31A2421')  (head turned to the left):")
    res = lookup_iconclass_code("31A2421")
    print(f"   valid={res['valid']}  desc={res['description']!r}")
    print(f"   parents: {[(p['code'], p['text'][:30]) for p in res['parents']]}")
    print(f"   siblings (n={res.get('n_siblings')}):")
    for s in res["siblings"][:6]:
        print(f"     {s['code']:<16} {s['text'][:50]}")

    print("\n3. get_parent_code('25F23'):")
    print(f"   {get_parent_code('25F23')}")

    print("\n4. validate_iconclass_code:")
    print(f"   '25F23' -> {validate_iconclass_code('25F23')}")
    print(f"   '11H(FRANCIS)' -> {validate_iconclass_code('11H(FRANCIS)')}")
    print(f"   'ZZZZ' -> {validate_iconclass_code('ZZZZ')}")

    print("\n5. define_codes(['31A2421', '31A2422', '25F23']):")
    for k, v in define_codes(["31A2421", "31A2422", "25F23"]).items():
        print(f"   {k:<12} {v[:50]}")

    print("\nOK")