single-file example HTTP MCP tool server in Python 3.12.11 (for use with `llama-ui` or any other MCP-capable inference client)

serve.py

# serve.py
# Python 3.12.11

import os
import re
import sys
import math
import datetime

from typing import Optional, Literal
from dataclasses import dataclass
from urllib.parse import quote, unquote

# NEEDS PACKAGES: run this command:
# `pip install fastmcp numpy requests beautifulsoup4 lxml curl_cffi>=0.6.0 trafilatura>=1.6.0`

import numpy as np
import requests as _requests
from fastmcp import FastMCP

from bs4 import BeautifulSoup

try:
    from curl_cffi import requests as _curl_requests
    HAS_CURL_CFFI = True
except ImportError:
    HAS_CURL_CFFI = False

try:
    import trafilatura as _trafilatura
    HAS_TRAFILATURA = True
except ImportError:
    HAS_TRAFILATURA = False

serve = FastMCP(name="Local MCP Server (Python Tools)")

#
# Constants
#

_N_PREFETCH = 256
_DTYPE_NP = np.float32
_DDG_ENDPOINT = "https://html.duckduckgo.com/html/"
_LEGACY_UA = "Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1)"

class PowersOfTwo:
    """On-demand cached sequence of powers of two."""

    def _compute_to(self, n: int) -> None:
        while self._computed_through < n:
            self._values.append(self._values[-1] * 2)
            self._computed_through += 1

    def __getitem__(self, n: int) -> int:
        if not isinstance(n, int):
            raise TypeError("index must be a non-negative integer")
        if n < 0:
            raise IndexError("negative indices not supported")
        self._compute_to(n)
        return self._values[n]

    def __iter__(self):
        index = 0
        while True:
            yield self[index]
            index += 1

    def get(self, n: int) -> int:
        """Explicit getter; same as indexing."""
        return self[n]

    def __init__(self):
        self._values: list[int] = [1]
        self._computed_through: int = 0
        self._compute_to(_N_PREFETCH)

POWERS = PowersOfTwo()

class Primes:
    """On-demand cached sequence of prime numbers."""

    def _compute_to(self, n: int) -> None:
        while self._computed_through < n:
            candidate = self._candidate
            if all(candidate % p != 0 for p in self._primes):
                self._primes.append(candidate)
                self._computed_through += 1
            self._candidate += 1

    def __getitem__(self, n: int) -> int:
        if not isinstance(n, int):
            raise TypeError("index must be a non-negative integer")
        if n < 0:
            raise IndexError("negative indices not supported")
        self._compute_to(n)
        return self._primes[n]

    def get(self, n: int) -> int:
        """Explicit getter; same as indexing."""
        return self[n]

    def __iter__(self):
        index = 0
        while True:
            yield self[index]
            index += 1

    def __init__(self):
        self._primes: list[int] = []
        self._computed_through: int = -1
        self._candidate: int = 2
        self._compute_to(_N_PREFETCH)

PRIMES = Primes()

@serve.tool
def add(a: int | float, b: int | float) -> int | float:
    """Return the sum of `a` and `b`."""
    return a + b

@serve.tool
def sub(a: int | float, b: int | float) -> int | float:
    """Return the difference of `a` minus `b`."""
    return a - b

@serve.tool
def mul(a: int | float, b: int | float) -> int | float:
    """Return the product of `a` and `b`."""
    return a * b

@serve.tool
def div(a: int | float, b: int | float) -> float:
    """Return the quotient of `a` divided by `b`. Raises ZeroDivisionError if `b` is zero."""
    if b == 0:
        raise ZeroDivisionError("division by zero")
    return a / b

@serve.tool
def log(n: int | float) -> int | float:
    """Return the natural logarithim of `n`."""
    return math.log(n)

@serve.tool
def log2(n: int | float) -> int | float:
    """Return the binary logarithim of `n`."""
    return math.log2(n)

@serve.tool
def log10(n: int | float) -> int | float:
    """Return the base-10 logarithim of `n`."""
    return math.log10(n)

@serve.tool
def sqrt(n: int | float) -> int | float:
    """Return the square root of `n`."""
    return math.sqrt(n)

@serve.tool
def isqrt(n: int | float) -> int:
    """Return the integer square root of `n`."""
    return math.isqrt(n)

@serve.tool
def sin(n: int | float) -> int | float:
    """Return the sine of `n`."""
    return math.sin(n)

@serve.tool
def cos(n: int | float) -> int | float:
    """Return the cosine of `n`."""
    return math.cos(n)

@serve.tool
def tan(n: int | float) -> int | float:
    """Return the tangent of `n`."""
    return math.tan(n)

@serve.tool
def sinh(n: int | float) -> int | float:
    """Return the hyperbolic sine of `n`."""
    return math.sinh(n)

@serve.tool
def cosh(n: int | float) -> int | float:
    """Return the hyperbolic cosine of `n`."""
    return math.cosh(n)

@serve.tool
def tanh(n: int | float) -> int | float:
    """Return the hyperbolic tangent of `n`."""
    return math.tanh(n)

@serve.tool
def exp(n: int | float) -> int | float:
    """Return `e^n`."""
    return math.exp(n)

@serve.tool
def erf(n: int | float) -> int | float:
    """Return the result of the error function with parameter `n`."""
    return math.erf(n)

@serve.tool
def factorial(n: int | float) -> int | float:
    """Return the factorial of a number `n`."""
    return math.factorial(n)

@serve.tool
def is_prime(n: int) -> bool:
    """Return whether `n` is a prime number. Uses trial division up to sqrt(n), checking
    divisibility by 2, then only odd numbers. Leverages the cached `PRIMES` sequence for
    efficient repeated calls."""

    if not isinstance(n, int):
        raise TypeError("n must be an integer")
    n = abs(n)
    if n < 2:
        return False
    if n < 4: # 2, 3
        return True
    if n % 2 == 0:
        return False
    limit = math.isqrt(n)
    # check against cached primes first
    idx = 0
    while True:
        try:
            p = PRIMES[idx]
        except IndexError:
            # PRIMES shouldn't raise IndexError (infinite iterator), but be safe
            break
        if p > limit:
            break
        if n % p == 0:
            return False
        idx += 1
        # fallback: if PRIMES hasn't computed far enough, compute manually
        if idx >= _N_PREFETCH and p < limit:
            # iterate odd numbers from last cached prime + 2 up to limit
            candidate = PRIMES[_N_PREFETCH - 1] + 2
            while candidate <= limit:
                if n % candidate == 0:
                    return False
                candidate += 2
            break
    return True

@serve.tool
def eval_expr(expression: str) -> object:
    """Evaluate a mathematical expression string and return the result.

    Provides a namespace with `math` module functions and common constants
    (pi, e, tau, inf, nan) for convenience.

    Args:
        expression: A string containing a valid Python expression, e.g.
                    "sqrt(3**2 + 4**2)" or "sin(pi/4)".

    Returns:
        The result of evaluating the expression.

    Raises:
        Various exceptions on invalid expressions, undefined names, etc."""

    import math as _math

    # build a safe-ish namespace with common math bindings
    _namespace = {
        # constants
        "pi":      _math.pi,
        "e":       _math.e,
        "tau":     _math.tau,
        "inf":     float("inf"),
        "nan":     float("nan"),
        "True":    True,
        "False":   False,
        # math module functions
        "sqrt":    _math.sqrt,
        "isqrt":   _math.isqrt,
        "log":     _math.log,
        "log2":    _math.log2,
        "log10":   _math.log10,
        "exp":     _math.exp,
        "sin":     _math.sin,
        "cos":     _math.cos,
        "tan":     _math.tan,
        "asin":    _math.asin,
        "acos":    _math.acos,
        "atan":    _math.atan,
        "atan2":   _math.atan2,
        "sinh":    _math.sinh,
        "cosh":    _math.cosh,
        "tanh":    _math.tanh,
        "asinh":   _math.asinh,
        "acosh":   _math.acosh,
        "atanh":   _math.atanh,
        "degrees": _math.degrees,
        "radians": _math.radians,
        "floor":   _math.floor,
        "ceil":    _math.ceil,
        "trunc":   _math.trunc,
        "factorial": _math.factorial,
        "comb":    _math.comb,
        "perm":    _math.perm,
        "gcd":     _math.gcd,
        "lcm":     _math.lcm,
        "fabs":    _math.fabs,
        "fmod":    _math.fmod,
        "modf":    _math.modf,
        "frexp":   _math.frexp,
        "ldexp":   _math.ldexp,
        "pow":     _math.pow,
        "prod":    _math.prod,
        "dist":    _math.dist,
        "hypot":   _math.hypot,
        "erf":     _math.erf,
        "erfc":    _math.erfc,
        "gamma":   _math.gamma,
        "lgamma":  _math.lgamma,
        # builtins (limited set)
        "abs":     abs,
        "min":     min,
        "max":     max,
        "round":   round,
        "sum":     sum,
        "int":     int,
        "float":   float,
        "complex": complex,
        "bool":    bool,
        "len":     len,
        "range":   range,
        "list":    list,
        "tuple":   tuple,
        "dict":    dict,
        "set":     set,
        "sorted":  sorted,
        "reversed": reversed,
        "enumerate": enumerate,
        "zip":     zip,
        "map":     map,
        "filter":  filter,
        "any":     any,
        "all":     all,
        "isinstance": isinstance,
        "type":    type,
        "str":     str,
        "repr":    repr,
        "hex":     hex,
        "oct":     oct,
        "bin":     bin,
        "ord":     ord,
        "chr":     chr,
    }

    try:
        result = eval(expression, {"__builtins__": {}}, _namespace)
        return result
    except Exception as e:
        # re-raise with a clear message
        raise type(e)(f"eval_expr: error evaluating '{expression}': {e}") from e

#
# ANSI codes (primarily used in the logger)
#

class ANSI:
    """ANSI codes for terminal emulators"""

    # standard colors
    FG_BLACK   = '\x1b[30m'
    FG_RED     = '\x1b[31m'
    FG_GREEN   = '\x1b[32m'
    FG_YELLOW  = '\x1b[33m'
    FG_BLUE    = '\x1b[34m'
    FG_MAGENTA = '\x1b[35m'
    FG_CYAN    = '\x1b[36m'
    FG_WHITE   = '\x1b[37m'
    BG_BLACK   = '\x1b[40m'
    BG_RED     = '\x1b[41m'
    BG_GREEN   = '\x1b[42m'
    BG_YELLOW  = '\x1b[43m'
    BG_BLUE    = '\x1b[44m'
    BG_MAGENTA = '\x1b[45m'
    BG_CYAN    = '\x1b[46m'
    BG_WHITE   = '\x1b[47m'

    # bright colors
    FG_BRIGHT_BLACK   = '\x1b[90m'
    FG_BRIGHT_RED     = '\x1b[91m'
    FG_BRIGHT_GREEN   = '\x1b[92m'
    FG_BRIGHT_YELLOW  = '\x1b[93m'
    FG_BRIGHT_BLUE    = '\x1b[94m'
    FG_BRIGHT_MAGENTA = '\x1b[95m'
    FG_BRIGHT_CYAN    = '\x1b[96m'
    FG_BRIGHT_WHITE   = '\x1b[97m'
    BG_BRIGHT_BLACK   = '\x1b[100m'
    BG_BRIGHT_RED     = '\x1b[101m'
    BG_BRIGHT_GREEN   = '\x1b[102m'
    BG_BRIGHT_YELLOW  = '\x1b[103m'
    BG_BRIGHT_BLUE    = '\x1b[104m'
    BG_BRIGHT_MAGENTA = '\x1b[105m'
    BG_BRIGHT_CYAN    = '\x1b[106m'
    BG_BRIGHT_WHITE   = '\x1b[107m'

    # text modes
    MODE_RESET_ALL           = '\x1b[0m'
    MODE_BOLD_SET            = '\x1b[1m'
    MODE_DIM_SET             = '\x1b[2m'
    MODE_ITALIC_SET          = '\x1b[3m'
    MODE_UNDERLINE_SET       = '\x1b[4m'
    MODE_BLINKING_SET        = '\x1b[5m'
    MODE_REVERSE_SET         = '\x1b[7m'
    MODE_HIDDEN_SET          = '\x1b[8m'
    MODE_STRIKETHROUGH_SET   = '\x1b[9m'
    MODE_BOLD_RESET          = '\x1b[22m'
    MODE_DIM_RESET           = '\x1b[22m'
    MODE_ITALIC_RESET        = '\x1b[23m'
    MODE_UNDERLINE_RESET     = '\x1b[24m'
    MODE_BLINKING_RESET      = '\x1b[25m'
    MODE_REVERSE_RESET       = '\x1b[27m'
    MODE_HIDDEN_RESET        = '\x1b[28m'
    MODE_STRIKETHROUGH_RESET = '\x1b[29m'

    # special
    BELL             = '\a'
    TERMINAL_RESET   = '\x1bc'
    SCROLLBACK_CLEAR = '\x1b[3J'
    CLEAR            = TERMINAL_RESET + SCROLLBACK_CLEAR + MODE_RESET_ALL

#
# Logger
#

class Logger:

    _LogLvl = Literal[0,1,2,3]

    @staticmethod
    def _get_debug_lvl() -> _LogLvl:
        return 0

    @staticmethod
    def _get_info_lvl() -> _LogLvl:
        return 1

    @staticmethod
    def _get_warn_lvl() -> _LogLvl:
        return 2

    @staticmethod
    def _get_error_lvl() -> _LogLvl:
        return 3

    @staticmethod
    def _islvl(obj: object) -> bool:
        """Is this a valid log level?"""
        return obj in [0,1,2,3]

    @staticmethod
    def _get_prefix(lvl: _LogLvl) -> str:
        if lvl == 0:
            lvltxt = f"{ANSI.FG_BRIGHT_WHITE}  DEBUG"
        elif lvl == 1:
            lvltxt = f"{ANSI.FG_BRIGHT_GREEN}   INFO"
        elif lvl == 2:
            lvltxt = f"{ANSI.FG_BRIGHT_YELLOW}WARNING"
        elif lvl == 3:
            lvltxt = f"{ANSI.FG_BRIGHT_RED} ERROR"
        else:
            raise ValueError
        return (
            f"{ANSI.MODE_RESET_ALL}{ANSI.MODE_BOLD_SET}{ANSI.FG_BRIGHT_BLACK}"
            f"{timestamp()}{ANSI.MODE_RESET_ALL}{ANSI.MODE_BOLD_SET} {lvltxt}"
            f"{ANSI.MODE_RESET_ALL}{ANSI.MODE_BOLD_SET}{ANSI.FG_BRIGHT_BLACK}:"
            f"{ANSI.MODE_RESET_ALL}"
        )

    def __init__(
        self,
        stdout: Optional[object] = None,
        stderr: Optional[object] = None
    ):
        self.stdout = stdout if stdout is not None else sys.stdout
        self.stderr = stderr if stderr is not None else sys.stderr
        self.last_lvl: __class__._LogLvl = 1

    def __call__(self, lvl: _LogLvl, msg: str) -> None:
        _lvl = lvl if __class__._islvl(lvl) else self.last_lvl
        out = self.stdout if _lvl in [0,1] else self.stderr
        print(f"{self._get_prefix(_lvl)} {msg}", file=out)
        self.last_lvl = _lvl

    def debug(self, msg: str) -> None:
        self(0, msg)

    def info(self, msg: str) -> None:
        self(1, msg)

    def warn(self, msg: str) -> None:
        self(2, msg)

    def error(self, msg: str) -> None:
        self(3, msg)

    def get_str(self, lvl: _LogLvl, msg: str) -> str:
        return f"{self._get_prefix(lvl if __class__._islvl(lvl) else self.last_lvl)} {msg}"

_log = Logger()

#
# Utils
#

@serve.tool
def timestamp() -> str:
    """Return a text-based timestamp representing the current time."""
    return datetime.datetime.now().strftime("[%Y-%m-%d %a %H:%M:%S]")

def isnum(obj: object) -> bool:
    return not isinstance(obj, bool) and isinstance(obj, int | float)

def exception_to_str(exc: Exception) -> str:
    """ExceptionType: Exception message"""
    return f"{type(exc).__name__}: {exc}"

def clear() -> None:
    """Clear the terminal"""
    print(ANSI.CLEAR, end='', flush=True)

def cls() -> None: # same as `clear()` for convenience
    """Clear the terminal"""
    print(ANSI.CLEAR, end='', flush=True)

#
# Math functions
#

@serve.tool
def softmax(z: list[float | int], T: float = 1.0, axis: int = -1) -> np.ndarray:
    """Numerically stable softmax over a specified axis of an array.
    Supports temperature scaling by parameter `T`."""
    z_arr = np.array(z, dtype=_DTYPE_NP)
    if z_arr.size == 0:
        return z_arr
    if T == 0.0:
        result = np.zeros_like(z_arr)
        indices = np.argmax(z_arr, axis=axis)
        np.put_along_axis(result, np.expand_dims(indices, axis), 1.0, axis=axis)
        return result
    if T < 0:
        z_arr = -z_arr
        T = -T
    max_val = np.max(z_arr, axis=axis, keepdims=True)
    exp_vals = np.exp((z_arr - max_val) / T)
    return exp_vals / np.sum(exp_vals, axis=axis, keepdims=True)

@serve.tool
def softmax_n(z: list[float | int], n: int, T: float = 1.0, axis: int = -1) -> np.ndarray:
    """Apply softmax `n` times."""
    for _ in range(n):
        z = softmax(z=z, T=T, axis=axis)
    return z

@serve.tool
def GEOMEAN(a: float | int, b: float | int) -> float:
    """Return the geometric mean of two numbers."""
    if not (isnum(a) and isnum(b)):
        raise TypeError
    if a == 0 or b == 0:
        return 0.0
    return math.exp((math.log(a) + math.log(b)) * 0.5)

@serve.tool
def geomean(arr: list[float | int]) -> float:
    """Calculates the geometric mean of an array."""
    if not isinstance(arr, np.ndarray):
        arr = np.array(arr, dtype=_DTYPE_NP)
    else:
        arr = arr.astype(_DTYPE_NP)
    if arr.size == 0:
        raise ValueError('array size is 0')
    if np.any(arr < 0):
        return _DTYPE_NP(np.nan)
    if np.any(arr == 0):
        return _DTYPE_NP(0.0)
    return float(np.exp(np.mean(np.log(arr))))

@serve.tool
def factors(n: int) -> list[int]:
    """Return the list of all factors of the absolute value of an integer `n`."""
    if not isinstance(n, int):
        raise TypeError
    _n = abs(n)
    factors: set[int] = set()
    for i in range(1, math.isqrt(_n) + 1):
        if _n % i == 0:
            factors.add(i)
            factors.add(_n // i)
    return sorted(factors)

@serve.tool
def GCD(a: int, b: int):
    """Return the greatest common divisor of two integers using an optimized Euclidean algorithm."""
    a = abs(a)
    b = abs(b)
    while (b != 0):
        a, b = b, a % b
    return a

@serve.tool
def LCM(a: int, b: int) -> int:
    """Return the least common multiple of two integers using the GCD relationship."""
    if a == 0 or b == 0:
        return 0
    return abs(a) // GCD(a, b) * abs(b)

@serve.tool
def gcd(nums: list[int]) -> int:
    """Return the greatest common divisor of a list of integers by iteratively applying `GCD()`."""
    if not nums:
        raise ValueError("nums is empty")
    if len(nums) == 1:
        return abs(nums[0])
    res = abs(nums[0])
    for i in range(1, len(nums)):
        res = GCD(res, abs(nums[i]))
    return res

@serve.tool
def lcm(nums: list[int]) -> int:
    """Return the least common multiple of a list of integers by iteratively applying `LCM()`."""
    if not nums:
        raise ValueError("nums is empty")
    res = abs(nums[0])
    for i in range(1, len(nums)):
        res = LCM(res, abs(nums[i]))
    return res

@serve.tool
def cofactors(nums: list[int]) -> list[int]:
    """Return the list of factors common to all integers in the given list."""
    if not nums:
        return []
    return factors(gcd(nums))

@serve.tool
def surprisal(p: float, base: Optional[int | float] = None) -> float:
    """Given a random event with probability `p`, calculate the surprisal
    (Shannon information content), as measured in nats (base `e`) by default,
    or as measured using the specified `base`."""
    if p > 1.0:
        _log.warn(f"p = {p} > 1.0")
    return log(1 / p, base) if isnum(base) else log(1 / p)

@serve.tool
def percent(a: int | float, b: int | float) -> float:
    """What percent is `a` of `b`?"""
    return (a / b) * 100

@serve.tool
def uniform(size: int, low: int | float, high: int | float) -> list:
    """Return a uniform random distribution of floats in the range [low, high) of a given size."""
    if size < 0:
        raise ValueError("size must be non-negative")
    return np.random.uniform(low, high, size).astype(_DTYPE_NP).tolist()

@serve.tool
def binary_decomposition(n: int) -> list[int]:
    """Return a list of integers representing the binary decomposition (factorization) of an integer.

    Examples:

        `binary_decomposition(0)`     -> `[]`        because                 == 0
        `binary_decomposition(1)`     -> `[0]`       because 2^0             == 1
        `binary_decomposition(4)`     -> `[2]`       because 2^2             == 4
        `binary_decomposition(9)`     -> `[3, 0]`    because 2^3 + 2^0       == 9
        `binary_decomposition(100)`   -> `[6, 5, 2]` because 2^6 + 2^5 + 2^2 == 100
        `binary_decomposition(81920)` -> `[16, 14]`  because 2^16 + 2^14     == 81920"""

    if not isinstance(n, int):
        raise TypeError("only supports integer input")
    if n < 0:
        raise ValueError("does not support negative input")
    decomp = []
    if n == 0:
        return decomp
    p = 0
    while n > 0:
        if n % 2:
            decomp.append(p)
        n //= 2
        p += 1
    return decomp[::-1]

@serve.tool
def binary_composition(nums: list[int]) -> int:
    """Reconstruct an integer from its binary decomposition (inverse of `binary_decomposition`).

    Given a list of exponents representing powers of two, returns their sum.

    Examples:

        `binary_composition([])`      -> `0`   because 0
        `binary_composition([0])`     -> `1`   because 2^0
        `binary_composition([2])`     -> `4`   because 2^2
        `binary_composition([3, 0])`  -> `9`   because 2^3 + 2^0
        `binary_composition([6, 5, 2])` -> `100` because 2^6 + 2^5 + 2^2"""

    res = 0
    for n in nums:
        res += 2**n
    return res

@serve.tool
def decomposition(n: int, base: int) -> list[int]:
    """Return the decomposition of a non-negative integer `n` into powers of a given `base`.

    Each element in the returned list is an exponent such that the sum of `base^exponent`
    over all elements equals `n`. This is a generalisation of `binary_decomposition`
    to arbitrary bases.

    Examples:

        `decomposition(0, 2)`      -> `[]`           because 0
        `decomposition(5, 2)`      -> `[2, 0]`       because 2^2 + 2^0 == 5
        `decomposition(5, 3)`      -> `[1, 1, 0]`    because 3^1 + 3^1 + 3^0 == 5
        `decomposition(42, 10)`    -> `[1, 0, 0, 0]` because 10^1 + 10^0 + 10^0 + 10^0 == 42"""

    if not isinstance(n, int):
        raise TypeError("n must be an integer")
    if n < 0:
        raise ValueError("n must be non-negative")
    if not isinstance(base, int) or base < 2:
        raise ValueError("base must be an integer >= 2")
    decomp = []
    if n == 0:
        return decomp
    exp = 0
    while n > 0:
        n, digit = divmod(n, base)
        decomp.extend([exp] * digit)
        exp += 1
    return decomp[::-1]

#
# Misc. functions
#

@serve.tool
def random_hex(n: int) -> str:
    """Generates a random hexadecimal string representing `n` bytes."""
    if not isinstance(n, int) or n <= 0:
        raise ValueError
    return os.urandom(n).hex()

@serve.tool
def gen_key(n_per_part: int = 4, n_parts: int = 8) -> str:
    """Generate a random alphanumeric key in the format `XXXX-XXXX-...-XXXX`.

    Each part consists of `n_per_part` uppercase ASCII letters (A-Z) generated
    from cryptographically secure random bytes. The parts are joined with hyphens.

    Args:
        n_per_part: Number of characters per part (default 4).
        n_parts:    Number of parts separated by hyphens (default 8).

    Returns:
        A string like `"ABCD-EFGH-IJKL-MNOP"`."""

    result = []
    for _ in range(n_parts):
        random_bytes = os.urandom(n_per_part)
        chars = []
        for byte in random_bytes:
            chars.append(chr((byte & 0x1F) % 26 + 65))
        result.append(''.join(chars))
    return '-'.join(result)

@serve.tool
def kv_cache_size_bytes(
    n_head_kv:      int,
    n_head_dim:     int,
    n_layer:        int,
    n_ctx:          int,
    bits_per_key:   int,
    bits_per_value: int
) -> int:
    """Calculate the total size (in bytes) of the KV cache for a transformer model.

    The KV cache stores the key and value tensors for each layer, head, token, and
    dimension, quantised to the specified bit widths.

    Args:
        n_head_kv:      Number of key/value heads.
        n_head_dim:     Dimension (size) of each head.
        n_layer:        Number of transformer layers.
        n_ctx:          Maximum context length (number of tokens).
        bits_per_key:   Number of bits used to store each key element.
        bits_per_value: Number of bits used to store each value element.

    Returns:
        Total KV cache size in bytes.

    Formula:
        bytes = n_head_kv * n_head_dim * n_layer * n_ctx * (bits_per_key + bits_per_value) / 8"""

    return int(n_head_kv * n_head_dim * n_layer * n_ctx * (bits_per_key + bits_per_value) / 8)

@serve.tool
def is_close(a: int | float, b: int | float, factor: float = 0.995) -> bool:
    """Return whether `a` and `b` are close to one another (symmetric).
    `factor` should be a float in the range [0.0, 1.0] (inclusive)."""
    if not (0.0 <= factor <= 1.0):
        raise ValueError(f"factor must be in range [0, 1] (got {factor})")
    return factor <= a/b <= 2 - factor

@serve.tool
def seconds_from_now(seconds: int):
    """Return a `datetime` object offset from the current time by the given number of seconds."""
    return datetime.datetime.now() + datetime.timedelta(seconds=seconds)

# ------------------------------------------------------------------------------------------------
#
# Web search, fetch, etc.
#
# ------------------------------------------------------------------------------------------------

@dataclass
class SearchConfig:
    """Settings for the web-search backend."""
    provider: str = "duckduckgo"
    user_agent: str = _LEGACY_UA
    region: str = "us-en"
    max_results: int = 25
    timeout_seconds: int = 20

@dataclass
class FetchConfig:
    """Settings for the page-fetch tool."""
    browser_profile: str = "chrome136"
    max_content_bytes: int = 16384
    timeout_seconds: int = 20

#
# Exceptions
#

class SearchError(Exception):
    """A caller-facing problem with a search request (empty query, rate-limit, etc.)."""

class FetchError(Exception):
    """A caller-facing problem with a page-fetch request (bad URL, unreachable, etc.)."""

#
# Search
#

def _unwrap_ddg_url(href: str) -> str:
    """Extract the real destination from a DuckDuckGo redirect URL.

    DDG wraps every result href in::

        //duckduckgo.com/l/?uddg=<encoded>&rut=...

    Pull out and decode the ``uddg`` parameter.  Also handles bare
    protocol-relative URLs and already-absolute URLs."""

    if not href:
        return ""

    # DDG redirect:  //duckduckgo.com/l/?uddg=https%3A%2F%2F...&rut=...
    m = re.search(r"[?&]uddg=([^&]+)", href)
    if m:
        return unquote(m.group(1))

    # protocol-relative:  //example.com/path
    if href.startswith("//"):
        return "https:" + href

    # already a normal URL
    return href

def _parse_ddg_results(html: str, count: int) -> list[dict]:
    """Parse DuckDuckGo's server-rendered HTML into a list of result dicts.

    Each result dict has keys: *rank*, *title*, *url*, *displayUrl*, *snippet*.
    Sponsored (``result--ad``) entries are skipped."""

    soup = BeautifulSoup(html, "lxml")
    results: list[dict] = []

    for body in soup.find_all("div", class_="result__body"):
        if len(results) >= count:
            break

        # skip sponsored blocks
        container = body.parent
        if container and _has_class(container, "result--ad"):
            continue

        # title anchor
        anchor = body.find("a", class_="result__a")
        if not anchor:
            continue

        title = anchor.get_text(strip=True)
        url = _unwrap_ddg_url(anchor.get("href", ""))
        if not title or not url:
            continue

        # snippet
        snippet_el = body.find("a", class_="result__snippet")
        snippet = snippet_el.get_text(strip=True) if snippet_el else ""

        # display URL
        url_el = body.find("a", class_="result__url")
        display_url = url_el.get_text(strip=True) if url_el else None
        if display_url and not display_url.strip():
            display_url = None

        results.append({
            "rank": len(results) + 1,
            "title": title,
            "url": url,
            "displayUrl": display_url,
            "snippet": snippet,
        })

    return results

def _has_class(tag, class_name: str) -> bool:
    """Check whether a BeautifulSoup tag has *class_name* among its classes."""
    classes = tag.get("class")
    return bool(classes) and class_name in classes

def search_web(query: str, count: int = 10, config: Optional[SearchConfig] = None) -> str:
    """Search the web via DuckDuckGo's no-JavaScript HTML endpoint.

    Returns a JSON string.  On success the top-level keys are:

        success, provider, query, count, results

    where *results* is a list of dicts with keys *rank*, *title*, *url*,
    *displayUrl*, *snippet*.

    On error the top-level keys are:

        success (``false``), error (description string)

    Args:
        query:  The search query.
        count:  Max results to return (clamped to 1-25, default 10).
        config: Optional overrides for user-agent, region, timeout, etc."""

    if config is None:
        config = SearchConfig()

    query = query.strip()
    if not query:
        return {"success": False, "error": "search query must not be empty"}

    count = max(1, min(count, config.max_results))
    url = f"{_DDG_ENDPOINT}?q={quote(query)}&kl={quote(config.region)}"
    headers = {"User-Agent": config.user_agent}

    try:
        resp = _requests.get(url, headers=headers, timeout=config.timeout_seconds)
        resp.raise_for_status()
        html = resp.text
    except Exception as exc:
        return {"success": False, "error": f"failed to fetch search results: {exc}"}

    # DDG rate-limit detection
    lower = html.lower()
    if "anomaly" in lower or "if this error persists" in lower:
        return {
            "success": False,
            "error": "DuckDuckGo temporarily rate-limited this request; wait a moment and retry",
        }

    try:
        results = _parse_ddg_results(html, count)
    except Exception as exc:
        return {"success": False, "error": f"Failed to parse search results: {exc}"}

    return {
        "success": True,
        "provider": "duckduckgo",
        "query": query,
        "count": len(results),
        "results": results,
    }

def _strip_html(html: str) -> str:
    """Last-resort text extraction: strip non-content nodes, keep inner text.

    Used when trafilatura (Readability) fails to find a coherent article."""

    soup = BeautifulSoup(html, "lxml")

    for tag in soup(
        ("script", "style", "noscript", "nav", "header", "footer", "svg", "form")
    ):
        tag.decompose()

    text = soup.get_text(separator=" ", strip=True)
    text = re.sub(r"\n[ \t\f\v]*(\n[ \t\f\v]*)+", "\n\n", text)
    text = re.sub(r"[ \t\f\v]+", " ", text)
    return text.strip()

def _fetch_url(url: str, config: FetchConfig) -> str:
    """Fetch `url` with browser impersonation, falling back to plain requests."""

    if HAS_CURL_CFFI:
        try:
            resp = _curl_requests.get(
                url,
                impersonate=config.browser_profile,
                timeout=config.timeout_seconds,
                allow_redirects=True,
            )
            resp.raise_for_status()
            return resp.text
        except Exception as exc:
            _log.warning("curl_cffi request failed (%s); falling back to stdlib", exc)

    # fallback with a modern browser UA
    headers = {
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/136.0.0.0 Safari/537.36"
        ),
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        "Accept-Language": "en-US,en;q=0.5",
    }
    resp = _requests.get(
        url, headers=headers, timeout=config.timeout_seconds, allow_redirects=True
    )
    resp.raise_for_status()
    return resp.text


def fetch_page(
    url: str, max_bytes: Optional[int] = None, config: Optional[FetchConfig] = None
) -> dict[str, object]:
    """Fetch a web page and return its main readable text with metadata.

    Uses ``curl_cffi`` for browser impersonation to bypass bot detection,
    then extracts the main article content via Readability (trafilatura),
    falling back to a plain HTML strip if extraction fails.

    Returns a JSON string.  On success the top-level keys are:

        success, url, title, author, siteName, published, language,
        minutesToRead, readable, length, truncated, text

    On error the top-level keys are:

        success (``false``), error (description string)

    Args:
        url:       The page URL to fetch (http or https).
        max_bytes: Max bytes of body text to return (default: 16_384).
                   Longer content is truncated with a ``[… truncated]`` marker.
        config:    Optional overrides for browser profile, timeout, etc."""

    if config is None:
        config = FetchConfig()

    cap = max_bytes if (max_bytes is not None and max_bytes > 0) else config.max_content_bytes

    # validate URL
    url = url.strip()
    if not url:
        return {"success": False, "error": "a page URL is required"}
    if not (url.startswith("http://") or url.startswith("https://")):
        return {"success": False, "error": f"{repr(url)} is not a valid HTTP(S) URL"}

    # fetch
    try:
        html = _fetch_url(url, config)
    except Exception as exc:
        return {"success": False, "error": f"failed to fetch {repr(url)}: {exc}"}

    # extract readable content
    title = author = site_name = language = published = None
    minutes_to_read = None
    readable = False
    text: Optional[str] = None

    if HAS_TRAFILATURA:
        print("has trafilatura")
        try:
            result = _trafilatura.bare_extraction(
                html,
                with_metadata=True,
                include_comments=False,
                include_tables=True,
                include_images=False,
                include_formatting=False,
                no_fallback=False,
            )
            if result:
                # support both dict-like and attribute access (varies by trafilatura version)
                def _get(obj, key):
                    return obj[key] if isinstance(obj, dict) else getattr(obj, key, None)

                text = _get(result, "text")
                title = _get(result, "title")
                author = _get(result, "author")
                site_name = _get(result, "site_name") or _get(result, "source")
                published = _get(result, "date")
                language = _get(result, "language")

                if text and text.strip():
                    readable = True
                    word_count = len(text.split())
                    minutes_to_read = max(1, word_count // 200)
        except Exception as exc:
            _log.warning("trafilatura extraction failed (%s)", exc)
    else:
        print("no trafilatura")

    # fallback to text extraction
    if not text or not text.strip():
        text = _strip_html(html)
        if not text.strip():
            return {
                "success": False, "error": "no readable content could be extracted from the page"
            }

    # truncate if needed
    truncated = False
    encoded = text.encode("utf-8")
    if len(encoded) > cap:
        encoded = encoded[:cap]
        text = encoded.decode("utf-8", errors="ignore").rstrip() + " [... truncated ...]"
        truncated = True

    return {
        "success": True,
        "url": url,
        "title": title,
        "author": author,
        "siteName": site_name,
        "published": published,
        "language": language,
        "minutesToRead": minutes_to_read,
        "readable": readable,
        "length": len(text.encode("utf-8")),
        "truncated": truncated,
        "text": text,
    }

@serve.tool
def web_search(query: str, count: int = 10) -> dict:
    """Search the web via DuckDuckGo. Returns JSON with titles, URLs, snippets."""
    return search_web(query, count)

@serve.tool
def web_fetch(url: str, max_bytes: int = 16384) -> dict:
    """Fetch a web page and extract its main readable content. Returns JSON."""
    return fetch_page(url, max_bytes)

if __name__ == "__main__":
    serve.run("http")