aethergraph 0.1.0a1__py3-none-any.whl → 0.1.0a2__py3-none-any.whl

aethergraph/services/mcp/service.py

@@ -15,45 +15,180 @@ class MCPService:
     """
 
     def __init__(self, clients: dict[str, MCPClientProtocol] | None = None, *, secrets=None):
+        """
+        Initialize the MCPService with optional clients and secrets provider.
+
+        Examples:
+            Basic usage with no clients:
+            ```python
+            service = MCPService()
+            ```
+
+            With pre-registered clients:
+            ```python
+            service = MCPService(clients={"default": my_client})
+            ```
+
+        Args:
+            clients: Optional dictionary mapping names to MCPClientProtocol instances.
+            secrets: Optional secrets provider (not implemented here).
+        """
         self._clients: dict[str, MCPClientProtocol] = clients or {}
         self._secrets = secrets  # optional (Secrets provider) Not implemented here
 
     # ---- registration ----
     def register(self, name: str, client: MCPClientProtocol) -> None:
+        """
+        Register a new MCP client under a given name.
+
+        Examples:
+            ```python
+            context.mcp().register("myserver", my_client)
+            ```
+
+        Args:
+            name: The name to register the client under.
+            client: The MCPClientProtocol instance to register.
+        """
         self._clients[name] = client
 
     def remove(self, name: str) -> None:
+        """
+        Remove a registered MCP client by name.
+
+        Examples:
+            ```python
+            context.mcp().remove("myserver")
+            ```
+
+        Args:
+            name: The name of the client to remove.
+        """
         self._clients.pop(name, None)
 
     def has(self, name: str) -> bool:
+        """
+        Check if a client with the given name is registered.
+
+        Examples:
+            ```python
+            if context.mcp().has("default"):
+                print("Client exists")
+            ```
+
+        Args:
+            name: The name to check.
+
+        Returns:
+            bool: True if the client exists, False otherwise.
+        """
         return name in self._clients
 
     def names(self) -> list[str]:
+        """
+        Get a list of all registered client names.
+
+        Examples:
+            ```python
+            names = context.mcp().names()
+            ```
+
+        Returns:
+            list[str]: List of registered client names.
+        """
         return list(self._clients.keys())
 
     def list_clients(self) -> list[str]:
+        """
+        List all registered client names.
+
+        Examples:
+            ```python
+            clients = context.mcp().list_clients()
+            ```
+
+        Returns:
+            list[str]: List of registered client names.
+        """
         return list(self._clients.keys())
 
     def get(self, name: str = "default") -> MCPClientProtocol:
+        """
+        Retrieve a registered MCP client by name.
+
+        Examples:
+            ```python
+            client = context.mcp().get("default")
+            ```
+
+        Args:
+            name: The name of the client to retrieve.
+
+        Returns:
+            MCPClientProtocol: The registered client.
+
+        Raises:
+            KeyError: If the client is not found.
+        """
         if name not in self._clients:
             raise KeyError(f"Unknown MCP server '{name}'")
         return self._clients[name]
 
     # ---- lifecycle ----
     async def open(self, name: str) -> None:
+        """
+        Open the connection for a specific MCP client.
+
+        Examples:
+            ```python
+            await context.mcp().open("default")
+            ```
+
+        Args:
+            name: The name of the client to open.
+        """
         await self.get(name).open()
 
     async def close(self, name: str) -> None:
+        """
+        Close the connection for a specific MCP client.
+
+        Examples:
+            ```python
+            await context.mcp().close("default")
+            ```
+
+        Args:
+            name: The name of the client to close.
+        """
         try:
             await self.get(name).close()
         except Exception:
             logger.warning(f"Failed to close MCP client '{name}'")
 
     async def open_all(self) -> None:
+        """
+        Open all registered MCP client connections.
+
+        Examples:
+            ```python
+            await context.mcp().open_all()
+            ```
+
+        """
         for n in self._clients:
             await self._clients[n].open()
 
     async def close_all(self) -> None:
+        """
+        Close all registered MCP client connections.
+
+        Examples:
+            ```python
+            await context.mcp().close_all()
+            ```
+
+        """
         for n in self._clients:
             try:
                 await self._clients[n].close()
@@ -64,29 +199,103 @@ class MCPService:
     async def call(
         self, name: str, tool: str, params: dict[str, Any] | None = None
     ) -> dict[str, Any]:
+        """
+        Call a tool on a specific MCP client, opening the connection if needed.
+
+        Examples:
+            ```python
+            result = await context.mcp().call("default", "sum", {"a": 1, "b": 2})
+            ```
+
+        Args:
+            name: The name of the client to use.
+            tool: The tool name to call.
+            params: Optional dictionary of parameters for the tool.
+
+        Returns:
+            dict[str, Any]: The result from the tool call.
+        """
         # lazy-open on first use; clients themselves also lazy-reconnect
         c = self.get(name)
         await c.open()
         return await c.call(tool, params or {})
 
     async def list_tools(self, name: str) -> list[MCPTool]:
+        """
+        List all tools available on a specific MCP client.
+
+        Examples:
+            ```python
+            tools = await context.mcp().list_tools("default")
+            ```
+
+        Args:
+            name: The name of the client.
+
+        Returns:
+            list[MCPTool]: List of available tools.
+        """
         c = self.get(name)
         await c.open()
         return await c.list_tools()
 
     async def list_resources(self, name: str) -> list[MCPResource]:
+        """
+        List all resources available on a specific MCP client.
+
+        Examples:
+            ```python
+            resources = await context.mcp().list_resources("default")
+            ```
+
+        Args:
+            name: The name of the client.
+
+        Returns:
+            list[MCPResource]: List of available resources.
+        """
         c = self.get(name)
         await c.open()
         return await c.list_resources()
 
     async def read_resource(self, name: str, uri: str) -> dict[str, Any]:
+        """
+        Read a resource from a specific MCP client.
+
+        Examples:
+            ```python
+            data = await context.mcp().read_resource("default", "resource://foo/bar")
+            ```
+
+        Args:
+            name: The name of the client.
+            uri: The URI of the resource to read.
+
+        Returns:
+            dict[str, Any]: The resource data.
+        """
         c = self.get(name)
         await c.open()
         return await c.read_resource(uri)
 
     # ---- optional secrets helpers  ----
     def set_header(self, name: str, key: str, value: str) -> None:
-        """For ws clients: set/override a header at runtime (demo/notebook UX)."""
+        """
+        Set or override a header for a websocket client at runtime.
+
+        Examples:
+            ```python
+            context.mcp().set_header("default", "Authorization", "Bearer token")
+            ```
+
+        Args:
+            name: The name of the client.
+            key: The header key.
+            value: The header value.
+
+        Raises:
+            RuntimeError: If the client does not support headers.
+        """
         c = self.get(name)
         # duck-typing for ws client
         if hasattr(c, "headers") and isinstance(c.headers, dict):  # type: ignore[attr-defined]
@@ -95,6 +304,21 @@ class MCPService:
             raise RuntimeError(f"MCP '{name}' does not support headers")
 
     def persist_secret(self, secret_name: str, value: str) -> None:
+        """
+        Persist a secret using the configured secrets provider.
+
+        Examples:
+            ```python
+            context.mcp().persist_secret("API_KEY", "my-secret-value")
+            ```
+
+        Args:
+            secret_name: The name of the secret.
+            value: The value to persist.
+
+        Raises:
+            RuntimeError: If the secrets provider is not writable.
+        """
         if not self._secrets or not hasattr(self._secrets, "set"):
             raise RuntimeError("Secrets provider is not writable")
         self._secrets.set(secret_name, value)  # type: ignore

aethergraph/services/mcp/stdio_client.py

@@ -7,13 +7,48 @@ from aethergraph.contracts.services.mcp import MCPClientProtocol, MCPResource, M
 
 
 class StdioMCPClient(MCPClientProtocol):
+    """
+    Initialize the MCP client service to communicate with a subprocess over stdio using JSON-RPC 2.0.
+
+    This class launches a subprocess (typically an MCP server), manages its lifecycle, and provides
+    asynchronous methods to interact with it using JSON-RPC 2.0 over standard input/output streams.
+    It handles command execution, environment setup, request/response serialization, and concurrency
+    control for safe multi-call usage.
+
+    Examples:
+        Basic usage with default environment:
+        ```python
+        from aethergraph.services.mcp import StdioMCPClient
+        client = StdioMCPClient(["python", "mcp_server.py"])
+        await client.open()
+        tools = await client.list_tools()
+        await client.close()
+        ```
+
+        Custom environment and timeout:
+        ```python
+        from aethergraph.services.mcp import StdioMCPClient
+        client = StdioMCPClient(
+            ["python", "mcp_server.py"],
+            env={"MY_ENV_VAR": "value"},
+            timeout=30.0
+        )
+        ```
+
+    Args:
+        cmd: Command to start the MCP server subprocess (list of str).
+        env: Optional dictionary of environment variables for the subprocess.
+        timeout: Timeout in seconds for each RPC call.
+
+    Returns:
+        None: Initializes the StdioMCPClient instance.
+
+    Notes:
+        - The subprocess should adhere to the JSON-RPC 2.0 specification over stdio.
+        - Ensure proper error handling in the subprocess to avoid deadlocks.
+    """
+
     def __init__(self, cmd: list[str], env: dict[str, str] | None = None, timeout: float = 60.0):
-        """MCP client that talks to a subprocess over stdio using JSON-RPC 2.0.
-        Args:
-            cmd: Command to start the MCP server subprocess (list of str).
-            env: Optional environment variables to set for the subprocess.
-            timeout: Timeout in seconds for each RPC call.
-        """
         self.cmd, self.env, self.timeout = cmd, env or {}, timeout
         self.proc = None
         self._id = 0

aethergraph/services/mcp/ws_client.py

@@ -3,15 +3,57 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import json
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import websockets
-from websockets.client import WebSocketClientProtocol
+
+if TYPE_CHECKING:
+    # only imported for static type checkers; no runtime import, no warning
+    from websockets.client import WebSocketClientProtocol
 
 from aethergraph.contracts.services.mcp import MCPClientProtocol, MCPResource, MCPTool
 
 
 class WsMCPClient(MCPClientProtocol):
+    """
+    Initialize the WebSocket MCP client with URL, headers, and connection parameters.
+    This class manages a WebSocket connection to an MCP (Modular Control Protocol) server,
+    handling connection lifecycle, pinging for keepalive, and concurrency for sending JSON-RPC
+    requests. It provides methods to list tools, call tools, list resources, and read resources
+    via the MCP protocol.
+
+    Examples:
+        Basic usage with default headers:
+        ```python
+        from aethergraph.services.mcp import WsMCPClient
+        client = WsMCPClient("wss://mcp.example.com/ws")
+        await client.open()
+        tools = await client.list_tools()
+        await client.close()
+        ```
+
+        Custom headers and ping interval:
+        ```python
+        from aethergraph.services.mcp import WsMCPClient
+        client = WsMCPClient(
+            "wss://mcp.example.com/ws",
+            headers={"Authorization": "Bearer <token>"},
+            ping_interval=10.0,
+            timeout=30.0
+        await client.open()
+        ```
+
+    Args:
+        url: The WebSocket URL of the MCP server (e.g., "wss://mcp.example.com/ws").
+        headers: Optional dictionary of additional HTTP headers to include in the WebSocket handshake.
+        timeout: Maximum time (in seconds) to wait for connection and responses.
+        ping_interval: Interval (in seconds) between WebSocket ping frames for keepalive.
+        ping_timeout: Maximum time (in seconds) to wait for a ping response before considering the connection dead.
+
+    Returns:
+        None: Initializes the WsMCPClient instance and prepares internal state.
+    """
+
     def __init__(
         self,
         url: str,

aethergraph/services/memory/distillers/llm_long_term.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+import json
+from typing import Any
+
+from aethergraph.contracts.services.llm import LLMClientProtocol
+from aethergraph.contracts.services.memory import Distiller, Event, HotLog, Indices, Persistence
+from aethergraph.contracts.storage.doc_store import DocStore
+
+# metering
+from aethergraph.core.runtime.runtime_metering import current_meter_context, current_metering
+from aethergraph.services.memory.facade.utils import now_iso, stable_event_id
+from aethergraph.services.memory.utils import _summary_doc_id
+
+
+class LLMLongTermSummarizer(Distiller):
+    """
+    LLM-based long-term summarizer.
+
+    Flow:
+      1) Pull recent events from HotLog.
+      2) Filter by kind/tag/signal.
+      3) Build a prompt that shows the most important events as a transcript.
+      4) Call LLM to generate a structured summary.
+      5) Save summary JSON via Persistence.save_json(uri).
+      6) Emit a long_term_summary Event pointing to summary_uri.
+
+    This is complementary to RAG:
+      - LLM distiller compresses sequences into a digest.
+      - RAG uses many such digests + raw docs for retrieval.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm: LLMClientProtocol,
+        summary_kind: str = "long_term_summary",
+        summary_tag: str = "session",
+        include_kinds: list[str] | None = None,
+        include_tags: list[str] | None = None,
+        max_events: int = 200,
+        min_signal: float = 0.0,
+        model: str | None = None,
+    ):
+        self.llm = llm
+        self.summary_kind = summary_kind
+        self.summary_tag = summary_tag
+        self.include_kinds = include_kinds
+        self.include_tags = include_tags
+        self.max_events = max_events
+        self.min_signal = min_signal
+        self.model = model  # optional model override
+
+    def _filter_events(self, events: Iterable[Event]) -> list[Event]:
+        out: list[Event] = []
+        kinds = set(self.include_kinds) if self.include_kinds else None
+        tags = set(self.include_tags) if self.include_tags else None
+
+        for e in events:
+            if kinds is not None and e.kind not in kinds:
+                continue
+            if tags is not None:
+                if not e.tags:
+                    continue
+                if not tags.issubset(set(e.tags)):
+                    continue
+            if (e.signal or 0.0) < self.min_signal:
+                continue
+            out.append(e)
+        return out
+
+    def _build_prompt(self, events: list[Event]) -> list[dict[str, str]]:
+        """
+        Convert events into a chat-style context for summarization.
+
+        We keep it model-agnostic: a list of {role, content} messages.
+        """
+        lines: list[str] = []
+
+        for e in events:
+            role = e.stage or e.kind or "event"
+            if e.text:
+                lines.append(f"[{role}] {e.text}")
+
+        transcript = "\n".join(lines)
+
+        system = (
+            "You are a log summarizer for an agent's memory. "
+            "Given a chronological transcript of events, produce a concise summary "
+            "of what happened, key themes, important user facts, and open TODOs."
+        )
+
+        user = (
+            "Here is the recent event transcript:\n\n"
+            f"{transcript}\n\n"
+            "Return a JSON object with keys: "
+            "`summary` (string), "
+            "`key_facts` (list of strings), "
+            "`open_loops` (list of strings)."
+            "Do not use markdown or include explanations or context outside the JSON."
+        )
+
+        return [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ]
+
+    async def distill(
+        self,
+        run_id: str,
+        timeline_id: str,
+        scope_id: str = None,
+        *,
+        hotlog: HotLog,
+        persistence: Persistence,
+        indices: Indices,
+        docs: DocStore,
+        **kw: Any,
+    ) -> dict[str, Any]:
+        # 1) fetch more events than needed, then filter
+        raw = await hotlog.recent(timeline_id, kinds=None, limit=self.max_events * 2)
+        kept = self._filter_events(raw)
+        if not kept:
+            return {}
+
+        kept = kept[-self.max_events :]
+        first_ts = kept[0].ts
+        last_ts = kept[-1].ts
+
+        # 2) Build prompt and call LLM
+        messages = self._build_prompt(kept)
+
+        # LLMClientProtocol: assume chat(...) returns (text, usage)
+        summary_json_str, usage = await self.llm.chat(
+            messages,
+        )
+
+        # 3) Parse LLM JSON response
+        try:
+            payload = json.loads(summary_json_str)
+        except Exception:
+            payload = {
+                "summary": summary_json_str,
+                "key_facts": [],
+                "open_loops": [],
+            }
+        ts = now_iso()
+
+        summary_obj = {
+            "type": self.summary_kind,
+            "version": 1,
+            "run_id": run_id,
+            "scope_id": scope_id or run_id,
+            "summary_tag": self.summary_tag,
+            "ts": ts,
+            "time_window": {"from": first_ts, "to": last_ts},
+            "num_events": len(kept),
+            "source_event_ids": [e.event_id for e in kept],
+            "summary": payload.get("summary", ""),
+            "key_facts": payload.get("key_facts", []),
+            "open_loops": payload.get("open_loops", []),
+            "llm_usage": usage,
+            "llm_model": self.llm.model if hasattr(self.llm, "model") else None,
+        }
+
+        scope = scope_id or run_id
+        doc_id = _summary_doc_id(scope, self.summary_tag, ts)
+        await docs.put(doc_id, summary_obj)
+
+        # 4) Emit summary Event with preview + uri in data
+        text = summary_obj["summary"] or ""
+        preview = text[:2000] + (" …[truncated]" if len(text) > 2000 else "")
+
+        evt = Event(
+            event_id="",
+            ts=ts,
+            run_id=run_id,
+            scope_id=scope,
+            kind=self.summary_kind,
+            stage="summary_llm",
+            text=preview,
+            tags=["summary", "llm", self.summary_tag],
+            data={
+                "summary_doc_id": doc_id,
+                "summary_tag": self.summary_tag,
+                "time_window": summary_obj["time_window"],
+                "num_events": len(kept),
+            },
+            metrics={"num_events": len(kept)},
+            severity=2,
+            signal=0.7,
+        )
+
+        evt.event_id = stable_event_id(
+            {
+                "ts": ts,
+                "run_id": run_id,
+                "kind": self.summary_kind,
+                "summary_tag": self.summary_tag,
+                "preview": preview[:200],
+            }
+        )
+
+        await hotlog.append(timeline_id, evt, ttl_s=7 * 24 * 3600, limit=1000)
+        await persistence.append_event(timeline_id, evt)
+
+        # Metering: record summary event
+        try:
+            meter = current_metering()
+            ctx = current_meter_context.get()
+            user_id = ctx.get("user_id")
+            org_id = ctx.get("org_id")
+
+            await meter.record_event(
+                user_id=user_id,
+                org_id=org_id,
+                run_id=run_id,
+                scope_id=scope,
+                kind=f"memory.{self.summary_kind}",  # e.g. "memory.long_term_summary"
+            )
+        except Exception:
+            import logging
+
+            logger = logging.getLogger("aethergraph.services.memory.distillers.llm_long_term")
+            logger.error("Failed to record metering event for long_term_summary")
+
+        return {
+            "summary_doc_id": doc_id,
+            "summary_kind": self.summary_kind,
+            "summary_tag": self.summary_tag,
+            "time_window": summary_obj["time_window"],
+            "num_events": len(kept),
+        }