xmlriver-mcp-toolkit 0.1.0__py3-none-any.whl

xmlriver_mcp/__init__.py

@@ -0,0 +1,3 @@
+"""Independent XMLRiver client and MCP server."""
+
+__version__ = "0.1.0"

xmlriver_mcp/batch.py

@@ -0,0 +1,190 @@
+import asyncio
+import time
+
+from .client import XmlRiverClient, resolve_xmlriver_creds
+from .config import Config
+from .errors import XmlRiverError
+from .models import DEFAULT_REGION, SerpArgs
+from .parsers import find_matches, normalize_domain
+from .services import collect_serp, utcnow
+from .storage import Storage, make_run_id, normalize_phrase
+
+
+def build_tasks(cfg: Config, phrases: list[str], regional: bool = False) -> list[dict]:
+    tasks = []
+    seen = set()
+    for phrase in phrases:
+        norm = normalize_phrase(phrase)
+        if not norm or norm in seen:
+            continue
+        seen.add(norm)
+        for engine in cfg.engines:
+            regions: list[tuple[int, list[str]]] = [(engine.region, [])]
+            if regional:
+                regions = [
+                    (getattr(region, engine.name), region.expected_hosts)
+                    for region in cfg.regions
+                    if getattr(region, engine.name) is not None
+                ]
+            for region, hosts in regions:
+                args = SerpArgs.model_validate(
+                    dict(
+                        query=phrase,
+                        engine=engine.name,
+                        region=region,
+                        device=engine.device,
+                        pages=cfg.depth // 10,
+                        target_domain=cfg.target_domain,
+                        include_subdomains=cfg.include_subdomains,
+                    )
+                )
+                tasks.append({**args.model_dump(), "expected_hosts": hosts})
+    if not tasks:
+        raise ValueError("No collection tasks")
+    return tasks
+
+
+async def run_batch(
+    cfg: Config,
+    tasks: list[dict],
+    store: Storage,
+    *,
+    dry_run: bool = False,
+    refresh: bool = False,
+    client=None,
+) -> dict:
+    started = time.monotonic()
+    records = []
+    pending = []
+    for task in tasks:
+        task["region"] = (
+            task.get("region") if task.get("region") is not None else DEFAULT_REGION[task["engine"]]
+        )
+        cached = None if refresh else store.cached(task, cfg.cache_ttl_hours)
+        if cached:
+            cached["expected_hosts"] = task.get("expected_hosts", [])
+            cached["expected_host_matches"] = [
+                doc
+                for doc in cached["results"]
+                if normalize_domain(doc["url"])
+                in {normalize_domain(h) for h in cached["expected_hosts"]}
+            ]
+            records.append(cached)
+        else:
+            pending.append(task)
+    estimate = (
+        None
+        if any(task["engine"] not in cfg.costs_per_1000 for task in pending)
+        else sum(task["pages"] * cfg.costs_per_1000[task["engine"]] / 1000 for task in pending)
+    )
+    if dry_run:
+        return dict(
+            status="dry_run",
+            tasks=len(tasks),
+            cached_hits=len(records),
+            successful_pages=0,
+            http_attempts=0,
+            estimated_cost_rub=estimate,
+            pages_without_retries=sum(task["pages"] for task in pending),
+            max_attempts=cfg.max_attempts,
+            network_requests=False,
+        )
+    manifest = dict(
+        run_id=make_run_id(),
+        created_at=utcnow(),
+        status="running",
+        http_attempts=0,
+        successful_pages=0,
+        estimated_cost_rub=None,
+        elapsed_seconds=0,
+        tasks=len(tasks),
+        cached_hits=len(records),
+    )
+    store.save_run(manifest, records)
+    own_client = client is None
+    client = client or XmlRiverClient(
+        *resolve_xmlriver_creds(),
+        rate_max=cfg.rate_max,
+        rate_window=cfg.rate_window,
+        max_attempts=cfg.max_attempts,
+    )
+    attempt_start = client.http_attempts
+    if own_client:
+        await client.__aenter__()
+    status = "ok"
+    try:
+        for task in pending:
+            progress = {
+                **task,
+                "phrase": task["query"],
+                "phrase_norm": normalize_phrase(task["query"]),
+            }
+            attempts_before = client.http_attempts
+            try:
+                async with asyncio.timeout(cfg.deadline):
+                    await collect_serp(
+                        client,
+                        SerpArgs.model_validate(
+                            {k: v for k, v in task.items() if k != "expected_hosts"}
+                        ),
+                        progress=progress,
+                    )
+            except (XmlRiverError, TimeoutError) as exc:
+                error = exc if isinstance(exc, XmlRiverError) else XmlRiverError(kind="deadline")
+                progress.update(status="error", complete=False, error=error.payload())
+                status = "partial"
+                if error.kind in {"auth", "balance", "ip", "budget"}:
+                    status = "aborted"
+            except asyncio.CancelledError:
+                status = "aborted"
+                progress.update(status="incomplete", complete=False)
+                raise
+            finally:
+                progress["http_attempts"] = client.http_attempts - attempts_before
+                matches = find_matches(
+                    progress.get("results", []), task["target_domain"], task["include_subdomains"]
+                )
+                progress.update(
+                    matches=matches,
+                    position=matches[0]["position"] if matches else None,
+                    url=matches[0]["url"] if matches else None,
+                )
+                if progress.get("complete"):
+                    progress["status"] = "found" if matches else "not_found"
+                progress["expected_host_matches"] = [
+                    doc
+                    for doc in progress.get("results", [])
+                    if normalize_domain(doc["url"])
+                    in {normalize_domain(h) for h in task.get("expected_hosts", [])}
+                ]
+                price = cfg.costs_per_1000.get(task["engine"])
+                progress["estimated_cost_rub"] = (
+                    None if price is None else progress["http_attempts"] * price / 1000
+                )
+                records.append(progress)
+                store.save_cache(progress)
+                manifest.update(
+                    http_attempts=client.http_attempts - attempt_start,
+                    elapsed_seconds=time.monotonic() - started,
+                )
+                store.save_run(manifest, records)
+            if status == "aborted":
+                break
+    finally:
+        if own_client:
+            await client.__aexit__()
+        fresh = [r for r in records if r.get("source") == "api"]
+        costs = [r.get("estimated_cost_rub") for r in fresh]
+        manifest.update(
+            status=status,
+            http_attempts=client.http_attempts - attempt_start,
+            successful_pages=sum(r.get("successful_pages", 0) for r in fresh),
+            estimated_cost_rub=None
+            if any(v is None for v in costs)
+            else sum(v for v in costs if v is not None),
+            elapsed_seconds=time.monotonic() - started,
+            completed_tasks=sum(bool(r.get("complete")) for r in records),
+        )
+        store.save_run(manifest, records)
+        store.append_history(records, manifest)
+    return manifest

xmlriver_mcp/cli.py

@@ -0,0 +1,117 @@
+import argparse
+import asyncio
+import json
+import logging
+import sys
+from pathlib import Path
+
+import yaml
+from pydantic import ValidationError
+
+from .batch import build_tasks, run_batch
+from .client import XmlRiverClient, resolve_xmlriver_creds
+from .config import load_config
+from .csv_io import read_phrases, update_csv
+from .errors import XmlRiverError
+from .models import INPUTS
+from .regions import validate_regions
+from .services import execute
+from .storage import Storage
+
+
+def parser():
+    ap = argparse.ArgumentParser(
+        description="Independent XMLRiver toolkit. API requests may cost money."
+    )
+    commands = ap.add_subparsers(dest="command", required=True)
+    commands.add_parser("balance")
+    serp = commands.add_parser("serp")
+    serp.add_argument("query")
+    serp.add_argument("--engine", choices=["yandex", "google"], default="yandex")
+    serp.add_argument("--region", type=int)
+    serp.add_argument("--device", choices=["desktop", "mobile", "tablet"], default="mobile")
+    serp.add_argument("--pages", type=int, default=1)
+    serp.add_argument("--target-domain")
+    serp.add_argument("--include-subdomains", action=argparse.BooleanOptionalAction, default=True)
+    ws = commands.add_parser("wordstat").add_subparsers(dest="action", required=True)
+    for name in ("top", "history"):
+        command = ws.add_parser(name)
+        command.add_argument("query")
+        command.add_argument("--regions", type=int, default=213)
+        command.add_argument("--device", choices=["desktop", "phone", "tablet"])
+        if name == "top":
+            command.add_argument("--limit", type=int, default=50)
+            command.add_argument("--min-frequency", type=int, default=0)
+            command.add_argument("--include-associations", action="store_true")
+            command.add_argument("--drop-self", action=argparse.BooleanOptionalAction, default=True)
+        else:
+            command.add_argument("--period", choices=["day", "week", "month"], default="month")
+            command.add_argument("--start")
+            command.add_argument("--end")
+    for group, actions in (("positions", ("run", "update-csv")), ("regions", ("run", "validate"))):
+        sub = commands.add_parser(group).add_subparsers(dest="action", required=True)
+        for action in actions:
+            command = sub.add_parser(action)
+            command.add_argument("--config", required=True, type=Path)
+            command.add_argument("--workspace", type=Path)
+            command.add_argument("--dry-run", action="store_true")
+            command.add_argument("--refresh", action="store_true")
+            command.add_argument("--csv", type=Path)
+            if action == "update-csv":
+                command.add_argument("--run", required=True)
+    return ap
+
+
+async def run(args):
+    if args.command in ("positions", "regions"):
+        cfg = load_config(args.config)
+        store = Storage(args.workspace or args.config.resolve().parent / "xmlriver-data")
+        path = args.csv or (Path(cfg.csv) if cfg.csv else None)
+        if args.action == "update-csv":
+            if path is None:
+                raise ValueError("CSV path required")
+            return update_csv(store, path, args.run, cfg.column, dry_run=args.dry_run)
+        if args.action == "validate":
+            async with XmlRiverClient(
+                *resolve_xmlriver_creds(), max_attempts=cfg.max_attempts
+            ) as client:
+                return await validate_regions(
+                    cfg, store, client, refresh=args.refresh, dry_run=args.dry_run
+                )
+        phrases = read_phrases(path, cfg.column) if path else cfg.phrases
+        tasks = build_tasks(cfg, phrases, regional=args.command == "regions")
+        return await run_batch(cfg, tasks, store, dry_run=args.dry_run, refresh=args.refresh)
+    name = {
+        "balance": "get_balance",
+        "serp": "serp",
+        "wordstat": "wordstat_" + getattr(args, "action", ""),
+    }[args.command]
+    data = {k: v for k, v in vars(args).items() if k not in {"command", "action"}}
+    INPUTS[name].model_validate(data)
+    async with XmlRiverClient(*resolve_xmlriver_creds()) as client:
+        return await execute(client, name, data)
+
+
+def main():
+    logging.getLogger("httpx").disabled = True
+    logging.getLogger("httpcore").disabled = True
+    args = parser().parse_args()
+    try:
+        result = asyncio.run(run(args))
+        code = 4 if result.get("status") in {"partial", "aborted"} else 0
+    except (ValueError, ValidationError, yaml.YAMLError):
+        result, code = XmlRiverError(kind="params").payload(), 2
+    except XmlRiverError as exc:
+        result, code = exc.payload(), 3
+    except KeyboardInterrupt:
+        result, code = {"status": "aborted"}, 4
+    except Exception:
+        result, code = XmlRiverError().payload(), 3
+    if code:
+        print("XMLRiver operation did not complete successfully.", file=sys.stderr)
+    print(json.dumps(result, ensure_ascii=False))
+    raise SystemExit(code)
+
+
+if __name__ == "__main__":
+    main()

xmlriver_mcp/client.py

@@ -0,0 +1,278 @@
+"""Reusable async client. Credentials come only from arguments or environment."""
+
+import asyncio
+import json
+import logging
+import math
+import os
+import random
+import re
+import time
+from collections import deque
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+
+import httpx
+
+from .errors import XmlRiverError
+from .models import Engine, HistoryArgs, SerpArgs, WordstatArgs
+from .parsers import extract_error, parse_serp, parse_wordstat
+
+GOOGLE_URL = "https://xmlriver.com/search/xml"
+YANDEX_URL = "https://xmlriver.com/search_yandex/xml"
+WORDSTAT_URL = "https://xmlriver.com/wordstat/new/json"
+BALANCE_URL = "https://xmlriver.com/api/get_balance/"
+COST_URL = "https://xmlriver.com/api/get_cost/{system}/"
+DEFAULT_TIMEOUT = 45.0
+CREDS_HELP = "Set XMLRIVER_USER and XMLRIVER_KEY in the environment."
+TEMP_CODES = {"110", "111", "115", "201", "202", "500"}
+WS_TEMP_CODES = {"101", "110", "115", "500"}
+NO_RESULTS_CODES = {"15"}
+
+
+class AuthenticationFilter(logging.Filter):
+    """Redact authenticated URLs even when an embedding application enables HTTP logs."""
+
+    def filter(self, record):
+        message = record.getMessage()
+        record.msg = re.sub(r"""(?i)([?&](?:key|user)=)[^&\s'"\\]+""", r"\1[redacted]", message)
+        record.args = ()
+        return True
+
+
+_auth_filter = AuthenticationFilter()
+for _logger in (
+    "httpx",
+    "httpcore.connection",
+    "httpcore.http11",
+    "httpcore.http2",
+    "httpcore.proxy",
+    "httpcore.socks",
+):
+    logging.getLogger(_logger).addFilter(_auth_filter)
+
+
+def classify_error_code(code: str) -> str:
+    return {
+        "31": "auth",
+        "42": "auth",
+        "45": "ip",
+        "200": "balance",
+        "2": "params",
+        "121": "params",
+        "400": "params",
+    }.get(code, "temp" if code in TEMP_CODES else "other")
+
+
+def resolve_xmlriver_creds() -> tuple[str, str]:
+    return os.environ.get("XMLRIVER_USER", "").strip(), os.environ.get("XMLRIVER_KEY", "").strip()
+
+
+class RateLimiter:
+    def __init__(self, max_requests: int = 5, window_seconds: float = 1.5):
+        if max_requests < 1 or window_seconds <= 0:
+            raise ValueError("Invalid rate limit")
+        self.maximum, self.window = max_requests, window_seconds
+        self.times: deque[float] = deque()
+        self.lock = asyncio.Lock()
+
+    async def acquire(self):
+        async with self.lock:
+            while True:
+                now = time.monotonic()
+                while self.times and self.times[0] <= now - self.window:
+                    self.times.popleft()
+                if len(self.times) < self.maximum:
+                    self.times.append(now)
+                    return
+                await asyncio.sleep(max(0, self.window - (now - self.times[0])))
+
+
+def retry_after(value: str | None) -> float | None:
+    if not value:
+        return None
+    try:
+        seconds = float(value)
+    except ValueError:
+        try:
+            date = parsedate_to_datetime(value)
+            seconds = (date - datetime.now(timezone.utc)).total_seconds()
+        except (ValueError, TypeError, OverflowError):
+            return None
+    return max(0, seconds) if math.isfinite(seconds) else None
+
+
+class XmlRiverClient:
+    def __init__(
+        self,
+        user: str = "",
+        key: str = "",
+        *,
+        timeout: float = DEFAULT_TIMEOUT,
+        rate_max: int = 5,
+        rate_window: float = 1.5,
+        max_concurrent: int = 5,
+        max_attempts: int | None = None,
+        transport=None,
+    ):
+        if timeout <= 0 or max_concurrent < 1 or (max_attempts is not None and max_attempts < 0):
+            raise ValueError("Invalid client limits")
+        self._user, self._key = user, key
+        self._timeout, self._transport = timeout, transport
+        self._limiter = RateLimiter(rate_max, rate_window)
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+        self._http: httpx.AsyncClient | None = None
+        self.max_attempts = max_attempts
+        self.http_attempts = 0
+
+    async def __aenter__(self):
+        self._http = httpx.AsyncClient(timeout=self._timeout, transport=self._transport)
+        return self
+
+    async def __aexit__(self, *exc):
+        if self._http:
+            await self._http.aclose()
+            self._http = None
+
+    async def _request(self, url: str, params: dict, parser, *, max_retries: int = 3):
+        if not self._user or not self._key:
+            raise XmlRiverError(kind="auth")
+        if self._http is None:
+            raise RuntimeError("Use XmlRiverClient inside async with")
+        if not 0 <= max_retries <= 10:
+            raise ValueError("max_retries must be 0..10")
+        for attempt in range(max_retries + 1):
+            delay = None
+            async with self._semaphore:
+                await self._limiter.acquire()
+                if self.max_attempts is not None and self.http_attempts >= self.max_attempts:
+                    raise XmlRiverError(kind="budget")
+                self.http_attempts += 1
+                try:
+                    response = await self._http.get(
+                        url, params={**params, "user": self._user, "key": self._key}
+                    )
+                    status = response.status_code
+                    if status == 429 or 500 <= status < 600:
+                        delay = retry_after(response.headers.get("Retry-After"))
+                        raise XmlRiverError(code=str(status), kind="temp")
+                    if not 200 <= status < 300:
+                        raise XmlRiverError(
+                            code=str(status), kind="auth" if status in (401, 403) else "params"
+                        )
+                    return parser(response)
+                except httpx.HTTPError:
+                    error = XmlRiverError(kind="temp")
+                except XmlRiverError as exc:
+                    error = exc
+            if error.kind != "temp" or attempt == max_retries:
+                raise error from None
+            await asyncio.sleep(
+                delay if delay is not None else min(2**attempt, 15) + random.uniform(0, 0.25)
+            )
+        raise AssertionError("unreachable")
+
+    async def fetch_serp(
+        self,
+        query: str,
+        *,
+        engine: Engine,
+        region: int,
+        device: str = "mobile",
+        page_index: int = 0,
+        max_retries: int = 3,
+    ) -> str:
+        args = SerpArgs.model_validate(
+            dict(query=query, engine=engine, region=region, device=device)
+        )
+        if type(page_index) is not int or not 0 <= page_index < 10:
+            raise ValueError("page_index must be 0..9")
+        params = {
+            "query": args.query,
+            "device": device,
+            "loc" if engine == "google" else "lr": region,
+            "page": page_index + (1 if engine == "google" else 0),
+        }
+
+        def parse(response):
+            text = response.text
+            error = extract_error(text)
+            if error and error[0] != "15":
+                raise XmlRiverError(code=error[0], kind=classify_error_code(error[0]))
+            parse_serp(text)
+            return text
+
+        return await self._request(
+            GOOGLE_URL if engine == "google" else YANDEX_URL, params, parse, max_retries=max_retries
+        )
+
+    async def fetch_wordstat(
+        self,
+        query: str,
+        *,
+        regions: int | None = 213,
+        device: str | None = None,
+        pagetype: str = "words",
+        period: str | None = None,
+        start: str | None = None,
+        end: str | None = None,
+        max_retries: int = 3,
+    ) -> dict:
+        args = dict(query=query, regions=regions, device=device)
+        validated: WordstatArgs
+        if pagetype == "history":
+            args.update(period=period or "month", start=start, end=end)
+            validated = HistoryArgs.model_validate(args)
+        elif pagetype == "words":
+            validated = WordstatArgs.model_validate(args)
+        else:
+            raise ValueError("Unknown pagetype")
+
+        def parse(response):
+            try:
+                # JSON is decoded as UTF-8. Do not replace undecodable bytes silently.
+                data = json.loads(response.content.decode("utf-8-sig"))
+            except (ValueError, UnicodeError):
+                raise XmlRiverError(kind="format") from None
+            if not isinstance(data, dict):
+                raise XmlRiverError(kind="format")
+            if "error" in data:
+                code = str(data.get("code", ""))
+                raise XmlRiverError(
+                    code=code, kind="temp" if code in WS_TEMP_CODES else classify_error_code(code)
+                )
+            if pagetype == "words":
+                parse_wordstat(data, include_associations=True)
+            return data
+
+        return await self._request(
+            WORDSTAT_URL,
+            {**validated.model_dump(exclude_none=True), "pagetype": pagetype},
+            parse,
+            max_retries=max_retries,
+        )
+
+    @staticmethod
+    def _number(response) -> float:
+        try:
+            value = float(response.text.strip().replace(",", "."))
+            if not math.isfinite(value) or value < 0:
+                raise ValueError
+            return value
+        except ValueError:
+            raise XmlRiverError(kind="format") from None
+
+    async def get_balance(self, *, max_retries: int = 3) -> float:
+        return await self._request(BALANCE_URL, {}, self._number, max_retries=max_retries)
+
+    async def get_cost(self, engine: Engine, *, max_retries: int = 3) -> float | None:
+        if engine not in ("google", "yandex"):
+            raise ValueError("Unknown engine")
+        try:
+            return await self._request(
+                COST_URL.format(system=engine), {}, self._number, max_retries=max_retries
+            )
+        except XmlRiverError as exc:
+            if exc.kind in {"format", "temp"}:
+                return None
+            raise

xmlriver_mcp/config.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+from typing import Annotated
+
+import yaml
+from pydantic import Field
+
+from .models import Arguments, Engine
+
+
+class EngineConfig(Arguments):
+    name: Engine
+    region: Annotated[int, Field(ge=0)]
+    device: str = "mobile"
+
+
+class RegionConfig(Arguments):
+    name: str
+    yandex: Annotated[int, Field(ge=0)] | None = None
+    google: Annotated[int, Field(ge=0)] | None = None
+    hints: list[str] = []
+    expected_hosts: list[str] = []
+
+
+class Config(Arguments):
+    target_domain: str
+    include_subdomains: bool = True
+    depth: Annotated[int, Field(ge=10, le=100, multiple_of=10)] = 50
+    engines: list[EngineConfig] = [EngineConfig(name="yandex", region=213)]
+    regions: list[RegionConfig] = []
+    phrases: list[str] = []
+    csv: str | None = None
+    column: str = "phrase"
+    cache_ttl_hours: Annotated[float, Field(ge=0)] = 24
+    max_attempts: Annotated[int, Field(ge=1)] = 1000
+    deadline: Annotated[float, Field(gt=0)] = 180
+    rate_max: Annotated[int, Field(ge=1)] = 5
+    rate_window: Annotated[float, Field(gt=0)] = 1.5
+    costs_per_1000: dict[Engine, Annotated[float, Field(ge=0)]] = {}
+
+
+def load_config(path: Path) -> Config:
+    data = yaml.safe_load(path.read_text(encoding="utf-8"))
+    cfg = Config.model_validate(data)
+    if cfg.csv:
+        cfg.csv = str((path.resolve().parent / cfg.csv).resolve())
+    return cfg