aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/plugins/mcp/http_server.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+import os
+import urllib.parse
+
+from fastapi import FastAPI, HTTPException, Request
+import httpx
+from pydantic import BaseModel
+import uvicorn
+
+# TODO: move it to tests/examples later
+DEMO_HTTP_TOKEN = os.getenv("DEMO_HTTP_TOKEN")
+
+app = FastAPI()
+
+TOOLS = [
+    {
+        "name": "search",
+        "description": "Search Wikipedia and return top hits.",
+        "input_schema": {
+            "type": "object",
+            "properties": {"q": {"type": "string"}, "k": {"type": "integer"}},
+            "required": ["q"],
+        },
+    }
+]
+
+
+class RPCReq(BaseModel):
+    jsonrpc: str
+    id: int | str | None = None
+    method: str
+    params: dict | None = None
+
+
+def ok(i, result):
+    return {"jsonrpc": "2.0", "id": i, "result": result}
+
+
+def err(i, msg, code=-32000, data=None):
+    e = {"jsonrpc": "2.0", "id": i, "error": {"code": code, "message": msg}}
+    if data is not None:
+        e["error"]["data"] = data
+    return e
+
+
+async def do_search(q: str, k: int = 5):
+    url = "https://en.wikipedia.org/w/api.php"
+    params = {
+        "action": "query",
+        "list": "search",
+        "format": "json",
+        "srsearch": q,
+        "srlimit": max(1, min(int(k or 5), 10)),
+    }
+    async with httpx.AsyncClient(timeout=10) as c:
+        r = await c.get(url, params=params)
+        r.raise_for_status()
+        data = r.json()
+    hits = []
+    for it in data.get("query", {}).get("search") or []:
+        title = it.get("title", "")
+        page = "https://en.wikipedia.org/wiki/" + urllib.parse.quote(title.replace(" ", "_"))
+        hits.append({"title": title, "url": page, "snippet": it.get("snippet", "")})
+    return {"hits": hits}
+
+
+@app.post("/rpc")
+async def rpc(req: RPCReq, request: Request):
+    if DEMO_HTTP_TOKEN:
+        auth = request.headers.get("authorization", "")
+        if auth != f"Bearer {DEMO_HTTP_TOKEN}":
+            raise HTTPException(status_code=401, detail="Unauthorized")
+    else:
+        logger = logging.getLogger("aethergraph.plugins.mcp.http_server")
+        logger.warning(
+            "No auth token DEMO_HTTP_TOKEN set, skipping auth check. Set up DEMO_HTTP_TOKEN in env for test."
+        )
+    try:
+        p = req.params or {}
+        if req.method == "tools/list":
+            return ok(req.id, TOOLS)
+        if req.method == "tools/call":
+            name = (p.get("name") or "").strip()
+            args = p.get("arguments") or {}
+            if name in ("search", "query"):
+                res = await do_search(args.get("q", ""), int(args.get("k", 5)))
+                return ok(req.id, res)
+            return err(req.id, f"Unknown tool: {name}")
+        if req.method == "resources/list":
+            return ok(req.id, [])
+        if req.method == "resources/read":
+            return ok(req.id, {"uri": p.get("uri"), "data": None})
+        return err(req.id, f"Unknown method: {req.method}")
+    except Exception as e:
+        return err(req.id, str(e))
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="127.0.0.1", port=8769)

aethergraph/plugins/mcp/ws_server.py

@@ -0,0 +1,180 @@
+# ws_mcp_server.py  (robust for websockets v15, with optional token auth)
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import urllib.parse
+
+import httpx
+from websockets import exceptions as ws_exceptions, serve
+from websockets.http import Headers
+
+# -------- Config --------
+DEMO_WS_TOKEN = os.getenv("DEMO_WS_TOKEN", "").strip()
+REQUIRE_HEADER_BEARER = True  # require Authorization header when token set
+ALLOW_FIRST_MESSAGE_AUTH = True  # also allow in-band JSON-RPC auth frame
+
+TOOLS = [
+    {
+        "name": "search",
+        "description": "Search Wikipedia and return top hits.",
+        "input_schema": {
+            "type": "object",
+            "properties": {"q": {"type": "string"}, "k": {"type": "integer"}},
+            "required": ["q"],
+        },
+    }
+]
+
+
+def ok(i, result):
+    return {"jsonrpc": "2.0", "id": i, "result": result}
+
+
+def err(i, msg, code=-32000, data=None):
+    e = {"jsonrpc": "2.0", "id": i, "error": {"code": code, "message": msg}}
+    if data is not None:
+        e["error"]["data"] = data
+    return e
+
+
+async def do_search(q: str, k: int = 5):
+    params = {
+        "action": "query",
+        "list": "search",
+        "format": "json",
+        "srsearch": q,
+        "srlimit": max(1, min(int(k or 5), 10)),
+    }
+    url = "https://en.wikipedia.org/w/api.php?" + urllib.parse.urlencode(params)
+    async with httpx.AsyncClient(timeout=15.0) as client:
+        r = await client.get(url)
+        r.raise_for_status()
+        data = r.json()
+    hits = []
+    for item in data.get("query", {}).get("search") or []:
+        title = item.get("title", "")
+        page = "https://en.wikipedia.org/wiki/" + urllib.parse.quote(title.replace(" ", "_"))
+        hits.append({"title": title, "url": page, "snippet": item.get("snippet", "")})
+    return {"hits": hits}
+
+
+# ---------- Handshake-time token check (recommended) ----------
+async def process_request(path: str, request_headers: Headers):
+    """If DEMO_WS_TOKEN is set, enforce Authorization: Bearer <token> at handshake."""
+    if not DEMO_WS_TOKEN or not REQUIRE_HEADER_BEARER:
+        return  # accept; continue with handshake
+
+    auth = request_headers.get("Authorization", "")
+    if auth == f"Bearer {DEMO_WS_TOKEN}":
+        return  # ok
+
+    # Reject handshake with 401
+    body = b"Unauthorized"
+    headers = [
+        ("Content-Type", "text/plain; charset=utf-8"),
+        ("Content-Length", str(len(body))),
+        ("WWW-Authenticate", 'Bearer realm="mcp-ws", error="invalid_token"'),
+    ]
+    return (401, headers, body)
+
+
+# ---------- Handler ----------
+async def handle(ws):
+    # Optional: in-band first-message auth if header was not used
+    if DEMO_WS_TOKEN and ALLOW_FIRST_MESSAGE_AUTH and (not REQUIRE_HEADER_BEARER):
+        try:
+            first_raw = await asyncio.wait_for(ws.recv(), timeout=5.0)
+            first = json.loads(first_raw)
+            if first.get("method") != "auth/bearer":
+                await ws.send(
+                    json.dumps(err(first.get("id"), "Unauthorized: expected auth/bearer"))
+                )
+                await ws.close()
+                return
+            tok = (first.get("params") or {}).get("token", "")
+            if tok != DEMO_WS_TOKEN:
+                await ws.send(json.dumps(err(first.get("id"), "Unauthorized: bad token")))
+                await ws.close()
+                return
+            # auth ok; optionally reply success
+            await ws.send(json.dumps(ok(first.get("id"), {"ok": True})))
+        except Exception:
+            # couldn't read/parse first frame or wrong shape
+            try:
+                await ws.send(json.dumps(err(None, "Unauthorized")))
+            finally:
+                await ws.close()
+            return
+
+    try:
+        async for raw in ws:
+            try:
+                req = json.loads(raw)
+                mid = req.get("id")
+                method = req.get("method")
+                params = req.get("params") or {}
+
+                if method == "tools/list":
+                    await ws.send(json.dumps(ok(mid, TOOLS)))
+                    continue
+
+                if method == "tools/call":
+                    name = (params.get("name") or "").strip()
+                    args = params.get("arguments") or {}
+                    if name in ("search", "query"):
+                        res = await do_search(args.get("q", ""), int(args.get("k", 5)))
+                        await ws.send(json.dumps(ok(mid, res)))
+                        continue
+                    await ws.send(json.dumps(err(mid, f"Unknown tool: {name}")))
+                    continue
+
+                if method == "resources/list":
+                    await ws.send(json.dumps(ok(mid, [])))
+                    continue
+
+                if method == "resources/read":
+                    await ws.send(json.dumps(ok(mid, {"uri": params.get("uri"), "data": None})))
+                    continue
+
+                await ws.send(json.dumps(err(mid, f"Unknown method: {method}")))
+            except Exception as e:
+                # Return JSON-RPC error but keep the session alive
+                try:
+                    rid = req.get("id") if isinstance(req, dict) else None
+                except Exception:
+                    rid = None
+                await ws.send(json.dumps(err(rid, str(e))))
+    except (ws_exceptions.ConnectionClosedOK, ws_exceptions.ConnectionClosedError):
+        return
+
+
+async def main(host="0.0.0.0", port=8765):
+    # If REQUIRE header-based auth and DISABLE in-band auth:
+    #   set REQUIRE_HEADER_BEARER=True and ALLOW_FIRST_MESSAGE_AUTH=False
+    async with serve(
+        handle,
+        host,
+        port,
+        ping_interval=20,
+        ping_timeout=10,
+        close_timeout=2,
+        max_queue=32,
+        process_request=process_request,  # <— handshake auth hook
+    ):
+        logger = logging.getLogger("aethergraph.plugins.mcp.ws_server")
+        logger.info(f"MCP WS server listening on ws://{host}:{port}")
+        if DEMO_WS_TOKEN:
+            mode = []
+            if REQUIRE_HEADER_BEARER:
+                mode.append("header")
+            if ALLOW_FIRST_MESSAGE_AUTH and not REQUIRE_HEADER_BEARER:
+                mode.append("first-message")
+            logger.info(f"Auth enabled: token set; modes: {', '.join(mode) or 'none'}")
+        await asyncio.Future()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

aethergraph/plugins/net/http.py

@@ -0,0 +1,10 @@
+# aethergraph/net/http.py
+from contextlib import asynccontextmanager
+
+import httpx
+
+
+@asynccontextmanager
+async def get_async_client(timeout_s: float = 10.0, headers: dict | None = None):
+    async with httpx.AsyncClient(timeout=timeout_s, headers=headers) as client:
+        yield client

aethergraph/plugins/utils/data_io.py

@@ -0,0 +1,359 @@
+# aethergraph/v2/utils/data_io.py
+from __future__ import annotations
+
+import csv
+import hashlib
+import io
+import os
+from typing import Any
+
+# Optional deps
+try:
+    _HAS_PANDAS = True
+except Exception:
+    _HAS_PANDAS = False
+
+try:
+    _HAS_PIL = True
+except Exception:
+    _HAS_PIL = False
+
+try:
+    import numpy as np
+
+    _HAS_NUMPY = True
+except Exception:
+    _HAS_NUMPY = False
+
+try:
+    from pypdf import PdfReader  # lightweight text extractor
+
+    _HAS_PYPDF = True
+except Exception:
+    _HAS_PYPDF = False
+
+
+# ---------- URI helpers ----------
+
+
+def _resolve_local_path(uri: str) -> str | None:
+    if uri.startswith("file://"):
+        return uri[len("file://") :]
+    return None
+
+
+def load_bytes(artifact_store, uri: str) -> bytes:
+    """Load raw bytes from artifact_store via URI (prefer local file path).
+    TODO: deprecate this function as artifact_store.get_bytes(uri) is preferred.
+    """
+    p = _resolve_local_path(uri)
+    if p and os.path.exists(p):
+        with open(p, "rb") as f:
+            return f.read()
+    # Optional: if later add artifact_store.get_bytes(uri), handle here.
+    raise FileNotFoundError(f"Cannot resolve bytes for URI: {uri}")
+
+
+# ---------- MIME normalization & classification ----------
+
+_EXTENSION_TO_MIME = {
+    # images
+    "png": "image/png",
+    "jpg": "image/jpeg",
+    "jpeg": "image/jpeg",
+    "gif": "image/gif",
+    "webp": "image/webp",
+    "tif": "image/tiff",
+    "tiff": "image/tiff",
+    "bmp": "image/bmp",
+    "svg": "image/svg+xml",
+    "heic": "image/heic",
+    "heif": "image/heif",
+    # text/docs
+    "txt": "text/plain",
+    "log": "text/plain",
+    "md": "text/markdown",
+    "csv": "text/csv",
+    "tsv": "text/tab-separated-values",
+    "json": "application/json",
+    "yaml": "text/yaml",
+    "yml": "text/yaml",
+    "xml": "application/xml",
+    "pdf": "application/pdf",
+    # archives
+    "zip": "application/zip",
+    "gz": "application/gzip",
+    "tar": "application/x-tar",
+    "7z": "application/x-7z-compressed",
+    # office
+    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+    "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    "pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+    # proprietary examples
+    "mph": "application/octet-stream",  # COMSOL
+}
+
+
+def normalize_mime(name: str | None, mimetype_hint: str | None) -> str:
+    mt = (mimetype_hint or "").lower().strip()
+    if mt:
+        return mt
+    if name:
+        n = name.lower()
+        if "." in n:
+            ext = n.rsplit(".", 1)[-1]
+            return _EXTENSION_TO_MIME.get(ext, "application/octet-stream")
+    return "application/octet-stream"
+
+
+def classify_for_processing(mime: str) -> str:
+    m = mime.lower()
+    if m.startswith("image/") and m != "image/svg+xml":
+        return "image"
+    if m in ("image/svg+xml", "application/xml", "text/xml"):
+        return "xml"
+    if m.startswith("text/") or m in ("application/json",):
+        return "text"
+    if m == "application/pdf":
+        return "pdf"
+    if m in (
+        "application/zip",
+        "application/gzip",
+        "application/x-tar",
+        "application/x-7z-compressed",
+    ):
+        return "archive"
+    if m in ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",):
+        return "xlsx"
+    return "binary"  # unknown/proprietary (e.g., COMSOL .mph)
+
+
+# ---------- Text helpers ----------
+
+
+def try_decode_text(b: bytes) -> str:
+    for enc in ("utf-8", "utf-8-sig", "utf-16", "utf-16le", "utf-16be"):
+        try:
+            return b.decode(enc)
+        except UnicodeDecodeError:
+            pass
+    return b.decode("latin-1", errors="replace")
+
+
+# ---------- CSV helpers ----------
+
+
+def read_csv_any(b: bytes) -> dict[str, Any]:
+    """
+    Returns a lightweight preview for CSV/TSV; if pandas is available, also return a DataFrame preview.
+    """
+    txt = try_decode_text(b)
+    # Dialect sniff
+    sniffer = csv.Sniffer()
+    try:
+        dialect = sniffer.sniff(txt.splitlines()[0] if txt else "")
+    except Exception:
+        dialect = csv.excel
+    rows = list(csv.reader(io.StringIO(txt), dialect=dialect))
+    preview_rows = rows[:10]
+    out: dict[str, Any] = {"rows_preview": preview_rows, "num_rows_previewed": len(preview_rows)}
+    if _HAS_PANDAS:
+        try:
+            import pandas as pd  # type: ignore
+
+            # Let pandas infer separator; fallback to comma
+            df = pd.read_csv(io.StringIO(txt))
+            out["pandas_head"] = df.head(10)  # DataFrame (caller may display)
+            out["columns"] = list(df.columns)
+            out["shape"] = tuple(df.shape)
+        except Exception:
+            pass
+    return out
+
+
+# ---------- Image helpers ----------
+
+
+def decode_image_pil(b: bytes, *, fix_orientation: bool = True, to_rgb: bool = True):
+    if not _HAS_PIL:
+        raise RuntimeError("Pillow not installed; cannot decode image.")
+    from PIL import Image, ImageCms, ImageOps  # local import for safety
+
+    im = Image.open(io.BytesIO(b))
+    if fix_orientation:
+        im = ImageOps.exif_transpose(im)
+    try:
+        if "icc_profile" in im.info:
+            src = ImageCms.ImageCmsProfile(io.BytesIO(im.info.get("icc_profile")))
+            dst = ImageCms.createProfile("sRGB")
+            im = ImageCms.profileToProfile(im, src, dst, outputMode=im.mode)
+    except Exception:
+        pass
+    if to_rgb and im.mode not in ("RGB", "RGBA"):
+        im = im.convert("RGB")
+    return im
+
+
+def pil_to_numpy(im) -> np.ndarray:
+    if not _HAS_NUMPY:
+        raise RuntimeError("NumPy not installed; cannot convert image to array.")
+    import numpy as np  # type: ignore
+
+    arr = np.asarray(im)
+    if arr.ndim == 2:
+        arr = arr[:, :, None]
+    return arr
+
+
+# ---------- PDF helpers ----------
+
+
+def extract_pdf_text(b: bytes, max_pages: int = 5) -> dict[str, Any]:
+    if not _HAS_PYPDF:
+        raise RuntimeError("pypdf not installed; cannot extract PDF text.")
+    reader = PdfReader(io.BytesIO(b))
+    pages = min(len(reader.pages), max_pages)
+    texts = []
+    for i in range(pages):
+        try:
+            texts.append(reader.pages[i].extract_text() or "")
+        except Exception:
+            texts.append("")
+    return {"num_pages": len(reader.pages), "preview_pages": pages, "text_preview": texts}
+
+
+# ---------- XLSX helpers (cheap preview without heavy engines) ----------
+
+
+def preview_xlsx(b: bytes) -> dict[str, Any]:
+    if not _HAS_PANDAS:
+        return {"note": "pandas not installed; cannot preview xlsx"}
+    try:
+        import pandas as pd  # type: ignore
+
+        with io.BytesIO(b) as bio:
+            xl = pd.ExcelFile(bio)
+            sheets = xl.sheet_names
+            out: dict[str, Any] = {"sheets": sheets, "previews": {}}
+            for s in sheets[:3]:
+                try:
+                    df = xl.parse(s, nrows=10)
+                    out["previews"][s] = df
+                except Exception:
+                    out["previews"][s] = "unreadable"
+            return out
+    except Exception as e:
+        return {"error": f"xlsx preview failed: {e}"}
+
+
+# ---------- Dispatcher ----------
+
+
+def quick_decode(
+    artifact_store, name: str | None, mimetype_hint: str | None, uri: str
+) -> dict[str, Any]:
+    """
+    Returns: {
+      'uri': str, 'mimetype': str, 'kind': 'image'|'text'|'pdf'|'archive'|'xlsx'|'binary',
+      'meta': {...},  # shape/columns/pages etc
+      'preview': ...  # small human-friendly preview (safe to log/send)
+    }
+    """
+    b = load_bytes(artifact_store, uri)
+    mime = normalize_mime(name, mimetype_hint)
+    kind = classify_for_processing(mime)
+    sha = hashlib.sha256(b).hexdigest()
+
+    if kind == "image":
+        if not _HAS_PIL:
+            return {
+                "uri": uri,
+                "mimetype": mime,
+                "kind": kind,
+                "meta": {"sha256": sha},
+                "preview": "Pillow missing",
+            }
+        im = decode_image_pil(b)
+        w, h = im.size
+        meta = {"width": w, "height": h, "mode": im.mode, "sha256": sha}
+        if _HAS_NUMPY:
+            arr = pil_to_numpy(im)
+            meta["array_shape"] = tuple(arr.shape)
+        return {
+            "uri": uri,
+            "mimetype": mime,
+            "kind": kind,
+            "meta": meta,
+            "preview": f"{w}x{h} {im.mode}",
+        }
+
+    if kind == "text":
+        txt = try_decode_text(b)
+        head = "\n".join(txt.splitlines()[:20])
+        return {
+            "uri": uri,
+            "mimetype": mime,
+            "kind": kind,
+            "meta": {"bytes": len(b), "sha256": sha},
+            "preview": head,
+        }
+
+    if kind == "pdf":
+        if not _HAS_PYPDF:
+            return {
+                "uri": uri,
+                "mimetype": mime,
+                "kind": kind,
+                "meta": {"bytes": len(b), "sha256": sha},
+                "preview": "pypdf missing",
+            }
+        meta = extract_pdf_text(b)
+        return {
+            "uri": uri,
+            "mimetype": mime,
+            "kind": kind,
+            "meta": {"bytes": len(b), **meta, "sha256": sha},
+            "preview": "\n---\n".join(meta["text_preview"]),
+        }
+
+    if kind == "archive":
+        # We don't auto-unpack; just list ZIP members if it's zip
+        import zipfile
+
+        bio = io.BytesIO(b)
+        if zipfile.is_zipfile(bio):
+            with zipfile.ZipFile(bio) as z:
+                names = z.namelist()[:20]
+            return {
+                "uri": uri,
+                "mimetype": mime,
+                "kind": kind,
+                "meta": {"bytes": len(b), "sha256": sha},
+                "preview": "\n".join(names),
+            }
+        return {
+            "uri": uri,
+            "mimetype": mime,
+            "kind": kind,
+            "meta": {"bytes": len(b), "sha256": sha},
+            "preview": "archive (non-zip)",
+        }
+
+    if kind == "xlsx":
+        meta = preview_xlsx(b)
+        return {
+            "uri": uri,
+            "mimetype": mime,
+            "kind": kind,
+            "meta": {"bytes": len(b), **meta, "sha256": sha},
+            "preview": f"sheets: {', '.join(meta.get('sheets', [])) if isinstance(meta, dict) else meta}",
+        }
+
+    # binary / unknown (e.g., COMSOL .mph)
+    return {
+        "uri": uri,
+        "mimetype": mime,
+        "kind": "binary",
+        "meta": {"bytes": len(b), "sha256": sha},
+        "preview": "opaque binary",
+    }

aethergraph/runner/__init__.py

@@ -0,0 +1,5 @@
+# Aethergraph Runner -- direct aethergraph.core.runtime.graph_runner for clean imports
+
+from aethergraph.core.runtime.graph_runner import run, run_async
+
+__all__ = ["run", "run_async"]