WaveGuardClient 2.2.0__tar.gz → 2.3.0__tar.gz

{waveguardclient-2.2.0 → waveguardclient-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: WaveGuardClient
-Version: 2.2.0
+Version: 2.3.0
 Summary: Python SDK for WaveGuard — physics-based anomaly detection API
 Author: Greg Partin
 License-Expression: MIT

{waveguardclient-2.2.0 → waveguardclient-2.3.0}/WaveGuardClient.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: WaveGuardClient
-Version: 2.2.0
+Version: 2.3.0
 Summary: Python SDK for WaveGuard — physics-based anomaly detection API
 Author: Greg Partin
 License-Expression: MIT

{waveguardclient-2.2.0 → waveguardclient-2.3.0}/WaveGuardClient.egg-info/SOURCES.txt

@@ -4,8 +4,11 @@ pyproject.toml
 WaveGuardClient.egg-info/PKG-INFO
 WaveGuardClient.egg-info/SOURCES.txt
 WaveGuardClient.egg-info/dependency_links.txt
+WaveGuardClient.egg-info/entry_points.txt
 WaveGuardClient.egg-info/requires.txt
 WaveGuardClient.egg-info/top_level.txt
+mcp_server/__init__.py
+mcp_server/server.py
 tests/test_client.py
 tests/test_mcp_server.py
 waveguard/__init__.py

waveguardclient-2.3.0/WaveGuardClient.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+waveguard-mcp = mcp_server.server:main

waveguardclient-2.3.0/WaveGuardClient.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+mcp_server
+waveguard

waveguardclient-2.3.0/mcp_server/server.py

@@ -0,0 +1,528 @@
+"""
+WaveGuard MCP Server — Model Context Protocol for Claude Desktop & AI Agents.
+
+Stateless anomaly detection via wave physics simulation.
+One tool: ``waveguard_scan``.  Send training + test data, get anomaly scores.
+
+Transports
+----------
+- **stdio** (default) — add to Claude Desktop config
+- **HTTP** — ``python -m mcp_server --http --port 3001``
+
+Claude Desktop config
+~~~~~~~~~~~~~~~~~~~~~
+Add to ``~/.config/claude/claude_desktop_config.json``
+(macOS/Linux) or ``%APPDATA%\\Claude\\claude_desktop_config.json`` (Windows)::
+
+    {
+      "mcpServers": {
+        "waveguard": {
+          "command": "python",
+          "args": ["/path/to/WaveGuardClient/mcp_server/server.py"],
+          "env": {
+            "WAVEGUARD_API_KEY": "your-key-here"
+          }
+        }
+      }
+    }
+
+Smithery / Glama config
+~~~~~~~~~~~~~~~~~~~~~~~
+::
+
+    {
+      "mcpServers": {
+        "waveguard": {
+          "url": "https://gpartin--waveguard-api-fastapi-app.modal.run/mcp",
+          "transport": "http"
+        }
+      }
+    }
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import json
+import argparse
+from typing import Any, Dict, List, Optional
+
+# ── Configuration ──────────────────────────────────────────────────────────
+
+API_URL = os.environ.get(
+    "WAVEGUARD_API_URL",
+    "https://gpartin--waveguard-api-fastapi-app.modal.run",
+)
+API_KEY = os.environ.get("WAVEGUARD_API_KEY", "")
+
+# ── HTTP client ───────────────────────────────────────────────────────────
+
+try:
+    import requests
+
+    _session = requests.Session()
+    if API_KEY:
+        _session.headers["X-API-Key"] = API_KEY
+except ImportError:
+    _session = None  # type: ignore[assignment]
+
+
+def _api_post(path: str, body: dict) -> dict:
+    if _session is None:
+        raise RuntimeError("requests library required: pip install requests")
+    resp = _session.post(f"{API_URL}{path}", json=body, timeout=90)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def _api_get(path: str) -> Any:
+    if _session is None:
+        raise RuntimeError("requests library required: pip install requests")
+    resp = _session.get(f"{API_URL}{path}", timeout=30)
+    resp.raise_for_status()
+    return resp.json()
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# MCP Tool Definitions
+# ═══════════════════════════════════════════════════════════════════════════
+
+TOOLS = [
+    {
+        "name": "waveguard_scan",
+        "description": (
+            "Detect anomalies in data using GPU-accelerated wave physics simulation. "
+            "Fully stateless — send training data (normal examples) and test data "
+            "(samples to check) in ONE call. Returns per-sample anomaly scores, "
+            "confidence levels, and the top features explaining WHY each anomaly "
+            "was flagged. Works on any data type: JSON objects, numbers, text, "
+            "time series, arrays. No separate training step required.\n\n"
+            "Example: to check if server metrics are anomalous, send 3-5 normal "
+            "readings as training, and the suspect readings as test."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "training": {
+                    "type": "array",
+                    "description": (
+                        "2+ examples of NORMAL/expected data. These define what "
+                        "'normal' looks like. All samples should be the same type "
+                        "and shape. More samples = better baseline (4-10 is ideal)."
+                    ),
+                    "minItems": 2,
+                },
+                "test": {
+                    "type": "array",
+                    "description": (
+                        "1+ data points to check for anomalies. Same type/shape "
+                        "as training data. Each sample is scored independently."
+                    ),
+                    "minItems": 1,
+                },
+                "sensitivity": {
+                    "type": "number",
+                    "description": (
+                        "Anomaly threshold multiplier (default: 2.0). Lower = more "
+                        "sensitive (flags more anomalies). Higher = less sensitive. "
+                        "Range: 0.5 to 5.0."
+                    ),
+                },
+                "encoder_type": {
+                    "type": "string",
+                    "enum": [
+                        "json",
+                        "numeric",
+                        "text",
+                        "timeseries",
+                        "tabular",
+                        "image",
+                        "correlation",
+                    ],
+                    "description": (
+                        "Data encoder type. Omit to auto-detect from data shape. "
+                        "Auto-detection works well for most data."
+                    ),
+                },
+            },
+            "required": ["training", "test"],
+        },
+    },
+    {
+        "name": "waveguard_scan_timeseries",
+        "description": (
+            "Detect anomalies in time-series data using GPU-accelerated wave "
+            "physics simulation. Send a flat array of numeric values and a "
+            "window size. The tool automatically creates overlapping windows, "
+            "uses the first N as training (normal baseline), and scores the "
+            "remaining windows as test samples. Returns per-window anomaly "
+            "scores, confidence, and p-values.\n\n"
+            "Example: send 100 CPU-usage readings with window_size=10. "
+            "The first 5 windows become training, the rest are tested."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "array",
+                    "items": {"type": "number"},
+                    "description": (
+                        "Flat array of numeric time-series values in "
+                        "chronological order."
+                    ),
+                    "minItems": 4,
+                },
+                "window_size": {
+                    "type": "integer",
+                    "description": (
+                        "Number of data points per window (default: 10). "
+                        "Smaller windows = finer resolution."
+                    ),
+                },
+                "test_windows": {
+                    "type": "integer",
+                    "description": (
+                        "Number of trailing windows to test (default: auto, "
+                        "uses last ~40%% of windows)."
+                    ),
+                },
+                "sensitivity": {
+                    "type": "number",
+                    "description": (
+                        "Anomaly threshold multiplier (default: 2.0). Lower = "
+                        "more sensitive. Range: 0.5 to 5.0."
+                    ),
+                },
+            },
+            "required": ["data"],
+        },
+    },
+    {
+        "name": "waveguard_health",
+        "description": (
+            "Check WaveGuard API health, GPU availability, version, and engine "
+            "status. No authentication required. Use this to verify the service "
+            "is running before scanning."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {},
+        },
+    },
+]
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Tool Execution
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+def _execute_timeseries(arguments: dict) -> dict:
+    """Sliding-window timeseries scan via the /v1/scan endpoint."""
+    data = arguments["data"]
+    window = int(arguments.get("window_size", 10))
+    sensitivity = arguments.get("sensitivity", 2.0)
+
+    # Build windows
+    windows = [data[i : i + window] for i in range(0, len(data) - window + 1)]
+    if len(windows) < 3:
+        return {
+            "content": [
+                {
+                    "type": "text",
+                    "text": (
+                        f"Not enough data: {len(data)} points with "
+                        f"window_size={window} gives {len(windows)} windows "
+                        f"(need at least 3)."
+                    ),
+                }
+            ],
+            "isError": True,
+        }
+
+    # Split into training / test
+    test_count = arguments.get("test_windows")
+    if test_count is None:
+        test_count = max(1, len(windows) * 2 // 5)
+    test_count = min(test_count, len(windows) - 2)
+
+    training = windows[: len(windows) - test_count]
+    test = windows[len(windows) - test_count :]
+
+    body: dict = {
+        "training": training,
+        "test": test,
+        "encoder_type": "timeseries",
+        "sensitivity": sensitivity,
+    }
+    result = _api_post("/v1/scan", body)
+
+    # Summarise
+    lines = [
+        f"Time-series scan: {len(windows)} windows "
+        f"(window_size={window}, {len(training)} train, {len(test)} test)",
+        "",
+    ]
+    for i, r in enumerate(result.get("results", [])):
+        idx = len(training) + i
+        is_anom = r.get("is_anomaly", False)
+        conf = r.get("confidence", 0)
+        pval = r.get("p_value", 1.0)
+        marker = "ANOMALY" if is_anom else "Normal"
+        lines.append(
+            f"  Window {idx}: {marker} (confidence: {conf:.0%}, "
+            f"p-value: {pval:.4f})"
+        )
+    summary = "\n".join(lines)
+    return {
+        "content": [
+            {"type": "text", "text": summary},
+            {"type": "text", "text": json.dumps(result, indent=2)},
+        ]
+    }
+
+
+def execute_tool(name: str, arguments: dict) -> dict:
+    """Execute an MCP tool and return the result."""
+    try:
+        if name == "waveguard_scan":
+            body: dict = {
+                "training": arguments["training"],
+                "test": arguments["test"],
+            }
+            if "sensitivity" in arguments:
+                body["sensitivity"] = arguments["sensitivity"]
+            if "encoder_type" in arguments:
+                body["encoder_type"] = arguments["encoder_type"]
+
+            result = _api_post("/v1/scan", body)
+
+            # Build human-readable summary for the agent
+            summary_data = result.get("summary", {})
+            total = summary_data.get("total_samples", len(arguments["test"]))
+            anomalies = summary_data.get("anomalies_found", 0)
+            rate = summary_data.get("anomaly_rate", 0)
+
+            lines = [
+                f"Scanned {total} samples: {anomalies} anomalies ({rate:.0%} anomaly rate)",
+                "",
+            ]
+
+            for i, r in enumerate(result.get("results", [])):
+                is_anom = r.get("is_anomaly", False)
+                conf = r.get("confidence", 0)
+                score = r.get("score", 0)
+                marker = "ANOMALY" if is_anom else "Normal"
+                line = f"  Sample {i + 1}: {marker} (confidence: {conf:.0%}, score: {score:.1f})"
+
+                if is_anom and r.get("top_features"):
+                    feats = r["top_features"][:3]
+                    feat_str = ", ".join(
+                        f"{f.get('label', '?')} (z={f.get('z_score', 0):.1f})"
+                        for f in feats
+                    )
+                    line += f"\n           Top features: {feat_str}"
+
+                lines.append(line)
+
+            summary = "\n".join(lines)
+
+            return {
+                "content": [
+                    {"type": "text", "text": summary},
+                    {"type": "text", "text": json.dumps(result, indent=2)},
+                ]
+            }
+
+        elif name == "waveguard_scan_timeseries":
+            return _execute_timeseries(arguments)
+
+        elif name == "waveguard_health":
+            result = _api_get("/v1/health")
+            status = (
+                f"Status: {result.get('status', '?')} | "
+                f"Version: {result.get('version', '?')} | "
+                f"GPU: {result.get('gpu', 'N/A')}"
+            )
+            return {"content": [{"type": "text", "text": status}]}
+
+        else:
+            return {
+                "content": [{"type": "text", "text": f"Unknown tool: {name}"}],
+                "isError": True,
+            }
+
+    except Exception as e:
+        return {
+            "content": [{"type": "text", "text": f"Error: {e}"}],
+            "isError": True,
+        }
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# MCP Protocol Handler (stdio JSON-RPC transport)
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+class MCPStdioServer:
+    """Minimal MCP server implementing JSON-RPC 2.0 over stdio."""
+
+    def __init__(self) -> None:
+        self.server_info = {
+            "name": "waveguard",
+            "version": "2.3.0",
+        }
+
+    def handle_message(self, msg: dict) -> Optional[dict]:
+        """Process a JSON-RPC 2.0 message and return the response."""
+        method = msg.get("method", "")
+        msg_id = msg.get("id")
+        params = msg.get("params", {})
+
+        if method == "initialize":
+            return {
+                "jsonrpc": "2.0",
+                "id": msg_id,
+                "result": {
+                    "protocolVersion": "2024-11-05",
+                    "capabilities": {
+                        "tools": {"listChanged": False},
+                    },
+                    "serverInfo": self.server_info,
+                },
+            }
+
+        elif method == "notifications/initialized":
+            return None
+
+        elif method == "tools/list":
+            return {
+                "jsonrpc": "2.0",
+                "id": msg_id,
+                "result": {"tools": TOOLS},
+            }
+
+        elif method == "tools/call":
+            tool_name = params.get("name", "")
+            arguments = params.get("arguments", {})
+            result = execute_tool(tool_name, arguments)
+            return {
+                "jsonrpc": "2.0",
+                "id": msg_id,
+                "result": result,
+            }
+
+        elif method == "ping":
+            return {
+                "jsonrpc": "2.0",
+                "id": msg_id,
+                "result": {},
+            }
+
+        else:
+            if msg_id is not None:
+                return {
+                    "jsonrpc": "2.0",
+                    "id": msg_id,
+                    "error": {
+                        "code": -32601,
+                        "message": f"Method not found: {method}",
+                    },
+                }
+            return None
+
+    def run_stdio(self) -> None:
+        """Run the MCP server on stdin/stdout."""
+        sys.stderr.write(
+            f"WaveGuard MCP server v2.3.0 started (API: {API_URL})\n"
+        )
+        sys.stderr.flush()
+
+        for line in sys.stdin:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                msg = json.loads(line)
+            except json.JSONDecodeError as e:
+                sys.stderr.write(f"Invalid JSON: {e}\n")
+                sys.stderr.flush()
+                continue
+
+            response = self.handle_message(msg)
+            if response is not None:
+                sys.stdout.write(json.dumps(response) + "\n")
+                sys.stdout.flush()
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# HTTP transport (for remote MCP clients / Smithery / Glama)
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+def run_http_server(port: int = 3001) -> None:
+    """Run MCP over HTTP for remote agent access."""
+    try:
+        from fastapi import FastAPI as FA
+        import uvicorn
+    except ImportError:
+        print("HTTP transport requires: pip install fastapi uvicorn")
+        sys.exit(1)
+
+    mcp_app = FA(title="WaveGuard MCP Server", version="2.3.0")
+    server = MCPStdioServer()
+
+    @mcp_app.post("/mcp")
+    async def mcp_endpoint(request: dict) -> dict:  # type: ignore[type-arg]
+        return server.handle_message(request)  # type: ignore[return-value]
+
+    @mcp_app.get("/mcp/tools")
+    async def mcp_tools() -> dict:  # type: ignore[type-arg]
+        return {"tools": TOOLS}
+
+    print(f"WaveGuard MCP HTTP server v2.3.0 on port {port}")
+    uvicorn.run(mcp_app, host="0.0.0.0", port=port)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Entry point
+# ═══════════════════════════════════════════════════════════════════════════
+
+def main():
+    """Entry point for `waveguard-mcp` console script."""
+    parser = argparse.ArgumentParser(
+        description="WaveGuard MCP Server v2.3.0"
+    )
+    parser.add_argument(
+        "--http",
+        action="store_true",
+        help="Use HTTP transport instead of stdio",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=3001,
+        help="HTTP port (default: 3001)",
+    )
+    parser.add_argument(
+        "--api-url",
+        type=str,
+        default=None,
+        help="WaveGuard API URL (overrides $WAVEGUARD_API_URL)",
+    )
+    args = parser.parse_args()
+
+    global API_URL
+    if args.api_url:
+        API_URL = args.api_url
+
+    if args.http:
+        run_http_server(args.port)
+    else:
+        server = MCPStdioServer()
+        server.run_stdio()
+
+
+if __name__ == "__main__":
+    main()

{waveguardclient-2.2.0 → waveguardclient-2.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "WaveGuardClient"
-version = "2.2.0"
+version = "2.3.0"
 description = "Python SDK for WaveGuard — physics-based anomaly detection API"
 readme = "README.md"
 license = "MIT"
@@ -59,9 +59,12 @@ Repository = "https://github.com/gpartin/WaveGuardClient"
 Issues = "https://github.com/gpartin/WaveGuardClient/issues"
 API = "https://gpartin--waveguard-api-fastapi-app.modal.run/docs"
 
+[project.scripts]
+waveguard-mcp = "mcp_server.server:main"
+
 [tool.setuptools.packages.find]
-include = ["waveguard*"]
-exclude = ["tests*", "examples*", "docs*", "mcp_server*"]
+include = ["waveguard*", "mcp_server*"]
+exclude = ["tests*", "examples*", "docs*"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

{waveguardclient-2.2.0 → waveguardclient-2.3.0}/tests/test_client.py

@@ -35,6 +35,7 @@ def _mock_response(status_code=200, json_data=None, text=""):
     resp.status_code = status_code
     resp.text = text or json.dumps(json_data or {})
     resp.json.return_value = json_data or {}
+    resp.headers = {}
     return resp
 
 
@@ -92,10 +93,10 @@ class TestScan:
     @patch("waveguard.client.requests.Session")
     def test_scan_parses_results(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(200, SCAN_RESPONSE)
+        session.request.return_value = _mock_response(200, SCAN_RESPONSE)
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         result = wg.scan(
             training=[{"a": 1}, {"a": 2}, {"a": 3}],
             test=[{"a": 2}, {"a": 100}],
@@ -113,10 +114,10 @@ class TestScan:
     @patch("waveguard.client.requests.Session")
     def test_scan_top_features(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(200, SCAN_RESPONSE)
+        session.request.return_value = _mock_response(200, SCAN_RESPONSE)
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         result = wg.scan(training=[{"a": 1}, {"a": 2}], test=[{"a": 100}])
 
         anomaly = result.results[1]
@@ -127,10 +128,10 @@ class TestScan:
     @patch("waveguard.client.requests.Session")
     def test_scan_sends_optional_params(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(200, SCAN_RESPONSE)
+        session.request.return_value = _mock_response(200, SCAN_RESPONSE)
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         wg.scan(
             training=[{"a": 1}, {"a": 2}],
             test=[{"a": 3}],
@@ -138,7 +139,7 @@ class TestScan:
             sensitivity=0.5,
         )
 
-        call_args = session.post.call_args
+        call_args = session.request.call_args
         body = call_args.kwargs.get("json") or call_args[1].get("json")
         assert body["encoder_type"] == "text"
         assert body["sensitivity"] == 0.5
@@ -151,40 +152,40 @@ class TestErrors:
     @patch("waveguard.client.requests.Session")
     def test_401_raises_auth_error(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(401, text="Unauthorized")
+        session.request.return_value = _mock_response(401, text="Unauthorized")
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="bad-key")
+        wg = WaveGuard(api_key="bad-key", max_retries=0)
         with pytest.raises(AuthenticationError):
             wg.scan(training=[{"a": 1}, {"a": 2}], test=[{"a": 3}])
 
     @patch("waveguard.client.requests.Session")
     def test_422_raises_validation_error(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(422, text="Empty training")
+        session.request.return_value = _mock_response(422, text="Empty training")
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         with pytest.raises(ValidationError):
             wg.scan(training=[], test=[{"a": 1}])
 
     @patch("waveguard.client.requests.Session")
     def test_429_raises_rate_limit_error(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(429, text="Rate limited")
+        session.request.return_value = _mock_response(429, text="Rate limited")
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         with pytest.raises(RateLimitError):
             wg.scan(training=[{"a": 1}, {"a": 2}], test=[{"a": 3}])
 
     @patch("waveguard.client.requests.Session")
     def test_500_raises_server_error(self, mock_session_cls):
         session = MagicMock()
-        session.post.return_value = _mock_response(500, text="Internal error")
+        session.request.return_value = _mock_response(500, text="Internal error")
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         with pytest.raises(ServerError):
             wg.scan(training=[{"a": 1}, {"a": 2}], test=[{"a": 3}])
 
@@ -193,10 +194,10 @@ class TestErrors:
         import requests as req
 
         session = MagicMock()
-        session.post.side_effect = req.ConnectionError("Cannot connect")
+        session.request.side_effect = req.ConnectionError("Cannot connect")
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         with pytest.raises(WaveGuardError, match="Cannot connect"):
             wg.scan(training=[{"a": 1}, {"a": 2}], test=[{"a": 3}])
 
@@ -214,7 +215,7 @@ class TestHealth:
     @patch("waveguard.client.requests.Session")
     def test_health_parses(self, mock_session_cls):
         session = MagicMock()
-        session.get.return_value = _mock_response(
+        session.request.return_value = _mock_response(
             200,
             {
                 "status": "healthy",
@@ -226,7 +227,7 @@ class TestHealth:
         )
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         h = wg.health()
         assert isinstance(h, HealthStatus)
         assert h.status == "healthy"
@@ -235,7 +236,7 @@ class TestHealth:
     @patch("waveguard.client.requests.Session")
     def test_tier_parses(self, mock_session_cls):
         session = MagicMock()
-        session.get.return_value = _mock_response(
+        session.request.return_value = _mock_response(
             200,
             {
                 "tier": "PRO",
@@ -248,7 +249,7 @@ class TestHealth:
         )
         mock_session_cls.return_value = session
 
-        wg = WaveGuard(api_key="test")
+        wg = WaveGuard(api_key="test", max_retries=0)
         t = wg.tier()
         assert isinstance(t, TierInfo)
         assert t.tier == "PRO"

{waveguardclient-2.2.0 → waveguardclient-2.3.0}/waveguard/client.py

@@ -26,6 +26,10 @@ Usage::
 
 from __future__ import annotations
 
+import logging
+import os
+import time
+import random
 import requests
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional
@@ -38,7 +42,9 @@ from .exceptions import (
     ServerError,
 )
 
-__version__ = "2.2.0"
+__version__ = "2.3.0"
+
+logger = logging.getLogger("waveguard")
 
 
 # ─────────────────────────────── Data Classes ─────────────────────────────
@@ -140,17 +146,23 @@ class WaveGuard:
 
     Parameters
     ----------
-    api_key : str
-        Your WaveGuard API key.
+    api_key : str, optional
+        Your WaveGuard API key.  If not provided, reads from the
+        ``WAVEGUARD_API_KEY`` environment variable.  Free-tier scans
+        work without a key (rate-limited).
     base_url : str, optional
         API base URL.  Defaults to the production Modal endpoint.
     timeout : float, optional
         Request timeout in seconds.  Default ``120`` (generous for GPU
         cold starts).
+    max_retries : int, optional
+        Number of automatic retries on transient errors (429, 500, 502,
+        503, 504, connection errors, timeouts).  Default ``2``.
+        Set to ``0`` to disable retries.
 
     Examples
     --------
-    >>> wg = WaveGuard(api_key="wg_test_key")
+    >>> wg = WaveGuard()  # reads WAVEGUARD_API_KEY from env
     >>> result = wg.scan(
     ...     training=[{"a": 1}, {"a": 2}, {"a": 3}],
     ...     test=[{"a": 100}],
@@ -161,22 +173,27 @@ class WaveGuard:
 
     DEFAULT_URL = "https://gpartin--waveguard-api-fastapi-app.modal.run"
 
+    # Status codes that trigger automatic retry
+    _RETRYABLE_STATUS = {429, 500, 502, 503, 504}
+
     def __init__(
         self,
         api_key: Optional[str] = None,
         base_url: str = DEFAULT_URL,
         timeout: float = 120.0,
+        max_retries: int = 2,
     ):
-        self.api_key = api_key
+        self.api_key = api_key or os.environ.get("WAVEGUARD_API_KEY")
         self.base_url = base_url.rstrip("/")
         self.timeout = timeout
+        self.max_retries = max_retries
         self._session = requests.Session()
         headers = {
             "Content-Type": "application/json",
             "User-Agent": f"waveguard-python/{__version__}",
         }
-        if api_key:
-            headers["X-API-Key"] = api_key
+        if self.api_key:
+            headers["X-API-Key"] = self.api_key
         self._session.headers.update(headers)
 
     # ── Core API ──────────────────────────────────────────────────────
@@ -302,66 +319,127 @@ class WaveGuard:
     # ── Internal HTTP ─────────────────────────────────────────────────
 
     def _post(self, path: str, body: dict) -> dict:
-        url = f"{self.base_url}{path}"
-        try:
-            r = self._session.post(url, json=body, timeout=self.timeout)
-        except requests.ConnectionError:
-            raise WaveGuardError(f"Cannot connect to {self.base_url}")
-        except requests.Timeout:
-            raise WaveGuardError(
-                f"Request timed out after {self.timeout}s"
-            )
-        return self._handle(r)
+        return self._request("POST", path, json=body)
 
     def _get(self, path: str) -> dict:
+        return self._request("GET", path)
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        json: Optional[dict] = None,
+    ) -> dict:
+        """Execute an HTTP request with automatic retry and backoff."""
         url = f"{self.base_url}{path}"
-        try:
-            r = self._session.get(url, timeout=self.timeout)
-        except requests.ConnectionError:
-            raise WaveGuardError(f"Cannot connect to {self.base_url}")
-        except requests.Timeout:
-            raise WaveGuardError(
-                f"Request timed out after {self.timeout}s"
-            )
-        return self._handle(r)
-
-    def _handle(self, r: requests.Response) -> dict:
-        if r.status_code == 401:
-            raise AuthenticationError(
-                "Invalid or missing API key",
-                status_code=401,
-                detail=r.text,
-            )
-        if r.status_code == 422:
-            raise ValidationError(
-                f"Validation failed: {r.text}",
-                status_code=422,
-                detail=r.text,
-            )
-        if r.status_code == 429:
-            raise RateLimitError(
-                f"Rate or tier limit exceeded. "
-                f"Upgrade at {RateLimitError.UPGRADE_URL}\n"
-                f"Detail: {r.text}",
-                status_code=429,
-                detail=r.text,
-            )
-        if r.status_code >= 500:
-            raise ServerError(
-                f"Server error {r.status_code}: {r.text}",
-                status_code=r.status_code,
-                detail=r.text,
-            )
-        if r.status_code >= 400:
-            raise WaveGuardError(
-                f"API error {r.status_code}: {r.text}",
-                status_code=r.status_code,
-                detail=r.text,
-            )
-        try:
-            return r.json()
-        except ValueError:
-            return {"raw": r.text}
+        last_exc: Optional[Exception] = None
+
+        for attempt in range(1 + self.max_retries):
+            try:
+                logger.debug(
+                    "%s %s (attempt %d/%d)",
+                    method, path, attempt + 1, 1 + self.max_retries,
+                )
+                r = self._session.request(
+                    method, url, json=json, timeout=self.timeout
+                )
+
+                # Non-retryable errors — raise immediately
+                if r.status_code in (401, 403):
+                    raise AuthenticationError(
+                        "Invalid or missing API key. "
+                        "Set WAVEGUARD_API_KEY or pass api_key= to WaveGuard().",
+                        status_code=r.status_code,
+                        detail=r.text,
+                    )
+                if r.status_code == 422:
+                    raise ValidationError(
+                        f"Validation failed: {r.text}",
+                        status_code=422,
+                        detail=r.text,
+                    )
+
+                # Retryable errors
+                if r.status_code in self._RETRYABLE_STATUS:
+                    retry_after = r.headers.get("Retry-After")
+                    if r.status_code == 429 and attempt == self.max_retries:
+                        raise RateLimitError(
+                            f"Rate or tier limit exceeded. "
+                            f"Upgrade at {RateLimitError.UPGRADE_URL}\n"
+                            f"Detail: {r.text}",
+                            status_code=429,
+                            detail=r.text,
+                        )
+                    if attempt < self.max_retries:
+                        delay = self._backoff_delay(attempt, retry_after)
+                        logger.info(
+                            "Retryable %d from %s — retrying in %.1fs",
+                            r.status_code, path, delay,
+                        )
+                        time.sleep(delay)
+                        continue
+                    # Final attempt — raise appropriate error
+                    if r.status_code >= 500:
+                        raise ServerError(
+                            f"Server error {r.status_code} after "
+                            f"{self.max_retries} retries",
+                            status_code=r.status_code,
+                            detail=r.text,
+                        )
+
+                # Other 4xx errors
+                if r.status_code >= 400:
+                    raise WaveGuardError(
+                        f"API error {r.status_code}: {r.text}",
+                        status_code=r.status_code,
+                        detail=r.text,
+                    )
+
+                # Success
+                try:
+                    return r.json()
+                except ValueError:
+                    raise WaveGuardError(
+                        f"Unexpected non-JSON response from {path}",
+                        status_code=r.status_code,
+                        detail=r.text,
+                    )
+
+            except (requests.ConnectionError, requests.Timeout) as e:
+                last_exc = e
+                if attempt < self.max_retries:
+                    delay = self._backoff_delay(attempt)
+                    logger.info(
+                        "%s on %s — retrying in %.1fs",
+                        type(e).__name__, path, delay,
+                    )
+                    time.sleep(delay)
+                    continue
+                if isinstance(e, requests.Timeout):
+                    raise WaveGuardError(
+                        f"Request timed out after {self.timeout}s "
+                        f"({self.max_retries} retries exhausted)"
+                    ) from e
+                raise WaveGuardError(
+                    f"Cannot connect to {self.base_url} "
+                    f"({self.max_retries} retries exhausted)"
+                ) from e
+
+        # Should not reach here, but just in case
+        raise WaveGuardError("Request failed") from last_exc
+
+    @staticmethod
+    def _backoff_delay(
+        attempt: int, retry_after: Optional[str] = None
+    ) -> float:
+        """Exponential backoff with jitter, respecting Retry-After."""
+        if retry_after:
+            try:
+                return min(float(retry_after), 60.0)
+            except ValueError:
+                pass
+        base = min(2 ** attempt, 30)  # 1, 2, 4, 8, ... capped at 30s
+        return base + random.uniform(0, base * 0.5)
 
     def __repr__(self) -> str:
         return f"WaveGuard(base_url='{self.base_url}')"

waveguardclient-2.2.0/WaveGuardClient.egg-info/top_level.txt

@@ -1 +0,0 @@
-waveguard