aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/core/runtime/runtime_services.py

@@ -0,0 +1,224 @@
+from collections.abc import Callable
+from contextlib import contextmanager
+from contextvars import ContextVar
+from pathlib import Path
+from typing import Any
+
+from aethergraph.contracts.services.llm import LLMClientProtocol
+from aethergraph.services.llm.generic_client import GenericLLMClient
+
+_current = ContextVar("aeg_services", default=None)
+# process-wide fallback (handles contextvar boundary issues)
+_services_global: Any = None
+
+
+def install_services(services: Any) -> None:
+    global _services_global
+    _services_global = services
+    return _current.set(services)
+
+
+def ensure_services_installed(factory: Callable[[], Any]) -> Any:
+    global _services_global
+    svc = _current.get() or _services_global
+    if svc is None:
+        svc = factory()
+        _services_global = svc
+    _current.set(svc)  # keep ContextVar in sync for this context
+    return svc
+
+
+def current_services() -> Any:
+    svc = _current.get() or _services_global
+    if svc is None:
+        raise RuntimeError(
+            "No services installed. Call install_services(container) at app startup."
+        )
+    return svc
+
+
+@contextmanager
+def use_services(services):
+    tok = _current.set(services)
+    try:
+        yield
+    finally:
+        _current.reset(tok)
+
+
+# --------- Channel service helpers ---------
+def get_channel_service() -> Any:
+    svc = current_services()
+    return svc.channels  # ChannelBus
+
+
+def set_default_channel(key: str) -> None:
+    svc = current_services()
+    svc.channels.set_default_channel_key(key)
+    return
+
+
+def get_default_channel() -> str:
+    svc = current_services()
+    return svc.channels.default_channel_key
+
+
+def set_channel_alias(alias: str, channel_key: str) -> None:
+    svc = current_services()
+    svc.channels.register_alias(alias, channel_key)
+
+
+def register_channel_adapter(name: str, adapter: Any) -> None:
+    svc = current_services()
+    svc.channel.register_adapter(name, adapter)
+
+
+# --------- LLM service helpers ---------
+def get_llm_service() -> Any:
+    svc = current_services()
+    return svc.llm
+
+
+def register_llm_client(
+    profile: str,
+    provider: str,
+    model: str,
+    embed_model: str | None = None,
+    base_url: str | None = None,
+    api_key: str | None = None,
+    timeout: float | None = None,
+) -> None:
+    svc = current_services()
+    client = svc.llm.configure_profile(
+        profile=profile,
+        provider=provider,
+        model=model,
+        embed_model=embed_model,
+        base_url=base_url,
+        api_key=api_key,
+        timeout=timeout,
+    )
+    return client
+
+
+# backend compatibility
+set_llm_client = register_llm_client
+
+
+def set_rag_llm_client(
+    client: LLMClientProtocol | None = None,
+    *,
+    provider: str | None = None,
+    model: str | None = None,
+    embed_model: str | None = None,
+    base_url: str | None = None,
+    api_key: str | None = None,
+    timeout: float | None = None,
+) -> LLMClientProtocol:
+    """Set the LLM client to use for RAG service.
+    If client is provided, use it directly.
+    Otherwise, create a new client using the provided parameters."""
+    svc = current_services()
+    if client is None:
+        if provider is None or model is None or embed_model is None:
+            raise ValueError(
+                "Must provide provider, model, and embed_model to create RAG LLM client"
+            )
+        try:
+            client = GenericLLMClient(
+                provider=provider,
+                model=model,
+                embed_model=embed_model,
+                base_url=base_url,
+                api_key=api_key,
+                timeout=timeout,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to create RAG LLM client: {e}") from e
+
+    svc.rag.set_llm_client(client=client)
+    return client
+
+
+def set_rag_index_backend(
+    *,
+    backend: str | None = None,  # "sqlite" | "faiss"
+    index_path: str | None = None,
+    dim: int | None = None,
+):
+    """
+    Configure the RAG index backend. If backend='faiss' but FAISS is missing,
+    we log a warning and fall back to SQLite automatically.
+    """
+    from aethergraph.services.rag.index_factory import create_vector_index
+
+    svc = current_services()
+    # resolve defaults from settings
+    s = svc.settings.rag  # AppSettings.rag bound into services
+    backend = backend or s.backend
+    index_path = index_path or s.index_path
+    dim = dim if dim is not None else s.dim
+    root = svc.settings.root
+
+    index = create_vector_index(
+        backend=backend, index_path=index_path, dim=dim, root=str(Path(root) / "rag")
+    )
+    svc.rag.set_index_backend(index)
+    return index
+
+
+# --------- Logger helpers ---------
+def current_logger_factory() -> Any:
+    svc = current_services()
+    return svc.logger
+
+
+# --------- External context services ---------
+def register_context_service(name: str, service: Any) -> None:
+    svc = current_services()
+    svc.ext_services[name] = service
+
+
+def get_ext_context_service(name: str) -> Any:
+    svc = current_services()
+    return svc.ext_services.get(name)
+
+
+def list_ext_context_services() -> list[str]:
+    svc = current_services()
+    return list(svc.ext_services.keys())
+
+
+# --------- MCP service helpers ---------
+def set_mcp_service(mcp_service: Any) -> None:
+    svc = current_services()
+    svc.mcp = mcp_service
+
+
+def get_mcp_service() -> Any:
+    svc = current_services()
+    return svc.mcp
+
+
+def register_mcp_client(name: str, client: Any) -> None:
+    svc = current_services()
+    if svc.mcp is None:
+        raise RuntimeError("No MCP service installed. Call set_mcp_service() first.")
+    svc.mcp.register(name, client)
+
+
+def list_mcp_clients() -> list[str]:
+    svc = current_services()
+    if svc.mcp:
+        return svc.mcp.list_clients()
+    return []
+
+
+# --------- Scheduler helpers --------- - (Not used)
+def ensure_global_scheduler_started() -> None:
+    svc = current_services()
+    sched = svc.schedulers.get("global")
+    if sched and not sched.is_running():
+        import asyncio
+
+        asyncio.create_task(sched.run_forever())

aethergraph/core/runtime/wakeup_watcher.py

@@ -0,0 +1,40 @@
+import asyncio
+import time
+
+"""This is a template implementation of a WakeupWatcher that periodically checks for
+Currenlty, we have not materialize the wakeup method in Aethergraph core.
+"""
+
+
+class WakeupWatcher:
+    def __init__(self, cont_store, resume_bus, poll_sec: int = 10):
+        self.cont_store = cont_store
+        self.resume_bus = resume_bus
+        self.poll_sec = poll_sec
+        self._task = None
+        self._stop = asyncio.Event()
+
+    async def start(self):
+        self._task = asyncio.create_task(self._loop())
+
+    async def stop(self):
+        self._stop.set()
+        if self._task:
+            await self._task
+
+    async def _loop(self):
+        while not self._stop.is_set():
+            now = time.time()
+            due = await self.cont_store.list_due_wakeups(now)
+            for cont in due:
+                # Publish to bus; bus routes to the right scheduler via run_id ↦ scheduler mapping
+                await self.resume_bus.post_wakeup(cont.run_id, cont.node_id)
+                # update next_wakeup in the store if needed
+                await self.cont_store.bump_wakeup(cont)
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self.poll_sec)
+            except asyncio.TimeoutError:
+                import logging
+
+                logger = logging.getLogger("aethergraph.core.runtime.wakeup_watcher")
+                logger.info("WakeupWatcher polling for due continuations...")

aethergraph/core/tools/__init__.py

@@ -0,0 +1,10 @@
+from .builtins.toolset import (
+    ask_approval as ask_approval,
+    ask_files as ask_files,
+    ask_text as ask_text,
+    get_latest_uploads as get_latest_uploads,
+    send_buttons as send_buttons,
+    send_file as send_file,
+    send_image as send_image,
+    send_text as send_text,
+)

aethergraph/core/tools/builtins/channel_tools.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from typing import Any
+
+from aethergraph.services.channel.wait_helpers import create_and_notify_continuation
+
+from ...execution.wait_types import WaitSpec
+from ..waitable import DualStageTool, waitable_tool
+
+
+def normalize_approval_result(payload: dict) -> dict:
+    """
+    Normalize approval result from various adapters into a consistent format.
+
+    It assumes the payload may contain:
+    - "approved": bool (explicit approval flag)
+    - "choice": str (the user's choice)
+    It infers "approved" from "choice" if "approved" is not present.
+    """
+    choice = payload.get("choice")
+
+    # infer from options (first = approved)
+    options = payload.get("options") or payload.get("buttons")
+    if not options:
+        prompt = payload.get("prompt")
+        if isinstance(prompt, dict):
+            options = prompt.get("buttons")
+
+    if not options or choice is None:
+        approved = False
+    else:
+        choice_norm = str(choice).strip().lower()
+        first_norm = str(options[0]).strip().lower()
+        approved = choice_norm == first_norm
+
+    return {"approved": approved, "choice": choice}
+
+
+# ----- AskTextTool -----
+class AskText(DualStageTool):
+    outputs = ["text"]
+
+    async def setup(
+        self,
+        prompt: str | None = None,
+        *,
+        silent: bool = False,
+        timeout_s: int = 3600,
+        channel: str | None = None,
+        context,
+    ) -> WaitSpec | dict[str, Any]:
+        token, inline = await create_and_notify_continuation(
+            context=context,
+            kind="user_input",
+            payload={"prompt": prompt, "_force_push": True, "_silent": silent}
+            if prompt
+            else {"_force_push": True, "_silent": silent},
+            timeout_s=timeout_s,
+            channel=channel,
+        )
+        return WaitSpec(
+            channel=channel,
+            token=token,
+            kind="user_input",
+            deadline=timeout_s,
+            notified=True,
+            inline_payload=inline,
+        )
+
+    async def on_resume(self, resume: dict[str, Any], *, context: Any) -> dict[str, Any]:
+        text = resume.get("text", "")
+        return {"text": text}
+
+
+ask_text_ds = waitable_tool(AskText)
+
+
+# ----- WaitText Tool  -----
+class WaitText(DualStageTool):
+    outputs = ["text"]
+
+    async def setup(
+        self,
+        *,
+        timeout_s: int = 3600,
+        channel: str | None = None,
+        context,
+    ) -> WaitSpec | dict[str, Any]:
+        token, inline = await create_and_notify_continuation(
+            context=context,
+            kind="user_input",
+            payload={"prompt": None, "_force_push": True, "_silent": True},
+            timeout_s=timeout_s,
+            channel=channel,
+        )
+        return WaitSpec(
+            channel=channel,
+            token=token,
+            kind="user_input",
+            deadline=timeout_s,
+            notified=True,
+            inline_payload=inline,
+        )
+
+    async def on_resume(self, resume: dict[str, Any], *, context: Any) -> dict[str, Any]:
+        text = resume.get("text", "")
+        return {"text": text}
+
+
+wait_text_ds = waitable_tool(WaitText)
+
+
+# ----- AskApprovalTool -----
+class AskApproval(DualStageTool):
+    outputs = ["approved", "choice"]
+
+    async def setup(
+        self,
+        prompt: str,
+        options: list[str] | tuple[str, ...] = ("Approve", "Reject"),
+        *,
+        timeout_s: int = 3600,
+        channel: str | None = None,
+        context: Any,
+    ) -> WaitSpec | dict[str, Any]:
+        token, inline = await create_and_notify_continuation(
+            context=context,
+            kind="approval",
+            payload={"prompt": {"title": prompt, "buttons": list(options)}, "_force_push": True},
+            timeout_s=timeout_s,
+            channel=channel,
+        )
+        return WaitSpec(
+            channel=channel,
+            token=token,
+            kind="approval",
+            deadline=timeout_s,
+            notified=True,
+            inline_payload=inline,
+        )
+
+    async def on_resume(self, resume: dict[str, Any], *, context: Any) -> dict[str, Any]:
+        return normalize_approval_result(resume)
+
+
+ask_approval_ds = waitable_tool(AskApproval)
+
+
+# ----- AskFiles Tool -----
+class AskFiles(DualStageTool):
+    outputs = ["text", "files"]
+
+    async def setup(
+        self,
+        *,
+        prompt: str,
+        accept: list[str] | None = None,
+        multiple: bool = True,
+        timeout_s: int = 3600,
+        channel: str | None = None,
+        context: Any,
+    ) -> WaitSpec | dict[str, Any]:
+        token, inline = await create_and_notify_continuation(
+            context=context,
+            kind="user_files",
+            payload={
+                "prompt": prompt,
+                "accept": accept or [],
+                "multiple": bool(multiple),
+                "_force_push": True,
+            },
+            timeout_s=timeout_s,
+            channel=channel,
+        )
+        return WaitSpec(
+            channel=channel,
+            token=token,
+            kind="user_files",
+            deadline=timeout_s,
+            notified=True,
+            inline_payload=inline,
+        )
+
+    async def on_resume(self, resume: dict[str, Any], *, context: Any) -> dict[str, Any]:
+        files = resume.get("files", [])
+        if not isinstance(files, list):
+            files = []
+        return {
+            "text": str(resume.get("text", "")),
+            "files": files,
+        }
+
+
+ask_files_ds = waitable_tool(AskFiles)

aethergraph/core/tools/builtins/toolset.py

@@ -0,0 +1,134 @@
+from typing import Any
+
+from aethergraph.contracts.services.channel import Button, FileRef
+
+from ..toolkit import tool
+from .channel_tools import ask_approval_ds, ask_files_ds, ask_text_ds, wait_text_ds
+
+
+@tool(name="ask_text", outputs=["text"])
+async def ask_text(
+    *,
+    resume=None,
+    context=None,
+    prompt: str | None = None,
+    silent: bool = False,
+    timeout_s: int = 3600,
+    channel: str | None = None,
+):
+    return await ask_text_ds(
+        resume=resume,
+        context=context,
+        prompt=prompt,
+        silent=silent,
+        timeout_s=timeout_s,
+        channel=channel,
+    )
+
+
+@tool(name="wait_text", outputs=["text"])
+async def wait_text(
+    *, resume=None, context=None, timeout_s: int = 3600, channel: str | None = None
+):
+    return await wait_text_ds(resume=resume, context=context, timeout_s=timeout_s, channel=channel)
+
+
+@tool(name="ask_approval", outputs=["approved", "choice"])
+async def ask_approval(
+    *,
+    resume=None,
+    context=None,
+    prompt: str,
+    options: list[str] | tuple[str, ...] = ("Approve", "Reject"),
+    timeout_s: int = 3600,
+    channel: str | None = None,
+):
+    return await ask_approval_ds(
+        resume=resume,
+        context=context,
+        prompt=prompt,
+        options=options,
+        timeout_s=timeout_s,
+        channel=channel,
+    )
+
+
+@tool(name="ask_files", outputs=["text", "files"])
+async def ask_files(
+    *,
+    resume=None,
+    context=None,
+    prompt: str,
+    accept: list[str] | None = None,
+    multiple: bool = True,
+    timeout_s: int = 3600,
+    channel: str | None = None,
+):
+    return await ask_files_ds(
+        resume=resume,
+        context=context,
+        prompt=prompt,
+        accept=accept,
+        multiple=multiple,
+        timeout_s=timeout_s,
+        channel=channel,
+    )
+
+
+@tool(name="send_text", outputs=["ok"])
+async def send_text(
+    *, text: str, meta: dict[str, Any] | None = None, channel: str | None = None, context=None
+):
+    ch = context.channel(channel)
+    await ch.send_text(text, meta=meta or {})
+    return {"ok": True}
+
+
+@tool(name="send_image", outputs=["ok"])
+async def send_image(
+    *,
+    url: str | None = None,
+    alt: str = "image",
+    title: str | None = None,
+    channel: str | None = None,
+    context=None,
+):
+    ch = context.channel(channel)
+    await ch.send_image(url=url, alt=alt, title=title)
+    return {"ok": True}
+
+
+@tool(name="send_file", outputs=["ok"])
+async def send_file(
+    *,
+    url: str | None = None,
+    file_bytes: bytes | None = None,
+    filename: str = "file.bin",
+    title: str | None = None,
+    channel: str | None = None,
+    context=None,
+):
+    ch = context.channel(channel)
+    await ch.send_file(url=url, file_bytes=file_bytes, filename=filename, title=title)
+    return {"ok": True}
+
+
+@tool(name="send_buttons", outputs=["ok"])
+async def send_buttons(
+    *,
+    text: str,
+    buttons: list[Button],
+    meta: dict[str, Any] | None = None,
+    channel: str | None = None,
+    context=None,
+):
+    ch = context.channel(channel)
+    await ch.send_buttons(text=text, buttons=buttons, meta=meta or {})
+    return {"ok": True}
+
+
+@tool(name="get_lastest_uploads", outputs=["files"])
+async def get_latest_uploads(*, clear: bool = True, context) -> list[FileRef]:
+    ch = context.channel()  # any channel session will expose the same get_latest_uploads
+    files = await ch.get_latest_uploads(clear=clear)
+    return {"files": files}