grpo-tools/iconclass_tools.py

"""
Iconclass tools for GRPO agent training.

These tools allow the model to validate, look up, and search iconclass codes
during generation rather than relying on post-hoc reward validation.
"""

import signal
from functools import lru_cache
from typing import Any

from iconclass import NotationNotFound, init

# Initialize iconclass once at module level
ic = init()


class TimeoutError(Exception):
    """Custom timeout exception for tool execution."""

    pass


def timeout_handler(signum, frame):
    raise TimeoutError("Tool execution timed out")


def with_timeout(timeout_seconds: int = 5):
    """Decorator to add timeout to tool functions."""

    def decorator(func):
        def wrapper(*args, **kwargs):
            # Set the signal handler
            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(timeout_seconds)
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)
                signal.signal(signal.SIGALRM, old_handler)
            return result

        return wrapper

    return decorator


@with_timeout(5)
def validate_iconclass_code(code: str) -> bool:
    """Check if an iconclass code exists in the system.

    Use this tool to verify that an iconclass code is valid before including
    it in your final answer.

    Args:
        code: The iconclass code to validate (e.g., "25F23", "41A").

    Returns:
        True if the code exists in iconclass, False otherwise.
    """
    if not code or not isinstance(code, str):
        return False
    try:
        _ = ic[code.strip()]
        return True
    except (NotationNotFound, KeyError):
        return False
    except Exception:
        return False


@with_timeout(5)
def lookup_iconclass_code(code: str) -> dict[str, Any]:
    """Look up an iconclass code to get its description and hierarchy.

    Use this tool to understand what an iconclass code represents and see
    its position in the classification hierarchy.

    Args:
        code: The iconclass code to look up (e.g., "25F23", "41A").

    Returns:
        A dictionary with:
        - valid: bool indicating if the code exists
        - description: the text description of the code
        - hierarchy: list of parent codes from root to this code
        - children: list of immediate child codes (if any)
        - related: list of related codes (if any)
    """
    code = code.strip() if code else ""

    if not code:
        return {
            "valid": False,
            "description": "Empty code provided",
            "hierarchy": [],
            "children": [],
            "related": [],
        }

    try:
        notation = ic[code]
        obj = notation.obj

        # Get description
        description = notation() if callable(notation) else str(notation)

        # Get hierarchy (parents)
        hierarchy = obj.get("p", [])

        # Get children (codes that start with this code)
        children = obj.get("c", [])[:10]  # Limit to 10

        # Get related codes
        related = obj.get("r", [])[:10]  # Limit to 10

        return {
            "valid": True,
            "description": description,
            "hierarchy": hierarchy,
            "children": children,
            "related": related,
        }

    except (NotationNotFound, KeyError):
        return {
            "valid": False,
            "description": f"Code '{code}' not found in iconclass",
            "hierarchy": [],
            "children": [],
            "related": [],
        }
    except Exception as e:
        return {
            "valid": False,
            "description": f"Error looking up code: {str(e)}",
            "hierarchy": [],
            "children": [],
            "related": [],
        }


@lru_cache(maxsize=1)
def _get_all_iconclass_codes() -> list[tuple[str, str]]:
    """Get all iconclass codes and their descriptions for search indexing.

    This is cached and only computed once.
    """
    codes_and_descriptions = []

    # Iconclass has a hierarchical structure starting with single digits
    top_level = [str(i) for i in range(10)]

    visited = set()

    def traverse(code: str, depth: int = 0):
        if depth > 8 or code in visited:  # Go deeper for better coverage
            return
        visited.add(code)

        try:
            notation = ic[code]
            desc = notation() if callable(notation) else str(notation)
            codes_and_descriptions.append((code, desc))

            # Get children and traverse
            children = notation.obj.get("c", [])
            for child in children:
                traverse(child, depth + 1)
        except (NotationNotFound, KeyError):
            pass
        except Exception:
            pass

    for top in top_level:
        traverse(top)

    return codes_and_descriptions


# Semantic search components (lazy loaded)
_semantic_search_index = None


def _get_semantic_search_index():
    """Lazily load and cache the semantic search index."""
    global _semantic_search_index

    if _semantic_search_index is not None:
        return _semantic_search_index

    try:
        # Import numpy first to avoid circular import issues
        import numpy as np

        # Delay sentence_transformers import
        from sentence_transformers import SentenceTransformer

        print("Building semantic search index for iconclass (first time only)...")

        # Load a fast, lightweight model
        model = SentenceTransformer("all-MiniLM-L6-v2")

        # Get all codes and descriptions
        codes_and_descriptions = _get_all_iconclass_codes()
        codes = [c[0] for c in codes_and_descriptions]
        descriptions = [c[1] for c in codes_and_descriptions]

        print(f"Encoding {len(descriptions)} iconclass descriptions...")

        # Encode all descriptions (no progress bar to avoid issues)
        embeddings = model.encode(descriptions, show_progress_bar=False, convert_to_numpy=True)

        _semantic_search_index = {
            "model": model,
            "codes": codes,
            "descriptions": descriptions,
            "embeddings": embeddings,
        }

        print(f"Semantic search index ready with {len(codes)} codes")
        return _semantic_search_index

    except ImportError as e:
        print(f"Warning: sentence-transformers not installed ({e}), falling back to keyword search")
        return None
    except Exception as e:
        print(f"Warning: Failed to build semantic index: {e}")
        return None


def search_iconclass_semantic(query: str, limit: int = 5) -> list[dict[str, str]]:
    """Search for iconclass codes using semantic similarity.

    Uses sentence-transformers to find codes whose descriptions are
    semantically similar to the query.

    Args:
        query: A text description to search for.
        limit: Maximum number of results to return (default 5, max 10).

    Returns:
        A list of dictionaries with code, description, and similarity score.
    """
    import numpy as np

    if not query or not isinstance(query, str):
        return []

    limit = min(max(1, limit), 10)
    index = _get_semantic_search_index()

    if index is None:
        # Fallback to keyword search
        return search_iconclass_keyword(query, limit)

    # Encode the query
    query_embedding = index["model"].encode([query], convert_to_numpy=True)[0]

    # Compute cosine similarities
    similarities = np.dot(index["embeddings"], query_embedding) / (
        np.linalg.norm(index["embeddings"], axis=1) * np.linalg.norm(query_embedding)
    )

    # Get top-k indices
    top_indices = np.argsort(similarities)[::-1][:limit]

    results = []
    for idx in top_indices:
        results.append({
            "code": index["codes"][idx],
            "description": index["descriptions"][idx],
            "similarity": float(similarities[idx]),
        })

    return results


def search_iconclass_keyword(query: str, limit: int = 5) -> list[dict[str, str]]:
    """Search for iconclass codes using keyword matching (fallback).

    Args:
        query: A text description to search for.
        limit: Maximum number of results to return.

    Returns:
        A list of dictionaries with code and description.
    """
    if not query or not isinstance(query, str):
        return []

    query = query.lower().strip()
    limit = min(max(1, limit), 10)
    results = []

    try:
        all_codes = _get_all_iconclass_codes()

        for code, description in all_codes:
            desc_lower = description.lower()
            query_terms = query.split()
            if all(term in desc_lower for term in query_terms):
                results.append({"code": code, "description": description})
                if len(results) >= limit:
                    break

    except Exception as e:
        return [{"code": "error", "description": f"Search failed: {str(e)}"}]

    return results


def search_iconclass(query: str, limit: int = 5) -> list[dict[str, str]]:
    """Search for iconclass codes matching a text description.

    Use this tool to find iconclass codes when you can describe what you see
    in the image but don't know the exact code.

    Uses semantic similarity search with sentence-transformers for better
    results. Falls back to keyword search if unavailable.

    Args:
        query: A text description to search for (e.g., "cat", "religious scene",
               "landscape with trees").
        limit: Maximum number of results to return (default 5, max 10).

    Returns:
        A list of dictionaries, each with:
        - code: the iconclass code
        - description: the text description
        - similarity: semantic similarity score (if using semantic search)
    """
    try:
        return search_iconclass_semantic(query, limit)
    except Exception:
        return search_iconclass_keyword(query, limit)


def build_search_index():
    """Pre-build the semantic search index. Call this at startup to avoid delay on first search."""
    return _get_semantic_search_index()


@with_timeout(5)
def get_parent_code(code: str) -> dict[str, Any]:
    """Get the parent code of an iconclass code.

    Use this tool to move up the hierarchy if you want a more general
    classification.

    Args:
        code: The iconclass code to get the parent for.

    Returns:
        A dictionary with:
        - valid: bool indicating if the operation succeeded
        - parent_code: the immediate parent code (or None if at root)
        - parent_description: description of the parent code
    """
    code = code.strip() if code else ""

    if not code:
        return {
            "valid": False,
            "parent_code": None,
            "parent_description": "Empty code provided",
        }

    try:
        notation = ic[code]
        hierarchy = notation.obj.get("p", [])

        if len(hierarchy) <= 1:
            return {
                "valid": True,
                "parent_code": None,
                "parent_description": "Code is at root level",
            }

        # Parent is second-to-last in hierarchy (last is current code)
        parent_code = hierarchy[-2] if len(hierarchy) >= 2 else hierarchy[0]

        try:
            parent_notation = ic[parent_code]
            parent_desc = (
                parent_notation() if callable(parent_notation) else str(parent_notation)
            )
        except Exception:
            parent_desc = f"Parent code: {parent_code}"

        return {
            "valid": True,
            "parent_code": parent_code,
            "parent_description": parent_desc,
        }

    except (NotationNotFound, KeyError):
        return {
            "valid": False,
            "parent_code": None,
            "parent_description": f"Code '{code}' not found",
        }
    except Exception as e:
        return {
            "valid": False,
            "parent_code": None,
            "parent_description": f"Error: {str(e)}",
        }


# List of all tools for easy import
ICONCLASS_TOOLS = [
    validate_iconclass_code,
    lookup_iconclass_code,
    search_iconclass,
    get_parent_code,
]


if __name__ == "__main__":
    # Quick test of tools
    print("Testing iconclass tools...\n")

    print("1. validate_iconclass_code('25F23'):")
    print(f"   {validate_iconclass_code('25F23')}")

    print("\n2. validate_iconclass_code('INVALID'):")
    print(f"   {validate_iconclass_code('INVALID')}")

    print("\n3. lookup_iconclass_code('25F23'):")
    result = lookup_iconclass_code("25F23")
    for k, v in result.items():
        print(f"   {k}: {v}")

    print("\n4. get_parent_code('25F23'):")
    result = get_parent_code("25F23")
    for k, v in result.items():
        print(f"   {k}: {v}")

    print("\n5. search_iconclass('cat') - SEMANTIC SEARCH:")
    results = search_iconclass("cat", limit=5)
    for r in results:
        sim = f" (sim: {r.get('similarity', 'N/A'):.3f})" if "similarity" in r else ""
        desc = r["description"][:60] if len(r["description"]) > 60 else r["description"]
        print(f"   {r['code']}: {desc}...{sim}")

    print("\n6. search_iconclass('blue sky with clouds'):")
    results = search_iconclass("blue sky with clouds", limit=5)
    for r in results:
        sim = f" (sim: {r.get('similarity', 'N/A'):.3f})" if "similarity" in r else ""
        desc = r["description"][:60] if len(r["description"]) > 60 else r["description"]
        print(f"   {r['code']}: {desc}...{sim}")

    print("\n7. search_iconclass('portrait of a woman'):")
    results = search_iconclass("portrait of a woman", limit=5)
    for r in results:
        sim = f" (sim: {r.get('similarity', 'N/A'):.3f})" if "similarity" in r else ""
        desc = r["description"][:60] if len(r["description"]) > 60 else r["description"]
        print(f"   {r['code']}: {desc}...{sim}")