cheapskate 0.1.0__py3-none-any.whl

cheapskate/__init__.py

@@ -0,0 +1,22 @@
+from cheapskate.callback import CheapSkateCallbackHandler
+from cheapskate.compressor import PromptCompressor
+from cheapskate.pruner import ToolNamespace, ToolPruner
+from cheapskate.router import CheapSkateRouter
+from cheapskate.token_counter import TokenCounter
+
+__all__ = [
+    "CheapSkateRouter",
+    "CheapSkateCallbackHandler",
+    "TokenCounter",
+    "ToolPruner",
+    "ToolNamespace",
+    "PromptCompressor",
+]
+
+__version__ = "0.1.0"
+
+
+def package_info() -> str:
+    info = f"cheapskate {__version__}"
+    print(info)
+    return info

cheapskate/callback.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.messages import BaseMessage
+from langchain_core.outputs import LLMResult
+from loguru import logger
+
+from cheapskate.token_counter import TokenCounter
+
+
+class CheapSkateCallbackHandler(BaseCallbackHandler):
+    def __init__(self, token_counter: Optional[TokenCounter] = None) -> None:
+        super().__init__()
+        self.token_counter = token_counter or TokenCounter()
+        self.routing_events: List[Dict[str, Any]] = []
+        self.prune_events: List[Dict[str, Any]] = []
+        self.compression_events: List[Dict[str, Any]] = []
+        self.fallback_events: List[Dict[str, Any]] = []
+        self._pending_baseline: Dict[UUID, int] = {}
+
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            flat_messages = [message for batch in messages for message in batch]
+            baseline = self.token_counter.record_baseline(flat_messages)
+            self._pending_baseline[run_id] = baseline
+            logger.debug("CheapSkate tracked chat start run_id={} baseline_tokens={}", run_id, baseline)
+        except Exception as exc:
+            logger.warning("CheapSkateCallbackHandler failed during chat start: {}", exc)
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._pending_baseline.pop(run_id, None)
+        except Exception as exc:
+            logger.warning("CheapSkateCallbackHandler failed during llm end: {}", exc)
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        logger.error("CheapSkate observed LLM error for run_id={}: {}", run_id, error)
+        self._pending_baseline.pop(run_id, None)
+
+    def record_routing(self, destination: str, reason: str, metadata: Optional[Dict[str, Any]] = None) -> None:
+        event = {"destination": destination, "reason": reason, "metadata": metadata or {}}
+        self.routing_events.append(event)
+        logger.info("CheapSkate routed to {} because {}", destination, reason)
+
+    def record_prune(
+        self,
+        retained: List[str],
+        dropped: List[str],
+        namespaces: List[str],
+    ) -> None:
+        event = {
+            "retained": list(retained),
+            "dropped": list(dropped),
+            "namespaces": list(namespaces),
+        }
+        self.prune_events.append(event)
+
+    def record_compression(self, before_tokens: int, after_tokens: int) -> None:
+        event = {
+            "before_tokens": before_tokens,
+            "after_tokens": after_tokens,
+            "tokens_saved": max(0, before_tokens - after_tokens),
+        }
+        self.compression_events.append(event)
+
+    def record_fallback(self, error: Union[str, BaseException], recovered: bool = True) -> None:
+        event = {"error": str(error), "recovered": recovered}
+        self.fallback_events.append(event)
+        logger.warning("CheapSkate fallback engaged: {}", error)
+
+    @property
+    def summary(self) -> Dict[str, Any]:
+        return {
+            "token_stats": self.token_counter.stats,
+            "routing_events": len(self.routing_events),
+            "prune_events": len(self.prune_events),
+            "compression_events": len(self.compression_events),
+            "fallback_events": len(self.fallback_events),
+            "last_route": self.routing_events[-1] if self.routing_events else None,
+        }
+
+    def reset(self) -> None:
+        self.routing_events.clear()
+        self.prune_events.clear()
+        self.compression_events.clear()
+        self.fallback_events.clear()
+        self._pending_baseline.clear()
+        self.token_counter.reset()

cheapskate/compressor.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import copy
+import json
+import re
+from typing import Dict, List, Optional, Sequence, Tuple
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from loguru import logger
+
+from cheapskate.token_counter import TokenCounter
+
+
+class PromptCompressor:
+    BOILERPLATE_PATTERNS: Tuple[re.Pattern[str], ...] = (
+        re.compile(r"\bplease\s+(note|be\s+advised|remember|ensure)\b[:\s]*", re.IGNORECASE),
+        re.compile(r"\bas\s+(an?\s+)?(ai|assistant|language\s+model)\b[,:\s]*", re.IGNORECASE),
+        re.compile(r"\bi\s+(hope\s+this\s+helps|understand\s+your\s+request)\b[!.\s]*", re.IGNORECASE),
+        re.compile(r"\b(just\s+to\s+clarify|to\s+be\s+clear|for\s+your\s+information)\b[,:\s]*", re.IGNORECASE),
+        re.compile(r"\b(kindly|basically|literally|actually|obviously|certainly)\b\s+", re.IGNORECASE),
+        re.compile(r"\b(in\s+order\s+to)\b", re.IGNORECASE),
+        re.compile(r"\b(it\s+is\s+important\s+to\s+note\s+that)\b\s*", re.IGNORECASE),
+        re.compile(r"\b(at\s+this\s+point\s+in\s+time)\b", re.IGNORECASE),
+        re.compile(r"[ \t]{2,}"),
+        re.compile(r"\n{3,}"),
+    )
+
+    FILLER_REPLACEMENTS: Tuple[Tuple[re.Pattern[str], str], ...] = (
+        (re.compile(r"\bin\s+order\s+to\b", re.IGNORECASE), "to"),
+        (re.compile(r"\bdue\s+to\s+the\s+fact\s+that\b", re.IGNORECASE), "because"),
+        (re.compile(r"\ba\s+large\s+number\s+of\b", re.IGNORECASE), "many"),
+        (re.compile(r"\bat\s+the\s+present\s+time\b", re.IGNORECASE), "now"),
+        (re.compile(r"\bin\s+the\s+event\s+that\b", re.IGNORECASE), "if"),
+        (re.compile(r"\bwith\s+regard\s+to\b", re.IGNORECASE), "about"),
+        (re.compile(r"\bprior\s+to\b", re.IGNORECASE), "before"),
+    )
+
+    JSON_BLOCK_PATTERN = re.compile(r"(\{[\s\S]*?\}|\[[\s\S]*?\])")
+    WHITESPACE_COLLAPSE = re.compile(r"[ \t]+")
+
+    def __init__(
+        self,
+        token_counter: Optional[TokenCounter] = None,
+        target_reduction_ratio: float = 0.35,
+        min_message_tokens: int = 24,
+        preserve_system_messages: bool = True,
+        preserve_recent_turns: int = 2,
+    ) -> None:
+        self._token_counter = token_counter or TokenCounter()
+        self._target_reduction_ratio = min(0.85, max(0.0, target_reduction_ratio))
+        self._min_message_tokens = max(1, min_message_tokens)
+        self._preserve_system_messages = preserve_system_messages
+        self._preserve_recent_turns = max(0, preserve_recent_turns)
+
+    def _shallow_copy_messages(self, messages: Sequence[BaseMessage]) -> List[BaseMessage]:
+        duplicated: List[BaseMessage] = []
+        for message in messages:
+            try:
+                duplicated.append(message.model_copy(deep=False))
+            except Exception:
+                duplicated.append(copy.copy(message))
+        return duplicated
+
+    def _protect_json_segments(self, text: str) -> Tuple[str, Dict[str, str]]:
+        placeholders: Dict[str, str] = {}
+        counter = 0
+
+        def _replacer(match: re.Match[str]) -> str:
+            nonlocal counter
+            candidate = match.group(0)
+            try:
+                json.loads(candidate)
+            except (json.JSONDecodeError, ValueError):
+                return candidate
+            token = f"__CHEAPSKATE_JSON_{counter}__"
+            placeholders[token] = candidate
+            counter += 1
+            return token
+
+        protected = self.JSON_BLOCK_PATTERN.sub(_replacer, text)
+        return protected, placeholders
+
+    @staticmethod
+    def _restore_json_segments(text: str, placeholders: Dict[str, str]) -> str:
+        restored = text
+        for token, original in placeholders.items():
+            restored = restored.replace(token, original)
+        return restored
+
+    def _scrub_text(self, text: str) -> str:
+        if not text or not text.strip():
+            return text
+        protected, placeholders = self._protect_json_segments(text)
+        scrubbed = protected
+        for pattern, replacement in self.FILLER_REPLACEMENTS:
+            scrubbed = pattern.sub(replacement, scrubbed)
+        for pattern in self.BOILERPLATE_PATTERNS:
+            if pattern.pattern.startswith("[") or pattern.pattern.startswith("\\n"):
+                scrubbed = pattern.sub("\n\n" if "\\n" in pattern.pattern else " ", scrubbed)
+            else:
+                scrubbed = pattern.sub("", scrubbed)
+        scrubbed = self.WHITESPACE_COLLAPSE.sub(" ", scrubbed)
+        scrubbed = re.sub(r"\n{3,}", "\n\n", scrubbed)
+        scrubbed = scrubbed.strip()
+        return self._restore_json_segments(scrubbed, placeholders)
+
+    def _token_density(self, text: str) -> float:
+        tokens = self._token_counter.count_text(text)
+        if tokens <= 0:
+            return 0.0
+        unique_terms = len({term.lower() for term in re.findall(r"[A-Za-z0-9_]+", text)})
+        return unique_terms / float(tokens)
+
+    def _truncate_by_density(self, text: str, budget_tokens: int) -> str:
+        if budget_tokens <= 0:
+            return text
+        current_tokens = self._token_counter.count_text(text)
+        if current_tokens <= budget_tokens:
+            return text
+        sentences = re.split(r"(?<=[.!?])\s+", text)
+        if len(sentences) <= 1:
+            words = text.split()
+            while words and self._token_counter.count_text(" ".join(words)) > budget_tokens:
+                words.pop()
+            return " ".join(words).strip()
+        ranked = sorted(
+            ((self._token_density(sentence), index, sentence) for index, sentence in enumerate(sentences)),
+            key=lambda item: item[0],
+            reverse=True,
+        )
+        selected_indices = set()
+        running = 0
+        for _, index, sentence in ranked:
+            sentence_tokens = self._token_counter.count_text(sentence)
+            if running + sentence_tokens > budget_tokens and selected_indices:
+                continue
+            selected_indices.add(index)
+            running += sentence_tokens
+            if running >= budget_tokens:
+                break
+        ordered = [sentences[index] for index in sorted(selected_indices)]
+        return " ".join(ordered).strip()
+
+    def _should_preserve(self, message: BaseMessage, absolute_index: int, total: int) -> bool:
+        if self._preserve_system_messages and isinstance(message, SystemMessage):
+            return True
+        if absolute_index >= max(0, total - self._preserve_recent_turns):
+            return True
+        return False
+
+    def _set_content(self, message: BaseMessage, content: str) -> BaseMessage:
+        try:
+            message.content = content
+            return message
+        except Exception:
+            if isinstance(message, SystemMessage):
+                return SystemMessage(content=content)
+            if isinstance(message, HumanMessage):
+                return HumanMessage(content=content)
+            if isinstance(message, AIMessage):
+                return AIMessage(content=content)
+            if isinstance(message, ToolMessage):
+                return ToolMessage(
+                    content=content,
+                    tool_call_id=getattr(message, "tool_call_id", "unknown"),
+                    name=getattr(message, "name", None),
+                )
+            return message
+
+    def compress(self, messages: Sequence[BaseMessage]) -> List[BaseMessage]:
+        if not messages:
+            return []
+
+        working_copy = self._shallow_copy_messages(messages)
+        baseline_tokens = self._token_counter.count_messages(working_copy)
+        target_tokens = int(baseline_tokens * (1.0 - self._target_reduction_ratio))
+        target_tokens = max(target_tokens, self._min_message_tokens)
+
+        compressed: List[BaseMessage] = []
+        total = len(working_copy)
+        for index, message in enumerate(working_copy):
+            raw_content = message.content
+            if isinstance(raw_content, list):
+                text_content = "\n".join(
+                    str(block.get("text", block) if isinstance(block, dict) else block)
+                    for block in raw_content
+                )
+            else:
+                text_content = str(raw_content or "")
+
+            if self._should_preserve(message, index, total):
+                scrubbed = self._scrub_text(text_content) if text_content else text_content
+                compressed.append(self._set_content(message, scrubbed))
+                continue
+
+            scrubbed = self._scrub_text(text_content)
+            message_tokens = self._token_counter.count_text(scrubbed)
+            if message_tokens > self._min_message_tokens * 2:
+                density = self._token_density(scrubbed)
+                budget = max(
+                    self._min_message_tokens,
+                    int(message_tokens * max(0.45, min(0.95, density + 0.25))),
+                )
+                scrubbed = self._truncate_by_density(scrubbed, budget)
+            compressed.append(self._set_content(message, scrubbed))
+
+        after_tokens = self._token_counter.count_messages(compressed)
+        if after_tokens > target_tokens and len(compressed) > 2:
+            overflow = after_tokens - target_tokens
+            for index in range(len(compressed) - self._preserve_recent_turns - 1, -1, -1):
+                if overflow <= 0:
+                    break
+                message = compressed[index]
+                if isinstance(message, SystemMessage):
+                    continue
+                content = str(message.content or "")
+                current = self._token_counter.count_text(content)
+                reduced_budget = max(self._min_message_tokens, current - overflow)
+                truncated = self._truncate_by_density(content, reduced_budget)
+                compressed[index] = self._set_content(message, truncated)
+                new_count = self._token_counter.count_text(truncated)
+                overflow -= max(0, current - new_count)
+
+        logger.debug(
+            "Prompt compression reduced tokens from {} to {}",
+            baseline_tokens,
+            self._token_counter.count_messages(compressed),
+        )
+        return compressed

cheapskate/pruner.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import hashlib
+import math
+import re
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple
+
+import numpy as np
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
+from loguru import logger
+
+
+ToolLike = Any
+NamespaceClassifier = Callable[[str, Sequence[str]], Sequence[str]]
+
+
+@dataclass(frozen=True)
+class ToolNamespace:
+    name: str
+    tools: Tuple[str, ...]
+    description: str = ""
+
+
+@dataclass
+class PruneResult:
+    retained_tools: List[ToolLike]
+    retained_names: List[str]
+    dropped_names: List[str]
+    selected_namespaces: List[str]
+    similarity_scores: Dict[str, float] = field(default_factory=dict)
+
+
+class ToolPruner:
+    DEFAULT_EMBEDDING_DIM = 384
+    DEFAULT_TOP_K = 8
+    DEFAULT_SIMILARITY_FLOOR = 0.12
+    TOKEN_PATTERN = re.compile(r"[a-z0-9_]+", re.IGNORECASE)
+
+    def __init__(
+        self,
+        namespaces: Optional[Sequence[ToolNamespace]] = None,
+        always_keep: Optional[Sequence[str]] = None,
+        top_k: int = DEFAULT_TOP_K,
+        similarity_floor: float = DEFAULT_SIMILARITY_FLOOR,
+        embedding_dim: int = DEFAULT_EMBEDDING_DIM,
+        namespace_classifier: Optional[NamespaceClassifier] = None,
+    ) -> None:
+        self._namespaces: List[ToolNamespace] = list(namespaces or [])
+        self._always_keep: Set[str] = {name.strip() for name in (always_keep or []) if name}
+        self._top_k = max(1, top_k)
+        self._similarity_floor = max(0.0, min(1.0, similarity_floor))
+        self._embedding_dim = max(64, embedding_dim)
+        self._namespace_classifier = namespace_classifier
+        self._namespace_lookup: Dict[str, ToolNamespace] = {
+            namespace.name: namespace for namespace in self._namespaces
+        }
+
+    def register_namespace(self, namespace: ToolNamespace) -> None:
+        self._namespaces.append(namespace)
+        self._namespace_lookup[namespace.name] = namespace
+
+    def set_always_keep(self, tool_names: Sequence[str]) -> None:
+        self._always_keep = {name.strip() for name in tool_names if name}
+
+    @property
+    def always_keep(self) -> Set[str]:
+        return set(self._always_keep)
+
+    def _tool_name(self, tool: ToolLike) -> str:
+        if isinstance(tool, dict):
+            if "function" in tool and isinstance(tool["function"], dict):
+                return str(tool["function"].get("name", "")).strip()
+            return str(tool.get("name", "")).strip()
+        name = getattr(tool, "name", None)
+        if name:
+            return str(name).strip()
+        return str(tool).strip()
+
+    def _tool_description(self, tool: ToolLike) -> str:
+        if isinstance(tool, dict):
+            if "function" in tool and isinstance(tool["function"], dict):
+                function_block = tool["function"]
+                return " ".join(
+                    [
+                        str(function_block.get("name", "")),
+                        str(function_block.get("description", "")),
+                        str(function_block.get("parameters", "")),
+                    ]
+                )
+            return " ".join(
+                [
+                    str(tool.get("name", "")),
+                    str(tool.get("description", "")),
+                    str(tool.get("parameters", tool.get("args_schema", ""))),
+                ]
+            )
+        pieces = [
+            str(getattr(tool, "name", "") or ""),
+            str(getattr(tool, "description", "") or ""),
+        ]
+        args_schema = getattr(tool, "args_schema", None)
+        if args_schema is not None:
+            schema_method = getattr(args_schema, "schema", None)
+            if callable(schema_method):
+                pieces.append(str(schema_method()))
+            else:
+                pieces.append(str(args_schema))
+        return " ".join(piece for piece in pieces if piece)
+
+    def _tokenize(self, text: str) -> List[str]:
+        return [token.lower() for token in self.TOKEN_PATTERN.findall(text or "") if token]
+
+    def _embed_text(self, text: str) -> np.ndarray:
+        vector = np.zeros(self._embedding_dim, dtype=np.float64)
+        tokens = self._tokenize(text)
+        if not tokens:
+            return vector
+        for index, token in enumerate(tokens):
+            digest = hashlib.sha256(token.encode("utf-8")).digest()
+            bucket = int.from_bytes(digest[:4], byteorder="big") % self._embedding_dim
+            sign = 1.0 if digest[4] % 2 == 0 else -1.0
+            weight = 1.0 + (1.0 / (1.0 + index))
+            vector[bucket] += sign * weight
+            if len(token) >= 4:
+                bigram = token[:4]
+                bigram_digest = hashlib.md5(bigram.encode("utf-8")).digest()
+                bigram_bucket = int.from_bytes(bigram_digest[:4], byteorder="big") % self._embedding_dim
+                vector[bigram_bucket] += 0.35 * sign
+        norm = np.linalg.norm(vector)
+        if norm > 0:
+            vector /= norm
+        return vector
+
+    def _message_corpus(self, messages: Sequence[BaseMessage]) -> str:
+        fragments: List[str] = []
+        for message in messages:
+            content = getattr(message, "content", "")
+            if isinstance(content, list):
+                content = " ".join(
+                    str(block.get("text", block) if isinstance(block, dict) else block)
+                    for block in content
+                )
+            fragments.append(str(content or ""))
+            if isinstance(message, AIMessage) and getattr(message, "tool_calls", None):
+                for tool_call in message.tool_calls or []:
+                    fragments.append(str(tool_call.get("name", "")))
+                    fragments.append(str(tool_call.get("args", "")))
+            if isinstance(message, ToolMessage):
+                fragments.append(str(getattr(message, "name", "") or ""))
+                fragments.append(str(getattr(message, "tool_call_id", "") or ""))
+            if isinstance(message, HumanMessage):
+                fragments.append(str(content or ""))
+        return "\n".join(fragment for fragment in fragments if fragment)
+
+    @staticmethod
+    def _cosine_similarity_matrix(query_vector: np.ndarray, matrix: np.ndarray) -> np.ndarray:
+        if matrix.size == 0:
+            return np.array([], dtype=np.float64)
+        query_norm = np.linalg.norm(query_vector)
+        if query_norm == 0:
+            return np.zeros(matrix.shape[0], dtype=np.float64)
+        matrix_norms = np.linalg.norm(matrix, axis=1)
+        safe_norms = np.where(matrix_norms == 0, 1.0, matrix_norms)
+        scores = (matrix @ query_vector) / (safe_norms * query_norm)
+        return np.nan_to_num(scores, nan=0.0, posinf=0.0, neginf=0.0)
+
+    def _classify_namespaces(self, query_text: str) -> List[str]:
+        if not self._namespaces:
+            return []
+        namespace_names = [namespace.name for namespace in self._namespaces]
+        if self._namespace_classifier is not None:
+            try:
+                classified = list(self._namespace_classifier(query_text, namespace_names))
+                valid = [name for name in classified if name in self._namespace_lookup]
+                if valid:
+                    return valid
+            except Exception as exc:
+                logger.warning("Namespace classifier failed; using embedding fallback: {}", exc)
+        namespace_matrix = np.vstack(
+            [
+                self._embed_text(f"{namespace.name} {namespace.description} {' '.join(namespace.tools)}")
+                for namespace in self._namespaces
+            ]
+        )
+        query_vector = self._embed_text(query_text)
+        scores = self._cosine_similarity_matrix(query_vector, namespace_matrix)
+        ranked_indices = np.argsort(scores)[::-1]
+        selected: List[str] = []
+        for index in ranked_indices:
+            if float(scores[index]) < self._similarity_floor and selected:
+                break
+            selected.append(self._namespaces[int(index)].name)
+            if len(selected) >= max(1, math.ceil(len(self._namespaces) * 0.5)):
+                break
+        return selected or [self._namespaces[int(ranked_indices[0])].name]
+
+    def _namespace_tool_allowlist(self, selected_namespaces: Sequence[str]) -> Optional[Set[str]]:
+        if not selected_namespaces or not self._namespaces:
+            return None
+        allowlist: Set[str] = set()
+        for namespace_name in selected_namespaces:
+            namespace = self._namespace_lookup.get(namespace_name)
+            if namespace is None:
+                continue
+            allowlist.update(namespace.tools)
+        return allowlist or None
+
+    def prune(
+        self,
+        tools: Sequence[ToolLike],
+        messages: Sequence[BaseMessage],
+        always_keep: Optional[Sequence[str]] = None,
+        top_k: Optional[int] = None,
+    ) -> PruneResult:
+        if not tools:
+            return PruneResult(
+                retained_tools=[],
+                retained_names=[],
+                dropped_names=[],
+                selected_namespaces=[],
+                similarity_scores={},
+            )
+
+        mandatory_names = set(self._always_keep)
+        if always_keep:
+            mandatory_names.update(name.strip() for name in always_keep if name)
+
+        query_text = self._message_corpus(messages)
+        selected_namespaces = self._classify_namespaces(query_text)
+        namespace_allowlist = self._namespace_tool_allowlist(selected_namespaces)
+
+        tool_entries: List[Tuple[ToolLike, str, str]] = []
+        for tool in tools:
+            name = self._tool_name(tool)
+            description = self._tool_description(tool)
+            tool_entries.append((tool, name, description))
+
+        if namespace_allowlist is not None:
+            candidate_entries = [
+                entry
+                for entry in tool_entries
+                if entry[1] in mandatory_names or entry[1] in namespace_allowlist
+            ]
+            if not candidate_entries:
+                candidate_entries = list(tool_entries)
+        else:
+            candidate_entries = list(tool_entries)
+
+        query_vector = self._embed_text(query_text)
+        description_matrix = np.vstack(
+            [self._embed_text(f"{name} {description}") for _, name, description in candidate_entries]
+        )
+        similarity_scores = self._cosine_similarity_matrix(query_vector, description_matrix)
+
+        scored_candidates: List[Tuple[float, ToolLike, str]] = []
+        for index, (tool, name, _) in enumerate(candidate_entries):
+            score = float(similarity_scores[index])
+            if name in mandatory_names:
+                score = max(score, 1.0)
+            scored_candidates.append((score, tool, name))
+
+        scored_candidates.sort(key=lambda item: item[0], reverse=True)
+        limit = top_k if top_k is not None else self._top_k
+        retained: List[ToolLike] = []
+        retained_names: List[str] = []
+        score_map: Dict[str, float] = {}
+        seen_names: Set[str] = set()
+
+        for score, tool, name in scored_candidates:
+            if name in mandatory_names:
+                if name not in seen_names:
+                    retained.append(tool)
+                    retained_names.append(name)
+                    score_map[name] = score
+                    seen_names.add(name)
+                continue
+            if len([item for item in retained_names if item not in mandatory_names]) >= limit:
+                continue
+            if score < self._similarity_floor and name not in mandatory_names:
+                continue
+            if name not in seen_names:
+                retained.append(tool)
+                retained_names.append(name)
+                score_map[name] = score
+                seen_names.add(name)
+
+        for tool, name, _ in tool_entries:
+            if name in mandatory_names and name not in seen_names:
+                retained.append(tool)
+                retained_names.append(name)
+                score_map[name] = 1.0
+                seen_names.add(name)
+
+        if not retained:
+            fallback_slice = tool_entries[:limit]
+            retained = [tool for tool, _, _ in fallback_slice]
+            retained_names = [name for _, name, _ in fallback_slice]
+            score_map = {name: 0.0 for name in retained_names}
+
+        dropped_names = [
+            name for _, name, _ in tool_entries if name and name not in set(retained_names)
+        ]
+
+        logger.info(
+            "Tool pruning retained {}/{} tools across namespaces {}",
+            len(retained_names),
+            len(tool_entries),
+            selected_namespaces,
+        )
+
+        return PruneResult(
+            retained_tools=retained,
+            retained_names=retained_names,
+            dropped_names=dropped_names,
+            selected_namespaces=list(selected_namespaces),
+            similarity_scores=score_map,
+        )