aethergraph 0.1.0a3__py3-none-any.whl → 0.1.0a4__py3-none-any.whl

aethergraph/services/memory/distillers/llm_long_term_v1.py

@@ -0,0 +1,180 @@
+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
+from aethergraph.contracts.storage.doc_store import DocStore
+
+# metering
+from aethergraph.services.memory.facade.utils import now_iso
+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,
+        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 "")
+
+        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),
+            "preview": preview,
+            "ts": ts,
+        }

aethergraph/services/memory/distillers/llm_meta_summary.py

@@ -1,70 +1,22 @@
 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.services.memory import Distiller, HotLog
 from aethergraph.contracts.storage.doc_store import DocStore
-from aethergraph.core.runtime.runtime_metering import current_meter_context, current_metering
-from aethergraph.services.memory.distillers.long_term import ar_summary_uri
-from aethergraph.services.memory.facade.utils import now_iso, stable_event_id
+from aethergraph.services.memory.facade.utils import now_iso
 from aethergraph.services.memory.utils import _summary_doc_id, _summary_prefix
 
-"""
-Meta-summary pipeline (multi-scale memory):
-
-  1) Raw events (chat_user / chat_assistant) are recorded via `mem.record(...)`.
-  2) `mem.distill_long_term(...)` compresses recent events into JSON summaries under:
-       mem/<scope_id>/summaries/<summary_tag>/...
-     e.g. summary_tag="session" → session-level long-term summaries.
-  3) `mem.distill_meta_summary(...)` loads those saved summaries from disk and asks the LLM
-     to produce a higher-level "summary of summaries" (meta summary), written under:
-       mem/<scope_id>/summaries/<meta_tag>/...
-
-ASCII view:
-
-    [events in HotLog + Persistence]
-              │
-              ▼
-     distill_long_term(...)
-              │
-              ▼
-    file://mem/<scope>/summaries/session/*.json   (long_term_summary)
-              │
-              ▼
-     distill_meta_summary(...)
-              │
-              ▼
-    file://mem/<scope>/summaries/meta/*.json      (meta_summary: summary of summaries)
-
-You control time scales via `summary_tag` (e.g. "session", "weekly", "meta") and
-`scope_id` (e.g. user+persona).
-"""
-
 
 class LLMMetaSummaryDistiller(Distiller):
-    """
-    LLM-based "summary of summaries" distiller.
-
-    Intended use:
-      - Input: previously generated summary Events (e.g. kind="long_term_summary").
-      - Output: higher-level meta summary (e.g. kind="meta_summary") for a broader time scale.
-
-    Example:
-      - Source: summary_tag="session"  (daily/session summaries)
-      - Target: summary_tag="meta"     (multi-session / weekly/monthly view)
-    """
-
     def __init__(
         self,
         *,
         llm: LLMClientProtocol,
-        # Source summaries (what we are compressing)
         source_kind: str = "long_term_summary",
         source_tag: str = "session",
-        # Target summary (what we produce)
         summary_kind: str = "meta_summary",
         summary_tag: str = "meta",
         max_summaries: int = 20,
@@ -78,115 +30,33 @@ class LLMMetaSummaryDistiller(Distiller):
         self.summary_tag = summary_tag
         self.max_summaries = max_summaries
         self.min_signal = min_signal
-        self.model = model  # optional model override
-
-    def _filter_source_summaries(self, events: Iterable[Event]) -> list[Event]:
-        """
-        Keep only summary Events matching:
-          - kind == source_kind
-          - tags include source_tag (and ideally 'summary')
-          - signal >= min_signal
-        """
-        out: list[Event] = []
-        for e in events:
-            if e.kind != self.source_kind:
-                continue
-            if (e.signal or 0.0) < self.min_signal:
-                continue
-            tags = set(e.tags or [])
-            if self.source_tag and self.source_tag not in tags:
-                continue
-            # Optional, but helps avoid mixing random summaries:
-            # require generic "summary" tag if present in your existing pipeline.
-            # if "summary" not in tags:
-            #     continue
-            out.append(e)
-        return out
-
-    def _build_prompt(self, summaries: list[Event]) -> list[dict[str, str]]:
-        """
-        Convert summary Events into a chat prompt for the LLM.
-
-        We use:
-          - e.text as the main human-readable summary preview.
-          - e.data.get("time_window") if present.
-        """
-
-        lines: list[str] = []
-
-        for idx, e in enumerate(summaries, start=1):
-            tw = (e.data or {}).get("time_window") if e.data else None
-            tw_from = (tw or {}).get("from", e.ts)
-            tw_to = (tw or {}).get("to", e.ts)
-            body = e.text or ""
-            lines.append(f"Summary {idx} [{tw_from} → {tw_to}]:\n{body}\n")
-
-        transcript = "\n\n".join(lines)
-
-        system = (
-            "You are a higher-level summarizer over an agent's existing summaries. "
-            "Given multiple prior summaries (each covering a period of time), you "
-            "should produce a concise, higher-level meta-summary capturing: "
-            "  - long-term themes and patterns, "
-            "  - important user facts that remain true, "
-            "  - long-running goals or open loops."
-        )
-
-        user = (
-            "Here are several previous summaries, each describing a time window:"
-            "\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 outside the JSON."
-        )
-
-        return [
-            {"role": "system", "content": system},
-            {"role": "user", "content": user},
-        ]
+        self.model = model
 
     def _build_prompt_from_saved(self, summaries: list[dict[str, Any]]) -> list[dict[str, str]]:
-        """
-        Build an LLM prompt from persisted summary JSONs.
-
-        Each summary dict is the JSON you showed:
-          {
-            "type": "long_term_summary",
-            "summary_tag": "session",
-            "summary": "...",
-            "time_window": {...},
-            ...
-          }
-        """
         lines: list[str] = []
-
         for idx, s in enumerate(summaries, start=1):
             tw = s.get("time_window") or {}
-            tw_from = tw.get("from", s.get("ts"))
-            tw_to = tw.get("to", s.get("ts"))
-            body = s.get("summary", "") or ""
+            tw_from = tw.get("from") or s.get("ts")
+            tw_to = tw.get("to") or s.get("ts")
 
-            # (Optional) strip ```json fences if present
-            stripped = body.strip()
-            if stripped.startswith("```"):
-                # very minimal fence strip; you can refine later
-                stripped = stripped.strip("`")
-                # fall back to original if this gets too messy
-                body_for_prompt = stripped or body
-            else:
-                body_for_prompt = body
+            # Support both "summary" (LLM distiller) and "text" (non-LLM distiller)
+            body = (s.get("summary") or s.get("text") or "").strip()
+
+            # Minimal fence stripping if someone stored fenced content
+            if body.startswith("```"):
+                body = body.strip().strip("`").strip()
 
-            lines.append(f"Summary {idx} [{tw_from} → {tw_to}]:\n{body_for_prompt}\n")
+            if len(body) > 2000:
+                body = body[:2000] + "…"
+
+            lines.append(f"Summary {idx} [{tw_from} → {tw_to}]:\n{body}\n")
 
         transcript = "\n\n".join(lines)
 
         system = (
             "You are a higher-level summarizer over an agent's existing long-term summaries. "
-            "Given multiple prior summaries (each describing a period), produce a meta-summary "
-            "that captures long-term themes, stable user facts, and persistent open loops."
+            "Given multiple prior summaries (each describing a time window), produce a meta-summary "
+            "capturing long-term themes, stable user facts, and persistent open loops."
         )
 
         user = (
@@ -199,36 +69,22 @@ class LLMMetaSummaryDistiller(Distiller):
             "Do not include any extra explanation outside the JSON."
         )
 
-        return [
-            {"role": "system", "content": system},
-            {"role": "user", "content": user},
-        ]
+        return [{"role": "system", "content": system}, {"role": "user", "content": user}]
 
     async def distill(
         self,
         run_id: str,
         timeline_id: str,
-        scope_id: str = None,
+        scope_id: str | None = None,
         *,
         hotlog: HotLog,
-        persistence: Persistence,
-        indices: Indices,
         docs: DocStore,
         **kw: Any,
     ) -> dict[str, Any]:
-        """
-        Distill method following the Distiller protocol.
-
-        IMPORTANT:
-          - This implementation is optimized for FSPersistence and reads
-            previously saved summary JSONs from:
-              mem/<scope_id>/summaries/<source_tag>/*.json
-          - If a different Persistence is used, we currently bail out.
-        """
         scope = scope_id or run_id
         prefix = _summary_prefix(scope, self.source_tag)
 
-        # 1) Load existing long-term summary JSONs from DocStore
+        # Load persisted long-term summaries from DocStore
         try:
             all_ids = await docs.list()
         except Exception:
@@ -239,68 +95,59 @@ class LLMMetaSummaryDistiller(Distiller):
             return {}
 
         chosen_ids = candidates[-self.max_summaries :]
-        summaries: list[dict[str, Any]] = []
+        loaded: list[dict[str, Any]] = []
         for doc_id in chosen_ids:
             try:
                 doc = await docs.get(doc_id)
                 if doc is not None:
-                    summaries.append(doc)  # type: ignore[arg-type]
+                    loaded.append(doc)  # type: ignore[arg-type]
             except Exception:
                 continue
 
-        if not summaries:
+        if not loaded:
             return {}
 
-        # Optional: filter by min_signal if present in saved JSON
-        filtered: list[dict[str, Any]] = []
-        for s in summaries:
-            sig = (
-                float(s.get("signal", 0.0)) if isinstance(s.get("signal"), int | float) else 1.0
-            )  # default 1.0
-            if sig < self.min_signal:
-                continue
-            # Also enforce type/tag consistency:
+        # Enforce consistency + min_signal if present
+        kept: list[dict[str, Any]] = []
+        for s in loaded:
             if s.get("type") != self.source_kind:
                 continue
             if s.get("summary_tag") != self.source_tag:
                 continue
-            filtered.append(s)
 
-        if not filtered:
-            return {}
+            sig_val = s.get("signal", None)
+            if isinstance(sig_val, (int, float)) and float(sig_val) < self.min_signal:  # noqa: UP038
+                continue
+            kept.append(s)
 
-        # Keep order as loaded (already sorted by filename)
-        kept = filtered
+        if not kept:
+            return {}
 
-        # 2) Derive aggregated time window
-        first_from = None
-        last_to = None
-        for s in kept:
+        # Derive aggregated time window safely
+        def _pick_time(s: dict[str, Any], key: str) -> str | None:
             tw = s.get("time_window") or {}
-            start = tw.get("from") or s.get("ts")
-            end = tw.get("to") or s.get("ts")
-            if start:
-                first_from = start if first_from is None else min(first_from, start)
-            if end:
-                last_to = end if last_to is None else max(last_to, end)
-        if first_from is None:
-            first_from = kept[0].get("ts")
-        if last_to is None:
-            last_to = kept[-1].get("ts")
+            return tw.get(key) or s.get("ts")
+
+        times_from = [t for t in (_pick_time(s, "from") for s in kept) if t]
+        times_to = [t for t in (_pick_time(s, "to") for s in kept) if t]
+
+        first_from = min(times_from) if times_from else (kept[0].get("ts") or now_iso())
+        last_to = max(times_to) if times_to else (kept[-1].get("ts") or now_iso())
 
-        # 3) Build prompt and call LLM
+        # Build prompt and call LLM (respect model override)
         messages = self._build_prompt_from_saved(kept)
-        summary_json_str, usage = await self.llm.chat(messages)
+        try:
+            if self.model:
+                summary_json_str, usage = await self.llm.chat(messages, model=self.model)  # type: ignore[arg-type]
+            else:
+                summary_json_str, usage = await self.llm.chat(messages)
+        except TypeError:
+            summary_json_str, usage = await self.llm.chat(messages)
 
-        # 4) Parse LLM JSON response
         try:
             payload = json.loads(summary_json_str)
         except Exception:
-            payload = {
-                "summary": summary_json_str,
-                "key_facts": [],
-                "open_loops": [],
-            }
+            payload = {"summary": summary_json_str, "key_facts": [], "open_loops": []}
 
         ts = now_iso()
         summary_obj = {
@@ -314,85 +161,29 @@ class LLMMetaSummaryDistiller(Distiller):
             "ts": ts,
             "time_window": {"from": first_from, "to": last_to},
             "num_source_summaries": len(kept),
-            "source_summary_uris": [
-                # reconstruct the URI pattern we originally use
-                # (this assumes summaries were written under ar_summary_uri)
-                ar_summary_uri(scope, self.source_tag, s.get("ts", ts))
-                for s in kept
-            ],
+            # ✅ store doc_ids you actually read (truth)
+            "source_summary_doc_ids": chosen_ids[-len(kept) :],
             "summary": payload.get("summary", ""),
             "key_facts": payload.get("key_facts", []),
             "open_loops": payload.get("open_loops", []),
             "llm_usage": usage,
             "llm_model": getattr(self.llm, "model", None),
+            "llm_model_override": self.model,
+            "min_signal": self.min_signal,
         }
 
         doc_id = _summary_doc_id(scope, self.summary_tag, ts)
         await docs.put(doc_id, summary_obj)
 
-        # 5) Emit meta_summary Event
         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_meta",
-            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_source_summaries": len(kept),
-                "source_summary_kind": self.source_kind,
-                "source_summary_tag": self.source_tag,
-            },
-            metrics={"num_source_summaries": len(kept)},
-            severity=2,
-            signal=0.8,
-        )
-
-        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_meta_summary")
-            logger.error("Failed to record metering event for llm_meta_summary")
-
         return {
             "summary_doc_id": doc_id,
             "summary_kind": self.summary_kind,
             "summary_tag": self.summary_tag,
             "time_window": summary_obj["time_window"],
-            "num_source_summaries": len(kept),
+            "num_source_summaries": summary_obj["num_source_summaries"],
+            "preview": preview,
+            "ts": ts,
         }