ss-tools/backend/src/agent/_embedding_router.py

# backend/src/agent/_embedding_router.py
# #region AgentChat.EmbeddingRouter [C:3] [TYPE Module] [SEMANTICS agent-chat,tools,embedding,fallback]
# @defgroup AgentChat Embedding-based tool router — fallback when keyword matching yields <3 tools.
# @LAYER Service
# @BRIEF Lazy-loaded embedding model for cosine similarity between user query and tool descriptions.
# @RATIONALE Tool descriptions are auto-generated from @tool docstrings (enforced by LangChain).
#            Optional _TOOL_DESCRIPTIONS_OVERRIDES in tools.py provides RU/EN synonyms for
#            tools where the docstring alone isn't descriptive enough. This eliminates the
#            hardcoded _TOOL_DESCRIPTIONS dict as a separate source of truth.
# @REJECTED Hardcoded _TOOL_DESCRIPTIONS dict — drifts out of sync with get_all_tools().
# @INVARIANT Descriptions are derived from get_all_tools() docstrings — always 1:1 with tools.
# @INVARIANT embedding_top_k() returns empty list (never raises) when model unavailable.

import logging
import os
from typing import Optional

logger = logging.getLogger("superset_tools_app")

# ═══════════════════════════════════════════════════════════════════
#  Tool description — auto-generated from @tool docstrings.
#  Override via _TOOL_DESCRIPTIONS_OVERRIDES in tools.py for RU/EN synonyms.
# ═══════════════════════════════════════════════════════════════════

def _get_descriptions() -> tuple[list[str], list[str]]:
    """Return (descriptions, tool_names) from get_all_tools() docstrings.
    Uses _TOOL_DESCRIPTIONS_OVERRIDES from tools.py for optional RU/EN synonyms.
    """
    from src.agent.tools import get_all_tools, _TOOL_DESCRIPTIONS_OVERRIDES

    all_tools = get_all_tools()
    names = []
    descriptions = []
    for tool_obj in all_tools:
        name = tool_obj.name
        names.append(name)
        desc = _TOOL_DESCRIPTIONS_OVERRIDES.get(name) or (tool_obj.description or "").strip()
        if not desc:
            desc = name  # fallback: use tool name as description
        descriptions.append(desc)
    return descriptions, names

# ═══════════════════════════════════════════════════════════════════
#  Model state — lazy-loaded on first call
# ═══════════════════════════════════════════════════════════════════

_embedding_model: Optional[object] = None
_tool_embeddings: Optional[object] = None  # torch.Tensor or numpy array
_tool_names: list[str] = []

_THRESHOLD = float(os.getenv("EMBEDDING_SIMILARITY_THRESHOLD", "0.65"))
_TOP_K = int(os.getenv("EMBEDDING_TOP_K", "5"))
_MODEL_NAME = os.getenv(
    "EMBEDDING_MODEL",
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
)


def _load_model() -> bool:
    """Lazy-load the embedding model and pre-embed tool descriptions.

    Returns True on success, False on any failure (missing package, download error,
    OOM, etc.). On failure, the router degrades gracefully — embedding_top_k()
    returns an empty list and the caller falls back to keyword-only results.
    """
    global _embedding_model, _tool_embeddings, _tool_names

    if _embedding_model is not None:
        return True

    try:
        from sentence_transformers import SentenceTransformer
    except ImportError:
        logger.warning(
            "sentence-transformers not installed — embedding router disabled. "
            "Install with: pip install sentence-transformers"
        )
        return False

    try:
        logger.info("Loading embedding model: %s", _MODEL_NAME)
        _embedding_model = SentenceTransformer(_MODEL_NAME)

        descriptions, _tool_names[:] = _get_descriptions()

        _tool_embeddings = _embedding_model.encode(
            descriptions,
            convert_to_tensor=True,
            show_progress_bar=False,
        )
        logger.info(
            "Embedding model loaded. Tools: %d, model: %s",
            len(_tool_names), _MODEL_NAME,
        )
        return True
    except Exception as exc:
        logger.warning(
            "Failed to load embedding model '%s': %s — embedding router disabled",
            _MODEL_NAME, exc,
        )
        _embedding_model = None
        return False


def embedding_top_k(query: str) -> list[str]:
    """Return top-K tool names above cosine similarity threshold.

    Args:
        query: Raw user query text (any language).

    Returns:
        List of tool name strings ordered by descending similarity.
        Empty list if model unavailable, no descriptions match, or an error occurs.
    """
    if not query or not query.strip():
        return []

    if not _load_model():
        return []

    try:
        from sentence_transformers.util import cos_sim

        query_embedding = _embedding_model.encode(
            query,
            convert_to_tensor=True,
            show_progress_bar=False,
        )

        similarities = cos_sim(query_embedding, _tool_embeddings)[0]

        results: list[str] = []
        for idx in similarities.argsort(descending=True)[:_TOP_K]:
            score = float(similarities[idx])
            if score >= _THRESHOLD:
                results.append(_tool_names[idx])

        if results:
            logger.debug(
                "Embedding router: query='%s' → %d tools above %.2f: %s",
                query[:80], len(results), _THRESHOLD, results,
            )

        return results
    except Exception as exc:
        logger.warning("Embedding router failed for query '%s': %s", query[:80], exc)
        return []


def embedding_is_available() -> bool:
    """Check if embedding model is loaded and ready (non-blocking)."""
    return _load_model()


# #endregion AgentChat.EmbeddingRouter