tool-prune 0.1.0__tar.gz

tool_prune-0.1.0/PKG-INFO

@@ -0,0 +1,87 @@
+Metadata-Version: 2.4
+Name: tool-prune
+Version: 0.1.0
+Summary: Calibrated tool selection and schema pruning for AI agents. Zero dependencies.
+Keywords: tool-prune,tool-selection,mcp,agents,typesafe,routing,function-calling,llm
+Author-email: Hemanth HM <hemanth.hm@gmail.com>
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Software Development :: Libraries
+Project-URL: Homepage, https://github.com/hemanth/tool-prune
+Project-URL: Repository, https://github.com/hemanth/tool-prune
+
+# tool-prune
+
+Calibrated tool selection and schema pruning for AI agents. Zero dependencies.
+
+```bash
+pip install tool-prune
+```
+
+## Quick start
+
+```python
+from tool_prune import prune
+
+tools = {
+    "read_file": "Read raw text from local filesystem path",
+    "run_query": "Execute SQL queries against database",
+    "web_search": "Search public web for documentation or articles"
+}
+
+match = prune("what tables exist in the db?", tools)
+print(match.tool)       # 'run_query'
+print(match.confidence) # 1.0
+```
+
+`prune()` narrows schemas to Top-K candidates or identifies direct matches in ~140ms. `dispatch()` runs immediate tool execution if confidence clears the threshold. That's the whole API.
+
+## Schema pruning for LLMs
+
+```python
+from tool_prune import ToolPrune
+
+router = ToolPrune(tools)
+top_tools = router.filter(user_prompt, k=3)
+
+response = llm.chat(
+    tools=top_tools,
+    messages=[{"role": "user", "content": user_prompt}]
+)
+```
+
+Drops prompt tokens by 75-85% and eliminates context distraction without losing tools.
+
+## Fast-path direct dispatch
+
+```python
+result = router.dispatch("read ./pyproject.toml", {
+    "read_file": lambda q, m: open("pyproject.toml").read(),
+    "run_query": lambda q, m: db.query(q),
+    "fallback": lambda q, m: call_llm(q)
+})
+```
+
+Runs deterministic handlers in under 160ms with zero token cost.
+
+## Demo
+
+```bash
+python examples/quickstart.py
+```
+
+Runs live quickstart routing with your `TYPESAFE_API_KEY`.
+
+## License
+
+MIT © [Hemanth.HM](https://h3manth.com)
+

tool_prune-0.1.0/README.md

@@ -0,0 +1,65 @@
+# tool-prune
+
+Calibrated tool selection and schema pruning for AI agents. Zero dependencies.
+
+```bash
+pip install tool-prune
+```
+
+## Quick start
+
+```python
+from tool_prune import prune
+
+tools = {
+    "read_file": "Read raw text from local filesystem path",
+    "run_query": "Execute SQL queries against database",
+    "web_search": "Search public web for documentation or articles"
+}
+
+match = prune("what tables exist in the db?", tools)
+print(match.tool)       # 'run_query'
+print(match.confidence) # 1.0
+```
+
+`prune()` narrows schemas to Top-K candidates or identifies direct matches in ~140ms. `dispatch()` runs immediate tool execution if confidence clears the threshold. That's the whole API.
+
+## Schema pruning for LLMs
+
+```python
+from tool_prune import ToolPrune
+
+router = ToolPrune(tools)
+top_tools = router.filter(user_prompt, k=3)
+
+response = llm.chat(
+    tools=top_tools,
+    messages=[{"role": "user", "content": user_prompt}]
+)
+```
+
+Drops prompt tokens by 75-85% and eliminates context distraction without losing tools.
+
+## Fast-path direct dispatch
+
+```python
+result = router.dispatch("read ./pyproject.toml", {
+    "read_file": lambda q, m: open("pyproject.toml").read(),
+    "run_query": lambda q, m: db.query(q),
+    "fallback": lambda q, m: call_llm(q)
+})
+```
+
+Runs deterministic handlers in under 160ms with zero token cost.
+
+## Demo
+
+```bash
+python examples/quickstart.py
+```
+
+Runs live quickstart routing with your `TYPESAFE_API_KEY`.
+
+## License
+
+MIT © [Hemanth.HM](https://h3manth.com)

tool_prune-0.1.0/pyproject.toml

@@ -0,0 +1,31 @@
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "tool-prune"
+version = "0.1.0"
+description = "Calibrated tool selection and schema pruning for AI agents. Zero dependencies."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [
+    { name = "Hemanth HM", email = "hemanth.hm@gmail.com" },
+]
+keywords = ["tool-prune", "tool-selection", "mcp", "agents", "typesafe", "routing", "function-calling", "llm"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Software Development :: Libraries",
+]
+
+[project.urls]
+Homepage = "https://github.com/hemanth/tool-prune"
+Repository = "https://github.com/hemanth/tool-prune"

tool_prune-0.1.0/tool_prune/__init__.py

@@ -0,0 +1,5 @@
+from .router import ToolPrune, CandidateTool, SelectionResult, prune
+
+__version__ = "0.1.0"
+__all__ = ["ToolPrune", "CandidateTool", "SelectionResult", "prune"]
+

tool_prune-0.1.0/tool_prune/router.py

@@ -0,0 +1,173 @@
+import os
+import time
+import json
+import urllib.request
+import urllib.error
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Union
+
+@dataclass
+class CandidateTool:
+    name: str
+    probability: float
+    tool: Any = None
+
+@dataclass
+class SelectionResult:
+    tool: str
+    confidence: float
+    probability: float
+    top_k: List[CandidateTool]
+    requires_generation: float
+    latency_ms: float
+    usage: Optional[Dict[str, int]] = None
+    raw: Optional[Dict[str, Any]] = None
+
+def _normalize_tools(tools: Union[Dict[str, Any], List[Dict[str, Any]]]):
+    criteria: Dict[str, str] = {}
+    registry: Dict[str, Any] = {}
+
+    if isinstance(tools, list):
+        for t in tools:
+            if not isinstance(t, dict):
+                continue
+            name = t.get("name") or t.get("id")
+            if not name:
+                continue
+            desc = t.get("criteria") or t.get("description") or t.get("summary") or name
+            criteria[name] = str(desc)
+            registry[name] = t
+    elif isinstance(tools, dict):
+        for name, val in tools.items():
+            if isinstance(val, str):
+                criteria[name] = val
+                registry[name] = {"name": name, "description": val}
+            elif isinstance(val, dict):
+                desc = val.get("criteria") or val.get("description") or val.get("summary") or name
+                criteria[name] = str(desc)
+                registry[name] = {"name": name, **val}
+
+    return criteria, registry
+
+class ToolPrune:
+    def __init__(
+        self,
+        tools: Union[Dict[str, Any], List[Dict[str, Any]]],
+        api_key: Optional[str] = None,
+        model: str = "jev-latest",
+        endpoint: str = "https://api.typesafe.ai/v1/systemone",
+        threshold: float = 0.85,
+        top_k: int = 3,
+    ):
+        self.criteria, self.registry = _normalize_tools(tools)
+        self.api_key = api_key or os.getenv("TYPESAFE_API_KEY")
+        self.model = model
+        self.endpoint = endpoint
+        self.threshold = threshold
+        self.top_k = top_k
+
+    def _prepare_payload(self, query: Union[str, Dict[str, Any]], options: Optional[Dict[str, Any]] = None) -> bytes:
+        if not self.api_key:
+            raise ValueError("TypeSafe API key required. Set TYPESAFE_API_KEY environment variable or pass api_key.")
+
+        state = {"intent": query} if isinstance(query, str) else query
+        opts = options or {}
+        body = {
+            "model": opts.get("model", self.model),
+            "state": state,
+            "questions": {
+                "tool": {
+                    "type": "choice",
+                    "instructions": opts.get("instructions", "Which specific tool is required to satisfy this intent?"),
+                    "criteria": self.criteria,
+                },
+                "requires_generation": {
+                    "type": "noul",
+                    "instructions": "Does fulfilling this request require open-ended creative text or arbitrary code generation rather than a deterministic tool execution?",
+                },
+            },
+        }
+        return json.dumps(body).encode("utf-8")
+
+    def _process_response(self, data: Dict[str, Any], latency_ms: float, top_k_count: Optional[int] = None) -> SelectionResult:
+        tool_answer = data.get("answers", {}).get("tool", {})
+        probs = tool_answer.get("probabilities", {})
+        sorted_probs = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)
+
+        k = top_k_count or self.top_k
+        top_candidates = [
+            CandidateTool(name=name, probability=p, tool=self.registry.get(name))
+            for name, p in sorted_probs[:k]
+        ]
+
+        selected_name = tool_answer.get("choice", "")
+        probability = probs.get(selected_name, 0.0)
+        confidence = tool_answer.get("confidence", probability)
+        requires_gen = data.get("answers", {}).get("requires_generation", {}).get("probability", 0.0)
+
+        return SelectionResult(
+            tool=selected_name,
+            confidence=confidence,
+            probability=probability,
+            top_k=top_candidates,
+            requires_generation=requires_gen,
+            latency_ms=latency_ms,
+            usage=data.get("usage"),
+            raw=data,
+        )
+
+    def select(self, query: Union[str, Dict[str, Any]], **kwargs) -> SelectionResult:
+        payload = self._prepare_payload(query, kwargs)
+        req = urllib.request.Request(
+            self.endpoint,
+            data=payload,
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            method="POST",
+        )
+
+        start = time.perf_counter()
+        try:
+            with urllib.request.urlopen(req) as resp:
+                raw_data = resp.read()
+        except urllib.error.HTTPError as e:
+            err = e.read().decode("utf-8", errors="replace")
+            raise RuntimeError(f"tool-prune API error ({e.code}): {err}") from e
+
+        latency_ms = (time.perf_counter() - start) * 1000.0
+        data = json.loads(raw_data.decode("utf-8"))
+        return self._process_response(data, latency_ms, kwargs.get("top_k"))
+
+    def filter(self, query: Union[str, Dict[str, Any]], k: Optional[int] = None, **kwargs) -> List[Any]:
+        result = self.select(query, top_k=k or self.top_k, **kwargs)
+        target_k = k or self.top_k
+        return [c.tool if c.tool is not None else c.name for c in result.top_k[:target_k]]
+
+    def dispatch(
+        self,
+        query: Union[str, Dict[str, Any]],
+        handlers: Dict[str, Callable[..., Any]],
+        threshold: Optional[float] = None,
+        **kwargs,
+    ) -> Any:
+        selection = self.select(query, **kwargs)
+        t = threshold if threshold is not None else self.threshold
+
+        if selection.confidence >= t and selection.tool in handlers:
+            return handlers[selection.tool](query, selection)
+
+        if "fallback" in handlers:
+            return handlers["fallback"](query, selection)
+
+        return selection
+
+def prune(
+    query: Union[str, Dict[str, Any]],
+    tools: Union[Dict[str, Any], List[Dict[str, Any]]],
+    **kwargs,
+) -> SelectionResult:
+    """One-shot tool pruning and selection."""
+    return ToolPrune(tools, **kwargs).select(query)
+