aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/services/continuations/stores/inmem_store.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from datetime import datetime, timezone
+import hashlib
+import hmac
+import os
+import threading
+
+from ..continuation import Continuation, Correlator
+
+
+class InMemoryContinuationStore:  # implements AsyncContinuationStore
+    def __init__(self, secret: bytes | None = None):
+        self.secret = secret or os.urandom(32)
+        self._by_token: dict[str, Continuation] = {}
+        self._by_run_node: dict[tuple[str, str], Continuation] = {}
+        self._corr_index: dict[tuple[str, str, str, str], list[str]] = defaultdict(list)
+        self._lock = threading.RLock()
+
+    def _hmac(self, *parts: str) -> str:
+        h = hmac.new(self.secret, digestmod=hashlib.sha256)
+        for part in parts:
+            h.update(part.encode("utf-8"))
+        return h.hexdigest()
+
+    async def mint_token(self, run_id: str, node_id: str, attempts: int) -> str:
+        return self._hmac(run_id, node_id, str(attempts), os.urandom(8).hex())
+
+    async def save(self, cont: Continuation) -> None:
+        with self._lock:
+            self._by_token[cont.token] = cont
+            self._by_run_node[(cont.run_id, cont.node_id)] = cont
+
+    async def get(self, run_id: str, node_id: str) -> Continuation | None:
+        with self._lock:
+            return self._by_run_node.get((run_id, node_id))
+
+    async def delete(self, run_id: str, node_id: str) -> None:
+        with self._lock:
+            c = self._by_run_node.pop((run_id, node_id), None)
+            if c:
+                self._by_token.pop(c.token, None)
+
+    async def get_by_token(self, token: str) -> Continuation | None:
+        with self._lock:
+            return self._by_token.get(token)
+
+    async def mark_closed(self, token: str) -> None:
+        with self._lock:
+            c = self._by_token.get(token)
+            if c:
+                c.closed = True
+
+    async def verify_token(self, run_id: str, node_id: str, token: str) -> bool:
+        c = await self.get(run_id, node_id)
+        return bool(c and hmac.compare_digest(c.token, token))
+
+    async def bind_correlator(self, *, token: str, corr: Correlator) -> None:
+        key = corr.key()
+        with self._lock:
+            toks = self._corr_index[key]
+            if token not in toks:
+                toks.append(token)
+
+    async def find_by_correlator(self, *, corr: Correlator) -> Continuation | None:
+        with self._lock:
+            toks = list(self._corr_index.get(corr.key()) or [])
+        for t in reversed(toks):
+            c = await self.get_by_token(t)
+            if c and not c.closed:
+                if c.deadline and datetime.now(timezone.utc) > c.deadline.astimezone(timezone.utc):
+                    continue
+                return c
+        return None
+
+    async def last_open(self, *, channel: str, kind: str) -> Continuation | None:
+        with self._lock:
+            for c in reversed(list(self._by_token.values())):
+                if not c.closed and c.channel == channel and c.kind == kind:
+                    return c
+        return None
+
+    async def list_waits(self) -> list[dict]:
+        with self._lock:
+            return [c.to_dict() for c in self._by_token.values() if not c.closed]
+
+    async def clear(self) -> None:
+        with self._lock:
+            self._by_token.clear()
+            self._by_run_node.clear()
+            self._corr_index.clear()
+
+    async def alias_for(self, token: str) -> str | None:
+        return token[:24]

aethergraph/services/eventbus/inmem.py

@@ -0,0 +1,21 @@
+# services/eventbus/inmem.py
+import asyncio
+from collections import defaultdict
+import threading
+
+
+class InMemoryEventBus:
+    def __init__(self):
+        self._subs = defaultdict(list)  # topic -> [handlers]
+        self._lock = threading.RLock()  # TODO: check if we need thread safety
+
+    async def publish(self, topic, event):
+        # fanout without blocking publishers
+        with self._lock:
+            subs = list(self._subs.get(topic, []))
+        for h in subs:
+            asyncio.create_task(h(event))  # fire-and-forget
+
+    async def subscribe(self, topic, handler):
+        with self._lock:
+            self._subs[topic].append(handler)

aethergraph/services/features/static.py

@@ -0,0 +1,10 @@
+# services/features/static.py
+class StaticFeatures:
+    def __init__(self, flags: dict[str, bool]):
+        self._f = set(k for k, v in flags.items() if v)
+
+    def has(self, name: str) -> bool:
+        return name in self._f
+
+    def all(self) -> set[str]:
+        return set(self._f)

aethergraph/services/kv/ephemeral.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import threading
+import time
+from typing import Any
+
+
+@dataclass
+class KVEntry:
+    value: Any
+    expire_at: float | None = None
+
+
+class EphemeralKV:
+    """Process-local, transient KV (not for blobs)."""
+
+    def __init__(self, *, prefix: str = "") -> None:
+        self._data: dict[str, KVEntry] = {}
+        self._lock = threading.RLock()
+        self._prefix = prefix
+
+    def _k(self, k: str) -> str:
+        return f"{self._prefix}{k}" if self._prefix else k
+
+    async def get(self, key: str, default: Any = None) -> Any:
+        k = self._k(key)
+        with self._lock:
+            e = self._data.get(k)
+            if not e:
+                return default
+            if e.expire_at and e.expire_at < time.time():
+                del self._data[k]
+                return default
+            return e.value
+
+    async def set(self, key: str, value: Any, *, ttl_s: int | None = None) -> None:
+        k = self._k(key)
+        with self._lock:
+            self._data[k] = KVEntry(value=value, expire_at=(time.time() + ttl_s) if ttl_s else None)
+
+    async def delete(self, key: str) -> None:
+        k = self._k(key)
+        with self._lock:
+            self._data.pop(k, None)
+
+    async def list_append_unique(
+        self, key: str, items: list[dict], *, id_key: str = "id", ttl_s: int | None = None
+    ) -> list[dict]:
+        k = self._k(key)
+        with self._lock:
+            cur = list(self._data.get(k, KVEntry([])).value or [])
+            seen = {x.get(id_key) for x in cur if isinstance(x, dict)}
+            cur.extend([x for x in items if isinstance(x, dict) and x.get(id_key) not in seen])
+            self._data[k] = KVEntry(value=cur, expire_at=(time.time() + ttl_s) if ttl_s else None)
+            return cur
+
+    async def list_pop_all(self, key: str) -> list:
+        k = self._k(key)
+        with self._lock:
+            e = self._data.pop(k, None)
+            return list(e.value) if e and isinstance(e.value, list) else []
+
+    # Optional helpers
+    async def mget(self, keys: list[str]) -> list[Any]:
+        return [await self.get(k) for k in keys]
+
+    async def mset(self, kv: dict[str, Any], *, ttl_s: int | None = None) -> None:
+        for k, v in kv.items():
+            await self.set(k, v, ttl_s=ttl_s)
+
+    async def expire(self, key: str, ttl_s: int) -> None:
+        k = self._k(key)
+        with self._lock:
+            e = self._data.get(k)
+            if e:
+                e.expire_at = time.time() + ttl_s
+
+    async def purge_expired(self, limit: int = 1000) -> int:
+        n = 0
+        now = time.time()
+        with self._lock:
+            for k in list(self._data.keys()):
+                if n >= limit:
+                    break
+                e = self._data.get(k)
+                if e and e.expire_at and e.expire_at < now:
+                    self._data.pop(k, None)
+                    n += 1
+        return n

aethergraph/services/kv/factory.py

@@ -0,0 +1,27 @@
+import os
+
+from .ephemeral import EphemeralKV
+from .layered import LayeredKV
+from .sqlite_kv import SQLiteKV
+
+
+def make_kv():
+    kind = (os.getenv("KV_BACKEND", "layered")).lower()
+    prefix = os.getenv("KV_PREFIX", "")  # e.g., tenant/project scoping
+
+    if kind == "ephemeral":
+        return EphemeralKV(prefix=prefix)
+
+    if kind == "sqlite":
+        path = os.getenv("KV_SQLITE_PATH", "./artifacts/kv.sqlite")
+        return SQLiteKV(path, prefix=prefix)
+
+    if kind == "layered":
+        cache = EphemeralKV(prefix=prefix)
+        durable = SQLiteKV(os.getenv("KV_SQLITE_PATH", "./artifacts/kv.sqlite"), prefix=prefix)
+        return LayeredKV(cache, durable)
+
+    # (future) cloud:
+    # if kind == "redis": return RedisKV(...)
+
+    raise ValueError(f"Unknown KV_BACKEND={kind}")

aethergraph/services/kv/layered.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+
+class LayeredKV:
+    """Layered KV combining a fast ephemeral cache (e.g. in-memory) with a durable backend (e.g. SQLite, Redis).
+    Gets first check the cache, then the durable store; sets write to both.
+    List operations invalidate the cache to avoid staleness.
+    """
+
+    def __init__(self, cache, durable):
+        self.cache = cache  # EphemeralKV
+        self.durable = durable  # SQLiteKV / RedisKV
+
+    async def get(self, key: str, default: Any = None) -> Any:
+        v = await self.cache.get(key, None)
+        if v is not None:
+            return v
+        v = await self.durable.get(key, default)
+        if v is not None:
+            await self.cache.set(key, v, ttl_s=5)  # short cache
+        return v
+
+    async def set(self, key: str, value: Any, *, ttl_s: int | None = None) -> None:
+        await self.durable.set(key, value, ttl_s=ttl_s)
+        await self.cache.set(key, value, ttl_s=min(ttl_s or 5, 5))
+
+    async def delete(self, key: str) -> None:
+        await self.durable.delete(key)
+        await self.cache.delete(key)
+
+    async def list_append_unique(self, *a, **k):
+        cur = await self.durable.list_append_unique(*a, **k)
+        await self.cache.delete(a[0])  # invalidate
+        return cur
+
+    async def list_pop_all(self, key: str) -> list:
+        items = await self.durable.list_pop_all(key)
+        await self.cache.delete(key)
+        return items
+
+    # mget/mset/expire/purge can just forward as needed

aethergraph/services/kv/sqlite_kv.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+import threading
+import time
+from typing import Any
+
+
+class SQLiteKV:
+    """
+    Durable KV with TTL (JSON values).
+    Thread-safe via RLock; async callers can await these methods safely.
+    """
+
+    def __init__(self, path: str, *, prefix: str = ""):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        self._db = sqlite3.connect(path, check_same_thread=False, isolation_level=None)
+        self._db.execute("PRAGMA journal_mode=WAL;")
+        self._db.execute("PRAGMA synchronous=NORMAL;")
+        self._db.execute("""
+            CREATE TABLE IF NOT EXISTS kv (
+              k TEXT PRIMARY KEY,
+              v TEXT,
+              expire_at REAL
+            )
+        """)
+        self._db.execute("CREATE INDEX IF NOT EXISTS kv_exp_idx ON kv(expire_at);")
+        self._lock = threading.RLock()
+        self._prefix = prefix
+
+    def _k(self, k: str) -> str:
+        return f"{self._prefix}{k}" if self._prefix else k
+
+    async def get(self, key: str, default: Any = None) -> Any:
+        k = self._k(key)
+        with self._lock:
+            row = self._db.execute("SELECT v, expire_at FROM kv WHERE k=?", (k,)).fetchone()
+            if not row:
+                return default
+            v_txt, exp = row
+            if exp and exp < time.time():
+                self._db.execute("DELETE FROM kv WHERE k=?", (k,))
+                return default
+            try:
+                return json.loads(v_txt)
+            except Exception:
+                return default
+
+    async def set(self, key: str, value: Any, *, ttl_s: int | None = None) -> None:
+        k = self._k(key)
+        with self._lock:
+            exp = (time.time() + ttl_s) if ttl_s else None
+            v_txt = json.dumps(value)
+            self._db.execute(
+                "INSERT INTO kv(k,v,expire_at) VALUES(?,?,?) "
+                "ON CONFLICT(k) DO UPDATE SET v=excluded.v, expire_at=excluded.expire_at",
+                (k, v_txt, exp),
+            )
+
+    async def delete(self, key: str) -> None:
+        k = self._k(key)
+        with self._lock:
+            self._db.execute("DELETE FROM kv WHERE k=?", (k,))
+
+    async def list_append_unique(
+        self, key: str, items: list, *, id_key: str = "id", ttl_s: int | None = None
+    ) -> list:
+        k = self._k(key)
+        with self._lock:
+            row = self._db.execute("SELECT v FROM kv WHERE k=?", (k,)).fetchone()
+            cur = []
+            if row and row[0]:
+                try:
+                    cur = json.loads(row[0])
+                except Exception:
+                    cur = []
+            seen = {x.get(id_key) for x in cur if isinstance(x, dict)}
+            cur.extend([x for x in items if isinstance(x, dict) and x.get(id_key) not in seen])
+            exp = (time.time() + ttl_s) if ttl_s else None
+            self._db.execute(
+                "INSERT INTO kv(k,v,expire_at) VALUES(?,?,?) "
+                "ON CONFLICT(k) DO UPDATE SET v=excluded.v, expire_at=excluded.expire_at",
+                (k, json.dumps(cur), exp),
+            )
+            return cur
+
+    async def list_pop_all(self, key: str) -> list:
+        k = self._k(key)
+        with self._lock:
+            row = self._db.execute("SELECT v FROM kv WHERE k=?", (k,)).fetchone()
+            self._db.execute("DELETE FROM kv WHERE k=?", (k,))
+            if not row or not row[0]:
+                return []
+            try:
+                val = json.loads(row[0])
+                return list(val) if isinstance(val, list) else []
+            except Exception:
+                return []
+
+    # Optional helpers
+    async def mget(self, keys: list[str]) -> list[Any]:
+        out = []
+        for k in keys:
+            out.append(await self.get(k))
+        return out
+
+    async def mset(self, kv: dict[str, Any], *, ttl_s: int | None = None) -> None:
+        for k, v in kv.items():
+            await self.set(k, v, ttl_s=ttl_s)
+
+    async def expire(self, key: str, ttl_s: int) -> None:
+        k = self._k(key)
+        with self._lock:
+            self._db.execute("UPDATE kv SET expire_at=? WHERE k=?", (time.time() + ttl_s, k))
+
+    async def purge_expired(self, limit: int = 1000) -> int:
+        with self._lock:
+            now = time.time()
+            # sqlite lacks DELETE .. LIMIT in older versions; do it in two steps
+            rows = self._db.execute(
+                "SELECT k FROM kv WHERE expire_at IS NOT NULL AND expire_at < ? LIMIT ?",
+                (now, limit),
+            ).fetchall()
+            for (k,) in rows:
+                self._db.execute("DELETE FROM kv WHERE k=?", (k,))
+            return len(rows)

aethergraph/services/llm/factory.py

@@ -0,0 +1,157 @@
+import logging
+import os
+
+from pydantic import SecretStr
+
+from aethergraph.config.llm import LLMProfile, LLMSettings
+
+from ..secrets.base import Secrets
+from .generic_client import GenericLLMClient
+from .providers import Provider
+
+
+def _resolve_key(direct: SecretStr | None, ref: str | None, secrets: Secrets) -> str | None:
+    if direct:
+        return direct.get_secret_value()
+    if ref:
+        return secrets.get(ref)
+    return None
+
+
+def _provider_default_base_url(provider: Provider) -> str | None:
+    # Fallback base URLs if not given in config or env
+    if provider == "openai":
+        return "https://api.openai.com/v1"
+    if provider == "azure":
+        # Must still rely on env/config for endpoint
+        return os.getenv("AZURE_OPENAI_ENDPOINT", "").rstrip("/") or None
+    if provider == "anthropic":
+        return "https://api.anthropic.com"
+    if provider == "google":
+        return "https://generativelanguage.googleapis.com"
+    if provider == "openrouter":
+        return "https://openrouter.ai/api/v1"
+    if provider == "lmstudio":
+        return os.getenv("LMSTUDIO_BASE_URL", "http://localhost:1234/v1")
+    if provider == "ollama":
+        return os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
+    return None
+
+
+def _apply_env_overrides_to_profile(
+    name: str,
+    p: LLMProfile,
+    *,
+    is_default: bool,
+    secrets: Secrets,
+) -> LLMProfile:
+    """
+    Mutate + return profile with env-based overrides.
+    - For the default profile, allow generic LLM_* env vars.
+    - For all profiles, fill missing base_url / api_key from provider-specific env.
+    """
+    # 1) Generic overrides for DEFAULT profile (if user wants a quick global switch)
+    if is_default:
+        provider_env = os.getenv("LLM_PROVIDER")
+        model_env = os.getenv("LLM_MODEL")
+        base_env = os.getenv("LLM_BASE_URL")
+        timeout_env = os.getenv("LLM_TIMEOUT")
+
+        if provider_env:
+            p.provider = provider_env.lower()  # type: ignore[assignment]
+        if model_env:
+            p.model = model_env
+        if base_env:
+            p.base_url = base_env
+        if timeout_env:
+            try:
+                p.timeout = float(timeout_env)
+            except ValueError:
+                logger = logging.getLogger("aethergraph.services.llm")
+                logger.warning(f"Invalid LLM_TIMEOUT value: {timeout_env}")
+
+    # 2) Provider-specific base_url fallback
+    if not p.base_url:
+        p.base_url = _provider_default_base_url(p.provider)  # type: ignore[arg-type]
+
+    # 3) API key resolution:
+    #    - prefer explicit api_key on profile
+    #    - else api_key_ref + Secrets
+    #    - else provider-specific env name
+    api_key = _resolve_key(p.api_key, p.api_key_ref, secrets)
+
+    if not api_key:
+        # Fallback to provider-specific env if nothing else was set
+        if p.provider == "openai":
+            api_key = os.getenv("OPENAI_API_KEY")
+        elif p.provider == "anthropic":
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+        elif p.provider == "google":
+            api_key = os.getenv("GOOGLE_API_KEY")
+        elif p.provider == "openrouter":
+            api_key = os.getenv("OPENROUTER_API_KEY")
+        elif p.provider == "azure":
+            api_key = os.getenv("AZURE_OPENAI_KEY")
+
+        # If we found one, and no api_key_ref was configured, we can
+        # optionally set api_key_ref so it's visible in config
+        if api_key and not p.api_key_ref:
+            # Optional: record that this profile is using that env key
+            p.api_key_ref = {
+                "openai": "OPENAI_API_KEY",
+                "anthropic": "ANTHROPIC_API_KEY",
+                "google": "GOOGLE_API_KEY",
+                "openrouter": "OPENROUTER_API_KEY",
+                "azure": "AZURE_OPENAI_KEY",
+            }.get(p.provider, None)  # type: ignore[index]
+
+    # Finally, store the resolved key back into api_key for the client factory
+    if api_key:
+        p.api_key = SecretStr(api_key)
+
+    return p
+
+
+def client_from_profile(p: LLMProfile, secrets: Secrets) -> GenericLLMClient:
+    # At this point, _apply_env_overrides_to_profile has already filled
+    # p.base_url, p.api_key, etc. as much as possible.
+    api_key = _resolve_key(p.api_key, p.api_key_ref, secrets)
+
+    return GenericLLMClient(
+        provider=p.provider,
+        model=p.model,
+        embed_model=p.embed_model,
+        base_url=p.base_url,
+        api_key=api_key,
+        azure_deployment=p.azure_deployment,
+        timeout=p.timeout,
+    )
+
+
+def build_llm_clients(cfg: LLMSettings, secrets: Secrets) -> dict[str, GenericLLMClient]:
+    """Returns dict of {profile_name: client}, always includes 'default' if enabled."""
+    if not cfg.enabled:
+        return {}
+
+    # Mutate cfg.llm.default in-place with env defaults
+    default_profile = _apply_env_overrides_to_profile(
+        name="default",
+        p=cfg.default,
+        is_default=True,
+        secrets=secrets,
+    )
+    clients: dict[str, GenericLLMClient] = {
+        "default": client_from_profile(default_profile, secrets)
+    }
+
+    # Extra profiles
+    for name, prof in (cfg.profiles or {}).items():
+        prof = _apply_env_overrides_to_profile(
+            name=name,
+            p=prof,
+            is_default=False,
+            secrets=secrets,
+        )
+        clients[name] = client_from_profile(prof, secrets)
+
+    return clients