aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/services/mcp/ws_client.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+from typing import Any
+
+import websockets
+from websockets.client import WebSocketClientProtocol
+
+from aethergraph.contracts.services.mcp import MCPClientProtocol, MCPResource, MCPTool
+
+
+class WsMCPClient(MCPClientProtocol):
+    def __init__(
+        self,
+        url: str,
+        *,
+        headers: dict[str, str] | None = None,
+        timeout: float = 60.0,
+        ping_interval: float = 20.0,
+        ping_timeout: float = 10.0,
+    ):
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.ping_interval = ping_interval
+        self.ping_timeout = ping_timeout
+
+        self._ws: WebSocketClientProtocol | None = None
+        self._id = 0
+        self._lock = asyncio.Lock()
+        self._ping_task: asyncio.Task | None = None
+
+    async def open(self):
+        if self._ws and not self._ws.closed:
+            return
+        try:
+            # websockets >=14
+            self._ws = await websockets.connect(
+                self.url, additional_headers=self.headers, open_timeout=self.timeout
+            )
+        except TypeError:
+            # likely on websockets <=13
+            self._ws = await websockets.connect(
+                self.url, extra_headers=self.headers, open_timeout=self.timeout
+            )
+
+        self._start_ping()
+
+    async def close(self):
+        if self._ping_task:
+            self._ping_task.cancel()
+            with contextlib.suppress(Exception):
+                await self._ping_task
+            self._ping_task = None
+        if self._ws and not self._ws.closed:
+            await self._ws.close()
+        self._ws = None
+
+    def _start_ping(self):
+        if self.ping_interval <= 0:
+            return
+
+        async def _pinger():
+            try:
+                while self._ws and not self._ws.closed:
+                    await asyncio.sleep(self.ping_interval)
+                    if not self._ws or self._ws.closed:
+                        break
+                    try:
+                        await asyncio.wait_for(self._ws.ping(), timeout=self.ping_timeout)  # type: ignore
+                    except Exception:
+                        break
+            except asyncio.CancelledError:
+                pass
+
+        self._ping_task = asyncio.create_task(_pinger())
+
+    async def _ensure(self):
+        if self._ws is None or self._ws.closed:
+            await self.open()
+
+    async def _rpc(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        await self._ensure()
+        async with self._lock:
+            self._id += 1
+            req = {"jsonrpc": "2.0", "id": self._id, "method": method, "params": params or {}}
+            data = json.dumps(req)
+            try:
+                assert self._ws is not None
+                await asyncio.wait_for(self._ws.send(data), timeout=self.timeout)
+                raw = await asyncio.wait_for(self._ws.recv(), timeout=self.timeout)
+            except Exception:
+                await self.close()
+                await self.open()
+                assert self._ws is not None
+                await asyncio.wait_for(self._ws.send(data), timeout=self.timeout)
+                raw = await asyncio.wait_for(self._ws.recv(), timeout=self.timeout)
+            resp = json.loads(raw)
+            if "error" in resp:
+                raise RuntimeError(str(resp["error"]))
+            return resp.get("result")
+
+    async def list_tools(self) -> list[MCPTool]:
+        return await self._rpc("tools/list")
+
+    async def call(self, tool: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
+        return await self._rpc("tools/call", {"name": tool, "arguments": params or {}})
+
+    async def list_resources(self) -> list[MCPResource]:
+        return await self._rpc("resources/list")
+
+    async def read_resource(self, uri: str) -> dict[str, Any]:
+        return await self._rpc("resources/read", {"uri": uri})

aethergraph/services/memory/bound.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from aethergraph.contracts.services.memory import Event
+from aethergraph.services.memory.facade import MemoryFacade
+from aethergraph.services.memory.resolver import ResolverContext, resolve_params
+
+Value = dict[str, Any]
+
+
+@dataclass
+class BoundMemory:
+    mem: MemoryFacade
+    defaults: dict[str, Any] = (
+        None  # run_id, graph_id, node_id, agent_id (usually injected by NodeContext)
+    )
+
+    async def record(
+        self,
+        *,
+        kind: str,
+        text: str | None = None,
+        severity: int = 2,
+        stage: str | None = None,
+        tags: list[str] | None = None,
+        entities: list[str] | None = None,
+        metrics: dict[str, Any] | None = None,
+        inputs_ref: dict[str, Any] | None = None,
+        outputs_ref: dict[str, Any] | None = None,
+        sources: list[str] | None = None,
+        signal: float | None = None,
+    ) -> Event:
+        base = {
+            **(self.defaults or {}),
+            "kind": kind,
+            "stage": stage,
+            "severity": severity,
+            "tags": tags or [],
+            "entities": entities or [],
+            "inputs_ref": inputs_ref,
+            "outputs_ref": outputs_ref,
+            "sources": sources,
+            "signal": float(
+                signal
+                if signal is not None
+                else self._estimate_signal(text=text, metrics=metrics, severity=severity)
+            ),
+        }
+        return await self.mem.record_raw(base=base, text=text, metrics=metrics)
+
+    async def user(self, text: str):
+        return await self.record(kind="user_msg", text=text, stage="observe")
+
+    async def assistant(self, text: str):
+        return await self.record(kind="assistant_msg", text=text, stage="act")
+
+    async def tool_start(self, note=None):
+        return await self.record(kind="tool_start", text=note, stage="act")
+
+    async def tool_ok(self, note=None, metrics=None):
+        return await self.record(
+            kind="tool_result", text=note, metrics=metrics, stage="observe", severity=3
+        )
+
+    async def tool_error(self, err: Exception):
+        return await self.record(
+            kind="error", text=f"{type(err).__name__}: {err}", severity=5, stage="observe"
+        )
+
+    async def write_result(
+        self,
+        *,
+        topic: str,
+        inputs: list[Value] | None = None,
+        outputs: list[Value] | None = None,
+        tags: list[str] | None = None,
+        metrics: dict[str, float] | None = None,
+        message: str | None = None,
+        severity: int = 3,
+    ) -> Event:
+        return await self.mem.write_result(
+            topic=topic,
+            inputs=inputs or [],
+            outputs=outputs or [],
+            tags=tags,
+            metrics=metrics,
+            message=message,
+            severity=severity,
+        )
+
+    async def resolve(self, params: dict[str, Any]) -> dict[str, Any]:
+        rctx = ResolverContext(mem=self.mem)
+        return await resolve_params(params, rctx)
+
+    # ---- helper ----
+    def _estimate_signal(
+        self, *, text: str | None, metrics: dict[str, Any] | None, severity: int
+    ) -> float:
+        score = 0.15 + 0.1 * severity
+        if text:
+            score += min(len(text) / 400.0, 0.4)
+        if metrics:
+            score += 0.2
+        return max(0.0, min(1.0, score))

aethergraph/services/memory/distillers/episode.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from typing import Any
+
+from aethergraph.contracts.services.memory import Distiller, Event, HotLog, Indices, Persistence
+
+
+def _now_iso() -> str:
+    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+
+
+def _stable_event_id(parts: dict[str, Any]) -> str:
+    blob = json.dumps(parts, sort_keys=True, ensure_ascii=False).encode("utf-8")
+    return hashlib.sha256(blob).hexdigest()[:24]
+
+
+def _episode_uri(sessiorun_idn_id: str, tool: str, run_id: str) -> str:
+    safe = tool.replace("/", "_")
+    return f"file://mem/{run_id}/episodes/{safe}/{run_id}.json"
+
+
+class EpisodeSummarizer(Distiller):
+    """
+    Aggregate all events for (tool, run_id) into a compact episode summary:
+      - sources: event_ids
+      - merged metrics (last-write-wins)
+      - notes: last N textual notes
+    Writes a JSON summary artifact and emits a lightweight run_summary event.
+    """
+
+    def __init__(
+        self, *, include_metrics: bool = True, note_limit: int = 8, note_chars: int = 2000
+    ):
+        self.include_metrics = include_metrics
+        self.note_limit = note_limit
+        self.note_chars = note_chars
+
+    async def distill(
+        self,
+        run_id: str,
+        *,
+        hotlog: HotLog,
+        persistence: Persistence,
+        indices: Indices,
+        tool: str,
+        **kw,
+    ) -> dict[str, Any]:
+        # Pull a reasonable window from hot memory; filter in-process.
+        # (If needed later, add a Persistence scan by day.)
+        events = await hotlog.recent(
+            run_id, kinds=["tool_start", "tool_result", "error", "run_summary"], limit=400
+        )
+
+        eps = [e for e in events if e.run_id == run_id and (e.tool or "") == tool]
+        if not eps:
+            return {}
+
+        srcs: list[str] = []
+        notes: list[str] = []
+        metrics: dict[str, float] = {}
+
+        for e in eps:
+            if e.event_id:
+                srcs.append(e.event_id)
+            if self.include_metrics and e.metrics:
+                metrics.update(e.metrics)  # simple merge; last-write-wins
+            if e.text:
+                notes.append(e.text)
+
+        ts = _now_iso()
+        summary = {
+            "kind": "episode_summary",
+            "run_id": run_id,
+            "tool": tool,
+            "ts": ts,
+            "sources": srcs,
+            "metrics": metrics,
+            "notes": notes[-self.note_limit :],
+        }
+
+        uri = _episode_uri(run_id, tool, run_id)
+        await persistence.save_json(uri, summary)
+
+        # Emit a compact run_summary event
+        compact_text = "\n".join(summary["notes"][-self.note_limit :])[: self.note_chars]
+        evt_base = {
+            "run_id": run_id,
+            "tool": tool,
+            "kind": "run_summary",
+            "severity": 1,
+            "tags": ["summary", "episode"],
+        }
+        eid = _stable_event_id(
+            {
+                "ts": ts,
+                "run_id": run_id,
+                "tool": tool,
+                "kind": "run_summary",
+                "text": compact_text[:200],
+            }
+        )
+        evt = Event(
+            event_id=eid,
+            ts=ts,
+            text=compact_text,
+            metrics={"notes": len(notes)},
+            signal=0.5,
+            **evt_base,
+        )
+        await hotlog.append(run_id, evt, ttl_s=7 * 24 * 3600, limit=1000)
+        await persistence.append_event(run_id, evt)
+
+        return {"uri": uri, "sources": srcs, "metrics": metrics}

aethergraph/services/memory/distillers/rolling.py

@@ -0,0 +1,74 @@
+import time
+from typing import Any
+
+from aethergraph.contracts.services.memory import Distiller, Event, HotLog, Indices, Persistence
+
+
+def _now_iso():
+    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+
+
+def ar_summary_uri(run_id: str, tag: str, ts: str) -> str:
+    # Save summaries under the same base "mem/<run_id>/..." tree as append_event,
+    # but using a file:// URI so FSPersistence can handle it.
+    safe_ts = ts.replace(":", "-")
+    return f"file://mem/{run_id}/summaries/{tag}/{safe_ts}.json"
+
+
+class RollingSummarizer(Distiller):
+    def __init__(
+        self, *, max_turns: int = 20, min_signal: float = 0.25, turn_kinds: list[str] | None = None
+    ):
+        self.max_turns = max_turns
+        self.min_signal = min_signal
+        self.turn_kinds = turn_kinds or ["user_msg", "assistant_msg"]
+
+    async def distill(
+        self, run_id: str, *, hotlog: HotLog, persistence: Persistence, indices: Indices, **kw
+    ) -> dict[str, Any]:
+        turns = await hotlog.recent(run_id, kinds=self.turn_kinds, limit=self.max_turns * 2)
+        kept = [t for t in turns if (t.signal or 0.0) >= self.min_signal]
+        if not kept:
+            return {}
+
+        lines = []
+        srcs: list[str] = []
+        for t in kept[-self.max_turns :]:
+            role = "User" if t.kind == "user_msg" else "Assistant"
+            if t.text:
+                lines.append(f"{role}: {t.text}")
+                srcs.append(t.event_id)
+        digest_text = "\n".join(lines)
+        ts = _now_iso()
+        summary = {
+            "kind": "rolling_summary",
+            "run_id": run_id,
+            "ts": ts,
+            "sources": srcs,
+            "text": digest_text,
+        }
+
+        uri = ar_summary_uri(run_id, "rolling", ts)
+
+        await persistence.save_json(uri=uri, obj=summary)
+
+        evt = Event(
+            event_id="",
+            ts=ts,
+            run_id=run_id,
+            kind="rolling_summary",
+            severity=1,
+            signal=0.5,
+            text=digest_text,
+            metrics={"num_turns": len(kept)},
+            tags=["summary"],
+        )
+
+        from aethergraph.services.memory.facade import stable_event_id
+
+        evt.event_id = stable_event_id(
+            {"ts": ts, "run_id": run_id, "kind": "rolling_summary", "text": digest_text[:200]}
+        )
+        await hotlog.append(run_id, evt, ttl_s=7 * 24 * 3600, limit=1000)
+        await persistence.append_event(run_id, evt)
+        return {"uri": uri, "sources": srcs}