diff --git a/v2/tutorials/cognee_memory_store/README.md b/v2/tutorials/cognee_memory_store/README.md new file mode 100644 index 00000000..96cfbc19 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/README.md @@ -0,0 +1,184 @@ +# Cognee + Flyte Memory Store + +A full-stack tutorial demonstrating a **sleep/wake memory architecture** for AI agents using [Cognee](https://github.com/topoteretes/cognee) (knowledge graph) and [Flyte v2](https://docs.union.ai/v2/). + +Inspired by Claude's Managed Agents memory store concepts: +- Many small focused memory files, addressed by path +- Access control via directory prefix (`memory/` is read-only) +- Immutable version history + append-only audit log +- Staged proposals: untrusted writes land in `staging/` before promotion +- Optimistic concurrency (expected SHA-256 preconditions) + +## Architecture + +``` +User question + │ + │ Wake cycle (per question — Flyte task) + │ ├─ Download shared state + │ ├─ Route question → topic slugs (Claude classifier) + │ ├─ cognee.search(datasets=[slugs]) ← semantic retrieval + │ ├─ Raw-file fallback (memory//*.txt) if graph incomplete + │ ├─ Inject user/sessions//memories/ context + │ └─ Claude answer with [preferences] + [retrieved] + │ + ▼ +Chat answer ──► Claude proposes a memory to stage ──► Accept/Edit/Deny + │ + staging/sessions//inbox/ + │ + Sleep cycle (every 6h — Flyte task) + ├─ Auto-promote staged proposals + ├─ Cluster + consolidate user memories + │ (flyte.map.aio — parallel pods) + ├─ Rebuild Cognee graph per topic + │ (flyte.map.aio — parallel pods) + └─ Upload updated state +``` + +**Flyte features demonstrated:** +| Feature | Where used | +|---|---| +| `flyte.map.aio` | Parallel cluster consolidation + per-topic Cognee rebuild | +| `flyte.io.Dir` | Sync memstore + Cognee DBs across pods via object storage | +| `cache="auto"` | `consolidate_cluster` skips unchanged clusters on retry | +| `report=True` | Live HTML progress streamed to Union UI during sleep cycle | +| `flyte.group()` | Per-phase spans visible in the execution timeline | +| `retries=N` | Transient network / Cognee failures auto-retried | +| `flyte.app.AppEnvironment` | Streamlit app served on Union | + +## Storage layout + +``` +memory/ ← READ-ONLY (written by ingest_url only) + _index.json ← topic index {slug: {label, sources}} + / + .txt ← scraped + stripped page content + +user/ + preferences.json ← shared across sessions + preferences.txt ← derived plain-text copy + sessions/ + _registry.json ← {session_name: {created_at_s, label}} + / + memories/ + _topic_map.json ← {memory_path: topic_slug} + .txt ← promoted user memories + chat/ + / + transcript.jsonl + summary.txt + +staging/ + sessions/ + / + inbox/.json ← untrusted staged proposals + archive// ← archived proposals with audit note + +audit/log.jsonl ← append-only mutation log +meta/ ← sha256 + timestamp per memory file +versions/ ← immutable snapshots of every write +``` + +## Files + +| File | Purpose | +|---|---| +| `memory_store.py` | Audited, versioned file-based memory store; access control, concurrency, versioning | +| `agent.py` | Proposal schema, staging, validation (anti-injection), promotion | +| `workflow.py` | All Flyte tasks: `init_memory_store`, `ingest_url`, `consolidate_cluster`, `rebuild_topic_dataset`, `sleep_cycle`, `wake_cycle`, `summarize_chat_session` | +| `app.py` | Streamlit UI: chat, session selector, URL ingestion, sleep cycle trigger, memory viewer | + +## Prerequisites + +1. **Union account** — [sign up at union.ai](https://union.ai) +2. **Anthropic API key** stored as a Union secret: + ```bash + union create secret internal-anthropic-api-key + # paste your key when prompted + ``` +3. **`uv`** installed — [docs.astral.sh/uv](https://docs.astral.sh/uv) + +## Quickstart (Union) + +```bash +# Clone and enter the repo +git clone https://github.com/unionai/unionai-examples +cd unionai-examples + +# Authenticate with Union +union login + +# Deploy and serve (builds image, registers sleep schedule, serves app) +uv run v2/tutorials/cognee_memory_store/app.py +``` + +Open the printed app URL. + +### Optional: configure image registry + +By default, images are pushed to `ghcr.io/flyteorg`. Override with: + +```bash +export AI_MEMORY_STORE_IMAGE_REGISTRY="ghcr.io/" +``` + +## Using the app + +### 1. Initialize +The app seeds a fresh memory store on first launch via the `init_memory_store` Flyte task. + +### 2. Ingest reference knowledge +In the sidebar under **Seed Knowledge from URL**: +- Enter a URL (e.g. `https://docs.union.ai/v2/union/user-guide/`) +- Set **Max pages** (up to 50) +- Click **Ingest URL** + +The `ingest_url` task crawls the site, classifies content into a topic, writes pages to `memory//`, and builds a Cognee knowledge graph. Check the **Topic knowledge base** expander to see ingested files. + +### 3. Ask questions +Type in the chat input. The app retrieves context from Cognee + raw memory files, then calls Claude. After each answer, Claude suggests a memory to stage — accept, edit, or deny inline. + +### 4. Manage sessions +Use the **Session** selector in the sidebar to create isolated sessions. Each session has its own promoted memories, staging inbox, and chat history. Reference knowledge (`memory/`) is shared across all sessions. + +### 5. Sleep cycle +Click **Trigger sleep now** (or wait — it fires automatically every 6 hours): +- Staged proposals are auto-promoted to `user/sessions//memories/` +- Related memories are clustered and consolidated by Claude +- Cognee knowledge graphs are rebuilt per topic +- Live HTML progress streams to the Union UI report panel + +### 6. Preferences +Under **Preferences** in the sidebar, set tone, format, and your name. Claude will follow these in every answer. You can also say things like *"always answer in bullet points"* in chat — Claude detects the preference and offers an inline approval card. + +## Debug mode + +```bash +export AI_MEMORY_STORE_DEBUG=1 +``` + +Enables per-message retrieval/answer timing and proposal detection details in the UI. + +## Local dev (no Union required) + +```bash +streamlit run v2/tutorials/cognee_memory_store/app.py -- --server +``` + +Staging, promotion, and memory viewer work locally. The `wake_cycle` and `sleep_cycle` Flyte tasks require remote object storage and will not run in pure local mode. + +Run the built-in self-checks: + +```bash +python v2/tutorials/cognee_memory_store/memory_store.py # storage self-check +SELF_CHECK=true python v2/tutorials/cognee_memory_store/app.py # app self-check +``` + +## What to look for + +- **Audit trail** — expand "Audit log (tail)" to see every stage/promote/consolidate event with timestamps and actors +- **Versioning** — every promoted write creates an immutable snapshot under `versions/` +- **Staged proposals** — the "Staging inbox" expander shows pending proposals before they're promoted +- **Parallel consolidation** — watch the sleep cycle's `flyte.map.aio` fan-out in the Union UI execution timeline +- **Raw-file fallback** — when Cognee's graph is incomplete (entity extraction can hit LLM output limits), the app falls back to reading `memory//*.txt` directly diff --git a/v2/tutorials/cognee_memory_store/agent.py b/v2/tutorials/cognee_memory_store/agent.py new file mode 100644 index 00000000..43e0feaa --- /dev/null +++ b/v2/tutorials/cognee_memory_store/agent.py @@ -0,0 +1,281 @@ +"""Agent helpers for the Cognee + Flyte memory-store tutorial. + +This module defines: +- A staged proposal format (untrusted writes land in staging/ first) +- A validator that rejects obvious memory-poisoning attempts +- Promotion helpers that write into user/ via MemoryStore + +The intention is to mimic Claude memory stores best practices: +- Separate untrusted staging from trusted memory +- Audit everything +- Use access modes (reference is read-only) +""" + +from __future__ import annotations + +import re +import time +import uuid +from typing import Literal, Optional + +from pydantic import BaseModel, Field + +from memory_store import ( + AccessDenied, + ConcurrencyError, + MemoryMeta, + MemoryStore, + session_staging_inbox_prefix, + session_staging_archive_prefix, +) + + +ProposalTarget = Literal["user"] +ProposalFormat = Literal["text", "json"] + + +class MemoryWriteProposal(BaseModel): + """An untrusted candidate write. + + - Stored under staging/inbox/.json + - Must be validated before promotion. + """ + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + created_at_s: float = Field(default_factory=lambda: time.time()) + + # Where the proposal intends to write + target: ProposalTarget + path: str + + # The content to store + format: ProposalFormat = "text" + content: str + + # Optional concurrency precondition + expected_sha256: Optional[str] = None + + # Provenance + author: str = "unknown" # user_id, session_id, etc. + reason: str = "" # why this memory is being written + + # Debuggable context (kept small) + source_question: str = "" + source_answer: str = "" + + # Topic classification — set at staging time, used by sleep_cycle to update topic_map + topic_slug: Optional[str] = None + + # Session this proposal belongs to — determines storage namespace + session: str = "default" + + +class ProposalDecision(BaseModel): + ok: bool + reason: str + normalized_path: str = "" + + +_DISALLOWED_PATH_PREFIXES = ( + "audit/", + "meta/", + "versions/", +) + + +def classify_proposal_topic( + content: str, + source_question: str, + topic_index: dict, + api_key: str, +) -> Optional[str]: + """Return the best-matching topic slug from the index, or None if no match. + + Used at proposal staging time so that sleep_cycle can link promoted user/ + memories to the right Cognee topic dataset without filename heuristics. + """ + if not topic_index or not api_key: + return None + + from anthropic import Anthropic + + topic_lines = "\n".join(f" {s}: {e.get('label', s)}" for s, e in list(topic_index.items())[:20]) + prompt = ( + f"Topics:\n{topic_lines}\n\n" + f"Question context: {source_question[:500]}\n\n" + f"Memory content:\n{content[:1000]}\n\n" + "Which topic slug does this memory belong to? " + "Return the exact slug string only, or null if none match well.\n" + "Return JSON only: \"topic_slug_name\" or null" + ) + try: + client = Anthropic(api_key=api_key, timeout=15.0) + msg = client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=40, + temperature=0, + messages=[{"role": "user", "content": prompt}], + ) + raw = msg.content[0].text.strip().strip('"').lower() + return raw if raw in topic_index else None + except Exception: + return None + + +def proposal_inbox_path(proposal_id: str, session: str = "default") -> str: + return f"{session_staging_inbox_prefix(session)}/{proposal_id}.json" + + +def proposal_archive_path(proposal_id: str, decision: str, session: str = "default") -> str: + safe = re.sub(r"[^a-zA-Z0-9_-]", "_", decision)[:24] or "unknown" + return f"{session_staging_archive_prefix(session)}/{safe}/{proposal_id}.json" + + +def stage_proposal(store: MemoryStore, proposal: MemoryWriteProposal) -> MemoryMeta: + """Write a proposal into the session's staging inbox (untrusted).""" + return store.write_json( + proposal_inbox_path(proposal.id, proposal.session), + proposal.model_dump(), + actor=proposal.author, + reason=proposal.reason or "stage-proposal", + op="stage", + ) + + +def list_staged_proposals(store: MemoryStore, limit: int = 50, session: str = "default") -> list[MemoryWriteProposal]: + paths = store.list_paths(session_staging_inbox_prefix(session)) + # Backward compat: also scan old-style staging/inbox for the default session + if session == "default": + for p in store.list_paths("staging/inbox"): + if p not in paths: + paths.append(p) + paths = sorted(paths)[:limit] + out: list[MemoryWriteProposal] = [] + for p in paths: + try: + raw = store.read_json(p, default=None) + if not raw: + continue + out.append(MemoryWriteProposal(**raw)) + except Exception: + continue + # newest first + out.sort(key=lambda x: x.created_at_s, reverse=True) + return out + + +def _normalize_target_path(proposal: MemoryWriteProposal) -> str: + session = proposal.session or "default" + memories_prefix = f"user/sessions/{session}/memories/" + path = proposal.path.lstrip("/") + if path.startswith(memories_prefix): + return path + # Strip any old-style "user/" prefix before routing to session namespace + path = path.removeprefix("user/") + return memories_prefix + path + + +def validate_proposal( + store: MemoryStore, + proposal: MemoryWriteProposal, +) -> ProposalDecision: + """Cheap, deterministic validator. + + This is intentionally conservative. If it rejects too often, tune it. + """ + + normalized = _normalize_target_path(proposal) + + if any(normalized.startswith(p) for p in _DISALLOWED_PATH_PREFIXES): + return ProposalDecision(ok=False, reason="attempt to write internal paths") + + if normalized.startswith("memory/"): + return ProposalDecision(ok=False, reason="memory/ is machine-managed, not user-writable") + + if proposal.target != "user": + return ProposalDecision(ok=False, reason="invalid target") + + # Basic size bounds: keep memories small and focused. + if len(proposal.content.encode("utf-8")) > 25_000: + return ProposalDecision(ok=False, reason="memory too large; split into smaller files") + + # Poisoning / injection heuristics: don't store instructions that would hijack later prompts. + lc = proposal.content.lower() + suspicious_markers = [ + "ignore previous", + "system prompt", + "developer message", + "you must obey", + "exfiltrate", + "api key", + "password", + "ssh-key", + ] + if any(m in lc for m in suspicious_markers): + return ProposalDecision(ok=False, reason="content looks like prompt injection / secret material") + + return ProposalDecision(ok=True, reason="ok", normalized_path=normalized) + + +def promote_proposal( + store: MemoryStore, + proposal: MemoryWriteProposal, + *, + actor: str = "promoter", + promotion_reason: str = "", +) -> MemoryMeta: + """Promote a validated proposal into trusted memory. + + IMPORTANT: This mutates trusted memory. Caller should run validate_proposal first. + """ + + decision = validate_proposal(store, proposal) + if not decision.ok: + raise AccessDenied(f"Proposal {proposal.id} rejected: {decision.reason}") + + path = decision.normalized_path + + try: + meta = store.write_text( + path, + proposal.content, + actor=actor, + reason=promotion_reason or proposal.reason or "promote", + expected_sha=proposal.expected_sha256, + op="promote", + extra_audit={"proposal_id": proposal.id, "proposal_author": proposal.author}, + ) + except ConcurrencyError: + raise + + return meta + + +def archive_proposal( + store: MemoryStore, + proposal: MemoryWriteProposal, + *, + actor: str, + decision: str, + note: str, +) -> None: + """Archive the staged proposal with an audit event. + + This does NOT delete the inbox entry (keeps the tutorial simple and append-only). + """ + store.ensure_layout() + archive_path = proposal_archive_path(proposal.id, decision, proposal.session) + store.write_json( + archive_path, + { + **proposal.model_dump(), + "archived_at_s": time.time(), + "archived_by": actor, + "decision": decision, + "note": note, + }, + actor=actor, + reason=f"archive:{decision}", + op="archive", + extra_audit={"proposal_id": proposal.id, "decision": decision}, + ) diff --git a/v2/tutorials/cognee_memory_store/app.py b/v2/tutorials/cognee_memory_store/app.py new file mode 100644 index 00000000..047927cb --- /dev/null +++ b/v2/tutorials/cognee_memory_store/app.py @@ -0,0 +1,1508 @@ +# /// script +# requires-python = "==3.13" +# dependencies = [ +# "flyte==2.1.5", +# "cognee==1.0.7", # see workflow.py for version rationale +# "streamlit>=1.42.0", +# "pydantic>=2.11.0", +# "anthropic>=0.40.0", +# "fastembed>=0.3.0", +# "packaging>=23.0", +# "charset-normalizer>=3.0", +# ] +# main = "main" +# /// + +"""Cognee + Flyte memory stores — Streamlit app. + +Run modes: + Local dev: streamlit run app.py + Serve on Union: uv run app.py + +Memory flow: + - Chat answered using Cognee semantic retrieval over promoted memories + - After each answer, Claude proposes a memory to stage (inline card: Accept/Edit/Deny) + - Accepted proposals → staging/inbox/ → auto-promoted on next sleep cycle + - Sleep cycle status visible in sidebar; manual trigger available +""" + +from __future__ import annotations + +import asyncio +import json +import re +import os +import sys +import time +import threading +import uuid +import shutil +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError, field_validator, model_validator + +import flyte +import flyte.app +from flyte.io import Dir + +from agent import MemoryWriteProposal, classify_proposal_topic, list_staged_proposals, stage_proposal +from workflow import ( + LOCAL_COGNEE_ROOT, + LOCAL_MEMSTORE_ROOT, + SHARED_COGNEE_DB_PREFIX, + SHARED_MEMSTORE_PATH, + GENERAL_TOPIC_SLUG, + _configure_cognee_runtime, + _setup_cognee_env, + _route_query_to_topics, + _strip_jina_header, + _topic_db_path, + init_memory_store, + ingest_url, + sleep_cycle, + summarize_chat_session, +) +from memory_store import ( + MemoryStore, + SESSION_REGISTRY_PATH, + TOPIC_INDEX_PATH, + load_topic_index, + list_sessions, + register_session, + session_memories_prefix, + session_staging_inbox_prefix, + session_staging_archive_prefix, + upsert_topic_index, + _parse_json_object, + _parse_json_array, +) + +MODEL = os.environ.get("AI_MEMORY_STORE_MODEL", "claude-haiku-4-5-20251001") +DEBUG = os.environ.get("AI_MEMORY_STORE_DEBUG", "").lower() in ("1", "true", "yes") + +PREFS_JSON_PATH = "user/preferences.json" +PREFS_TXT_PATH = "user/preferences.txt" + +CHAT_CONTEXT_MESSAGES = int(os.environ.get("AI_CHAT_CONTEXT_MESSAGES", "24")) +CHAT_TRANSCRIPT_MAX_LINES = int(os.environ.get("AI_CHAT_TRANSCRIPT_MAX_LINES", "1000")) + + +Scalar = str | int | float | bool + + +class ExtractedProposal(BaseModel): + model_config = ConfigDict(extra="forbid") + + type: Literal["preference", "memory"] + reason: str + + # preference + updates: dict[str, Scalar] | None = None + + # memory + path: str | None = None + content: str | None = None + + @field_validator("reason") + @classmethod + def _reason_nonempty(cls, v: str) -> str: + v = (v or "").strip() + if not v: + raise ValueError("reason is required") + if len(v) > 240: + raise ValueError("reason too long") + return v + + @field_validator("updates") + @classmethod + def _validate_updates(cls, v: dict[str, Any] | None) -> dict[str, Scalar] | None: + if v is None: + return None + if not isinstance(v, dict) or not v: + raise ValueError("updates must be a non-empty object") + if len(v) > 12: + raise ValueError("too many preference updates") + + cleaned: dict[str, Scalar] = {} + for k, val in v.items(): + if not isinstance(k, str) or not re.fullmatch(r"[a-zA-Z][a-zA-Z0-9_\-]{0,40}", k): + raise ValueError(f"invalid preference key: {k!r}") + + if isinstance(val, bool): + cleaned[k] = val + elif isinstance(val, (int, float)) and not isinstance(val, bool): + cleaned[k] = val + elif isinstance(val, str): + sval = val.strip() + if len(sval) > 400: + raise ValueError(f"preference value too long for {k!r}") + cleaned[k] = sval + else: + raise ValueError(f"invalid preference value for {k!r} (must be scalar)") + + return cleaned + + @model_validator(mode="after") + def _validate_shape(self) -> "ExtractedProposal": + if self.type == "preference": + if not self.updates: + raise ValueError("preference requires updates") + if self.path is not None or self.content is not None: + raise ValueError("preference must not include path/content") + elif self.type == "memory": + if not self.path or not isinstance(self.path, str) or not self.path.strip(): + raise ValueError("memory requires path") + if not self.content or not isinstance(self.content, str) or not self.content.strip(): + raise ValueError("memory requires content") + if self.updates is not None: + raise ValueError("memory must not include updates") + return self + +THIS_DIR = Path(__file__).resolve().parent +_file_name = Path(__file__).name + +app_env = flyte.app.AppEnvironment( + name="cognee-memory-store-chat", + image=( + flyte.Image.from_uv_script(__file__, name="cognee-memory-store-chat", pre=True) + .with_source_file(THIS_DIR / "memory_store.py", "/root") + .with_source_file(THIS_DIR / "agent.py", "/root") + .with_source_file(THIS_DIR / "workflow.py", "/root") + ), + args=["streamlit", "run", _file_name, "--server.port", "8080", "--", "--server"], + port=8080, + scaling=flyte.app.Scaling(replicas=(1, 1)), + resources=flyte.Resources(cpu=2, memory="4Gi"), + secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")], +) + + +# --------------------------------------------------------------------------- +# Flyte connection +# --------------------------------------------------------------------------- + +# --------------------------------------------------------------------------- +# Shared state sync +# --------------------------------------------------------------------------- + +async def _reset_shared_state() -> None: + """Wipe the shared memstore and every known per-topic cognee_db on remote storage. + + Old runs can leave behind cognee_db files keyed to an older library/schema + version (e.g. corrupted Kuzu/Ladybug state), which then crashes the next + ingest with "Could not map version_code". To avoid that, we overwrite each + remote prefix with an empty directory at app startup. + + We read the topic index BEFORE wiping memstore so we know which topic DBs + exist. For a brand-new install (no remote memstore yet) this is a no-op for + cognee_db and the memstore wipe is also a no-op upload. + """ + import tempfile + + # 1) Discover existing topic slugs from the remote memstore (best-effort). + slugs: list[str] = [] + try: + with tempfile.TemporaryDirectory() as snapshot: + await Dir(path=SHARED_MEMSTORE_PATH).download(local_path=snapshot) + slugs = list(load_topic_index(MemoryStore(Path(snapshot))).keys()) + except Exception: + pass # Remote memstore doesn't exist yet — nothing to enumerate. + + # 2) Overwrite memstore + each known topic DB + general catchall with empty dirs. + targets = [SHARED_MEMSTORE_PATH] + targets.extend(_topic_db_path(s) for s in slugs) + targets.append(_topic_db_path(GENERAL_TOPIC_SLUG)) + + with tempfile.TemporaryDirectory() as empty: + for path in targets: + try: + await Dir.from_local(empty, remote_destination=path) + except Exception: + pass # Best-effort wipe; don't block app startup. + + # 3) Wipe local caches so the next download starts clean. + if LOCAL_MEMSTORE_ROOT.exists(): + shutil.rmtree(LOCAL_MEMSTORE_ROOT, ignore_errors=True) + if LOCAL_COGNEE_ROOT.exists(): + shutil.rmtree(LOCAL_COGNEE_ROOT, ignore_errors=True) + + +async def _download_shared_state() -> None: + LOCAL_MEMSTORE_ROOT.mkdir(parents=True, exist_ok=True) + LOCAL_COGNEE_ROOT.mkdir(parents=True, exist_ok=True) + await Dir(path=SHARED_MEMSTORE_PATH).download(local_path=str(LOCAL_MEMSTORE_ROOT)) + # Download per-topic cognee DBs based on the current topic index + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + topic_index = load_topic_index(store) + for slug in topic_index: + local_cognee = LOCAL_COGNEE_ROOT / slug + local_cognee.mkdir(parents=True, exist_ok=True) + try: + await Dir(path=_topic_db_path(slug)).download(local_path=str(local_cognee)) + except Exception: + pass # Topic DB may not exist yet if not yet ingested + # Always download the general catchall DB (user memories not tied to any ingested topic) + general_cognee = LOCAL_COGNEE_ROOT / GENERAL_TOPIC_SLUG + general_cognee.mkdir(parents=True, exist_ok=True) + try: + await Dir(path=_topic_db_path(GENERAL_TOPIC_SLUG)).download(local_path=str(general_cognee)) + except Exception: + pass # General DB may not exist yet if no unclassified memories have been promoted + + +async def _upload_memstore() -> None: + await Dir.from_local(str(LOCAL_MEMSTORE_ROOT), remote_destination=SHARED_MEMSTORE_PATH) + + +def _ensure_seeded() -> None: + import concurrent.futures + + def _do_seed(): + # Wipe the shared memstore + cognee_db prefixes on every app start, then + # re-init from scratch. This prevents stale/corrupted state from older + # cognee versions (e.g. KeyError 'ladybug', kuzu version_code mismatch) + # from breaking subsequent ingests. See _reset_shared_state for details. + try: + asyncio.run(_reset_shared_state()) + except Exception: + pass # Best-effort; the seed step below still re-initialises memstore. + + run = flyte.run(init_memory_store) + run.wait(quiet=True) + run.sync() + asyncio.run(_download_shared_state()) + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + fut = pool.submit(_do_seed) + try: + fut.result(timeout=120) + except (concurrent.futures.TimeoutError, Exception): + # Let the app start even if seeding is still in progress. + pass + + +# --------------------------------------------------------------------------- +# LLM helpers +# --------------------------------------------------------------------------- + +def _call_llm(system: str, messages: list[dict], timeout_s: float = 30.0, max_tokens: int = 900) -> str: + import anthropic + + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + if not api_key: + return "[error] ANTHROPIC_API_KEY not set" + + client = anthropic.Anthropic(api_key=api_key, timeout=timeout_s) + msg = client.messages.create( + model=MODEL, + max_tokens=max_tokens, + system=system, + messages=messages, + temperature=0, + ) + return msg.content[0].text + + +def _extract_proposal_from_message(user_message: str) -> dict | None: + """Parallel Claude call — detects anything in the user's message worth staging. + + Runs concurrently with _call_llm so it adds zero latency to the response. + + Returns one of: + {"type": "preference", "updates": {key: value, ...}, "reason": "..."} + {"type": "memory", "content": "...", "path": "user/.txt", "reason": "..."} + None + """ + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + if not api_key: + return None + + import anthropic + + try: + client = anthropic.Anthropic(api_key=api_key, timeout=25.0) + msg = client.messages.create( + model=MODEL, + max_tokens=300, + system=( + "You detect whether a user message contains something worth persisting as a memory or preference.\n\n" + "Return exactly ONE of the following JSON objects (no extra keys), or exactly null.\n\n" + "Preference schema (dynamic):\n" + " {\"type\":\"preference\",\"updates\":{:,...},\"reason\":}\n" + " - updates must be a small object (1-12 entries)\n" + " - keys must match: [a-zA-Z][a-zA-Z0-9_\\-]{0,40}\n" + " - values must be JSON scalars only: string/number/boolean (no arrays/objects)\n\n" + "Memory schema:\n" + " {\"type\":\"memory\",\"path\":,\"content\":,\"reason\":}\n\n" + "Examples:\n" + " User: 'I’m working on a Flyte v2 workflow for batch inference this week.'\n" + " Return: {\"type\":\"memory\",\"path\":\"user/projects_flyte_batch_inference.txt\",\"content\":\"User is working on a Flyte v2 workflow for batch inference this week.\",\"reason\":\"current project\"}\n\n" + " User: 'Always answer in bullet points.'\n" + " Return: {\"type\":\"preference\",\"updates\":{\"answer_style\":\"bullet_points\"},\"reason\":\"format preference\"}\n\n" + "Skip (return null) for: questions, vague chat, or anything already covered by existing preferences.\n\n" + "Return JSON only, no explanation, no code fences." + ), + messages=[{"role": "user", "content": user_message}], + temperature=0, + ) + + raw = msg.content[0].text + result = _parse_json_object(raw) + + try: + proposal = ExtractedProposal.model_validate(result) if result else None + except ValidationError as e: + proposal = None + err = e + else: + err = None + + # Debug breadcrumbs for the UI. + try: + import streamlit as st + + st.session_state.last_proposal_raw = raw + st.session_state.last_proposal_parsed = result + st.session_state.last_proposal_error = ( + (err.errors()[0].get("msg") if err else "") + if "err" in locals() else "" + ) + except Exception: + pass + + if not proposal: + return None + + out = proposal.model_dump(exclude_none=True) + if DEBUG: + print(f"DEBUG proposal={out}", file=sys.stderr) + return out + except Exception as e: + if DEBUG: + print(f"DEBUG _extract_proposal_from_message error: {type(e).__name__}: {e}", file=sys.stderr) + return None + + +def _retrieve_context(question: str, session: str = "default", timeout_s: float = 15.0) -> str: + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + topic_index = load_topic_index(store) + target_slugs = _route_query_to_topics(question, topic_index, api_key) + if not target_slugs: + target_slugs = list(topic_index.keys()) + + async def _search() -> str: + import cognee + all_results = [] + for slug in target_slugs: + local_cognee = LOCAL_COGNEE_ROOT / slug + if not local_cognee.exists(): + continue + _setup_cognee_env(local_cognee) + _configure_cognee_runtime(cognee, local_cognee) + try: + results = await asyncio.wait_for( + cognee.search(query_text=question, datasets=[slug]), + timeout=timeout_s, + ) + all_results.extend(results or []) + except (asyncio.TimeoutError, Exception): + pass + + def _extract_text(r) -> str: + sr = getattr(r, "search_result", None) + if sr is not None: + s = str(sr).strip() + return "" if (s.startswith("<") and s.endswith(">")) else s + for attr in ("text", "content", "payload", "value"): + val = getattr(r, attr, None) + if isinstance(val, str) and val.strip(): + return val + s = str(r) + return "" if (s.startswith("<") and s.endswith(">")) else s + + return "\n\n".join(t for r in all_results[:5] if (t := _extract_text(r))) + + try: + cognee_ctx = asyncio.run(_search()).strip() + except Exception: + cognee_ctx = "" + + # Raw-file fallback. cognify's entity-extraction LLM calls regularly hit + # the 8192-token non-streaming output ceiling and leave the graph empty, + # so a clean cognee.search() can still return nothing useful. The crawler + # always writes the scraped page text to memory/topic_/*.txt, so we + # use those files as ground-truth context when the graph search comes back + # short. Cap per-slug bytes to keep the prompt bounded. + raw_parts: list[str] = [] + if len(cognee_ctx) < 400: + per_slug_cap = 20_000 + # Keyword-score filenames so the most specific page (e.g. task_environment.txt + # for a "TaskEnvironment" query) is read first. Alphabetical sort exhausts + # the total budget on index/overview files before reaching the relevant one. + query_words = set(re.sub(r"[^a-z0-9]", " ", question.lower()).split()) + + def _file_relevance(p: Path) -> int: + stem_words = set(re.sub(r"[^a-z0-9]", " ", p.stem.lower()).split()) + return -len(query_words & stem_words) # most overlap first + + for slug in target_slugs: + topic_dir = LOCAL_MEMSTORE_ROOT / "memory" / slug + if not topic_dir.exists(): + continue + for fpath in sorted(topic_dir.glob("*.txt"), key=_file_relevance): + try: + text = fpath.read_text(encoding="utf-8", errors="replace") + except Exception: + continue + if not text.strip(): + continue + # Strip Jina header so the window covers actual page content, + # not Title/URL/navigation boilerplate. + text = _strip_jina_header(text) + raw_parts.append(f"--- {slug}/{fpath.name} ---\n{text[:per_slug_cap]}") + if sum(len(p) for p in raw_parts) > 40_000: + break + if sum(len(p) for p in raw_parts) > 40_000: + break + + # Session user memories (raw file read — no Cognee needed) + session_mem_dir = LOCAL_MEMSTORE_ROOT / "user" / "sessions" / session / "memories" + user_mem_parts: list[str] = [] + if session_mem_dir.exists(): + for fpath in sorted(session_mem_dir.glob("*.txt")): + if fpath.name.startswith("_"): + continue + try: + uc = fpath.read_text(encoding="utf-8", errors="replace") + if uc.strip(): + user_mem_parts.append(f"[USER_MEMORY from {fpath.name}]\n{uc[:5000]}") + except Exception: + pass + + parts = [p for p in (cognee_ctx, "\n\n".join(raw_parts), "\n\n".join(user_mem_parts[:10])) if p] + return "\n\n".join(parts) + + +# --------------------------------------------------------------------------- +# Preference helpers +# --------------------------------------------------------------------------- + +def _load_prefs(store: MemoryStore) -> dict: + obj = store.read_json(PREFS_JSON_PATH, default=None) + if isinstance(obj, dict): + return obj + text = store.read_text(PREFS_TXT_PATH, default="") + prefs: dict[str, str] = {} + for ln in text.splitlines(): + if "=" in ln: + k, v = ln.split("=", 1) + k, v = k.strip(), v.strip() + if k: + prefs[k] = v + return prefs + + +def _prefs_to_text(prefs: dict) -> str: + lines = [f"{k}={v}" for k, v in sorted(prefs.items()) if v is not None and v != ""] + return "\n".join(lines).strip() + ("\n" if lines else "") + + +def _chat_transcript_path(session_id: str, session: str = "default") -> str: + return f"user/sessions/{session}/chat/{session_id}/transcript.jsonl" + + +def _chat_summary_path(session_id: str, session: str = "default") -> str: + return f"user/sessions/{session}/chat/{session_id}/summary.txt" + + +def _clip_text(s: str, max_chars: int = 4000) -> str: + s = s or "" + return s if len(s) <= max_chars else s[: max_chars - 1] + "…" + + +def _build_llm_messages(history: list[dict], max_messages: int) -> list[dict]: + # Keep only roles Anthropic expects. + msgs = [m for m in (history or []) if m.get("role") in ("user", "assistant")] + msgs = msgs[-max_messages:] + out: list[dict] = [] + for m in msgs: + out.append({"role": m["role"], "content": _clip_text(str(m.get("content", "")), 8000)}) + return out + + +def _append_transcript(store: MemoryStore, session_id: str, entries: list[dict], *, actor: str, reason: str, session: str = "default") -> None: + """Append JSONL entries to the per-session transcript (stored under user/). + + Keeps the transcript bounded so it doesn't grow forever. + """ + store.ensure_layout() + path = _chat_transcript_path(session_id, session) + existing = store.read_text(path, default="") + lines = existing.splitlines() if existing.strip() else [] + + for e in entries: + lines.append(json.dumps(e, ensure_ascii=False)) + + if len(lines) > CHAT_TRANSCRIPT_MAX_LINES: + lines = lines[-CHAT_TRANSCRIPT_MAX_LINES :] + + content = "\n".join(lines).strip() + ("\n" if lines else "") + expected = store.current_sha(path) or None + store.write_text(path, content, actor=actor, reason=reason, expected_sha=expected, op="append") + + +def _save_prefs(store: MemoryStore, prefs: dict, *, actor: str, reason: str) -> None: + expected = store.current_sha(PREFS_JSON_PATH) or None + store.write_json(PREFS_JSON_PATH, prefs, actor=actor, reason=reason, expected_sha=expected, op="update") + try: + expected_txt = store.current_sha(PREFS_TXT_PATH) or None + store.write_text( + PREFS_TXT_PATH, _prefs_to_text(prefs), + actor=actor, reason=f"derived:{reason}", expected_sha=expected_txt, op="update", + ) + except Exception: + pass + asyncio.run(_upload_memstore()) + + + + +# --------------------------------------------------------------------------- +# Session state +# --------------------------------------------------------------------------- + +def _init_session() -> None: + import streamlit as st + + defaults: dict = { + "active_session": "default", + "_last_active_session": "", + "messages": [], + "last_answer": "", + "sleep_run": None, + "sleep_run_id": "", + "sleep_run_url": "", + "summary_run": None, + "summary_run_id": "", + "summary_run_url": "", + "ingest_run": None, + "ingest_run_url": "", + "promote_run": None, + # keyed by assistant message index → proposal dict with "status": "pending"|"accepted"|"denied" + "memory_proposals": {}, + # UI helpers + "pending_pref_dialog": None, # assistant message index + "toast_queue": [], + "chat_session_id": "", + # Debugging + "last_proposal_raw": "", + "last_proposal_parsed": None, + "last_proposal_error": "", + } + for k, v in defaults.items(): + if k not in st.session_state: + st.session_state[k] = v + + if not st.session_state.get("chat_session_id"): + st.session_state.chat_session_id = uuid.uuid4().hex[:12] + + +def _active_session() -> str: + """Return the currently active session name.""" + import streamlit as st + return st.session_state.get("active_session", "default") + + +def _switch_session(new_session: str) -> None: + """Save the current session's state and load the new session's state.""" + import streamlit as st + + current = st.session_state.get("active_session", "default") + if new_session == current: + return + + # Save current session's messages and chat ID under session-keyed backup keys + st.session_state[f"_messages_{current}"] = list(st.session_state.get("messages", [])) + st.session_state[f"_chat_id_{current}"] = st.session_state.get("chat_session_id", "") + st.session_state[f"_proposals_{current}"] = dict(st.session_state.get("memory_proposals", {})) + + # Load new session's saved state (or empty defaults) + st.session_state["messages"] = list(st.session_state.get(f"_messages_{new_session}", [])) + st.session_state["chat_session_id"] = st.session_state.get(f"_chat_id_{new_session}", "") or uuid.uuid4().hex[:12] + st.session_state["memory_proposals"] = dict(st.session_state.get(f"_proposals_{new_session}", {})) + st.session_state["pending_pref_dialog"] = None + st.session_state["active_session"] = new_session + st.session_state["_last_active_session"] = new_session + + +def _queue_toast(text: str, *, icon: str | None = None) -> None: + import streamlit as st + + st.session_state.toast_queue.append({"text": text, "icon": icon}) + + +def _drain_toasts(max_toasts: int = 3) -> None: + import streamlit as st + + q = st.session_state.get("toast_queue") or [] + if not q: + return + + for item in q[:max_toasts]: + st.toast(item.get("text", ""), icon=item.get("icon")) + + st.session_state.toast_queue = [] + + +def _maybe_open_preference_dialog(store: MemoryStore) -> None: + import streamlit as st + + msg_idx = st.session_state.get("pending_pref_dialog") + if msg_idx is None: + return + + proposal = st.session_state.memory_proposals.get(msg_idx) + if not proposal or proposal.get("status") != "pending" or proposal.get("type") != "preference": + st.session_state.pending_pref_dialog = None + return + + updates = proposal.get("updates", {}) + if not isinstance(updates, dict) or not updates: + st.session_state.pending_pref_dialog = None + return + + @st.dialog("Preference detected") + def _dialog() -> None: + st.caption(proposal.get("reason", "")) + + raw = st.text_area( + "Updates (JSON)", + value=json.dumps(updates, indent=2, sort_keys=True), + key=f"pref_dialog_updates_{msg_idx}", + height=180, + ) + + try: + edited_obj = json.loads(raw) + except Exception: + edited_obj = None + + if isinstance(edited_obj, dict): + st.json(edited_obj, expanded=True) + else: + st.warning("Updates must be valid JSON object") + + col1, col2 = st.columns(2) + if col1.button("Save preference", type="primary", use_container_width=True): + if not isinstance(edited_obj, dict): + st.error("Cannot save: updates JSON must be an object") + return + + try: + validated = ExtractedProposal.model_validate( + {"type": "preference", "updates": edited_obj, "reason": proposal.get("reason", "preference")} + ) + except ValidationError as e: + st.error(f"Cannot save: {e.errors()[0].get('msg', 'invalid updates')}") + return + + current = _load_prefs(store) + current.update(validated.updates or {}) + _save_prefs(store, current, actor="chat", reason=proposal.get("reason", "preference")) + st.session_state.memory_proposals[msg_idx]["status"] = "accepted" + st.session_state.pending_pref_dialog = None + _queue_toast("Preference saved", icon="⚙️") + st.rerun() + + if col2.button("Dismiss", use_container_width=True): + st.session_state.memory_proposals[msg_idx]["status"] = "denied" + st.session_state.pending_pref_dialog = None + st.rerun() + + _dialog() + + +# --------------------------------------------------------------------------- +# Run polling helper +# --------------------------------------------------------------------------- + +def _try_finish_run(run) -> tuple[bool, str, list]: + """Return (done, phase, outputs). outputs is empty if not done/succeeded.""" + terminal = {"SUCCEEDED", "FAILED", "ABORTED", "TIMED_OUT"} + run.sync() # let exceptions propagate so callers can surface them via st.warning + try: + phase = getattr(run.phase, "name", str(run.phase)) + except Exception: + # run.phase raises ValueError("Cannot convert UNSPECIFIED phase") when the + # run object stored in session state loses its cluster connection after sync. + # Fall back to reading outputs from local metadata: if they exist the run + # must have completed successfully. + try: + outs = list(run.outputs()) + if len(outs) == 1 and isinstance(outs[0], (list, tuple)): + outs = list(outs[0]) + if outs: + return True, "SUCCEEDED", outs + except Exception: + pass + return False, "UNSPECIFIED", [] + if phase not in terminal: + return False, phase, [] + if phase != "SUCCEEDED": + return True, phase, [] + outs = list(run.outputs()) + if len(outs) == 1 and isinstance(outs[0], (list, tuple)): + outs = list(outs[0]) + return True, phase, outs + + +# --------------------------------------------------------------------------- +# Inline memory proposal card (shown below each assistant message) +# --------------------------------------------------------------------------- + +def _render_proposal_card(msg_idx: int, proposal: dict, store: MemoryStore) -> None: + import streamlit as st + + status = proposal.get("status") + if status != "pending": + if status == "accepted": + ptype = proposal.get("type", "memory") + label = "Preference saved" if ptype == "preference" else "Memory staged — promoted on next sleep cycle" + st.caption(f"✅ {label}") + return + + ptype = proposal.get("type", "memory") + is_pref = ptype == "preference" + + with st.container(border=True): + if is_pref: + st.caption("⚙️ Preference detected — accept to save immediately") + updates = proposal.get("updates", {}) + st.json(updates, expanded=True) + st.caption(proposal.get("reason", "")) + else: + st.caption("💡 Memory suggestion — accept to stage for the next sleep cycle") + edited = st.text_area( + "Content (edit before accepting)", + value=proposal.get("content", ""), + key=f"proposal_content_{msg_idx}", + height=80, + label_visibility="collapsed", + ) + st.caption(f"`{proposal.get('path', '')}` · {proposal.get('reason', '')}") + + col1, col2, _ = st.columns([1, 1, 5]) + if col1.button("Accept", key=f"accept_{msg_idx}", type="primary"): + if is_pref: + current = _load_prefs(store) + current.update(proposal.get("updates", {})) + _save_prefs(store, current, actor="chat", reason=proposal.get("reason", "preference")) + st.toast("Preference saved", icon="⚙️") + else: + _api_key = os.environ.get("ANTHROPIC_API_KEY", "") + _active_sess = st.session_state.get("active_session", "default") + _topic_slug = classify_proposal_topic( + content=edited, + source_question=proposal.get("source_question", ""), + topic_index=load_topic_index(store), + api_key=_api_key, + ) + prop = MemoryWriteProposal( + target="user", + path=proposal["path"], + content=edited, + author="chat", + reason=proposal.get("reason", ""), + source_question=proposal.get("source_question", ""), + topic_slug=_topic_slug, + session=_active_sess, + ) + stage_proposal(store, prop) + asyncio.run(_upload_memstore()) + st.toast(f"Memory staged in '{_active_sess}' — auto-promoted on next sleep cycle", icon="💡") + st.session_state.memory_proposals[msg_idx]["status"] = "accepted" + if st.session_state.get("pending_pref_dialog") == msg_idx: + st.session_state.pending_pref_dialog = None + st.rerun() + if col2.button("Deny", key=f"deny_{msg_idx}"): + st.session_state.memory_proposals[msg_idx]["status"] = "denied" + if st.session_state.get("pending_pref_dialog") == msg_idx: + st.session_state.pending_pref_dialog = None + st.rerun() + + +# --------------------------------------------------------------------------- +# Sidebar sections +# --------------------------------------------------------------------------- + +def _render_session_selector(store: MemoryStore) -> None: + """Sidebar widget to pick or create a named session.""" + import streamlit as st + + st.subheader("🗃️ Session") + sessions = list_sessions(store) + if not sessions: + register_session(store, "default", label="Default Session") + asyncio.run(_upload_memstore()) + sessions = ["default"] + + active = st.session_state.get("active_session", "default") + if active not in sessions: + active = sessions[0] + st.session_state["active_session"] = active + + selected = st.selectbox( + "Active session", + sessions, + index=sessions.index(active), + key="session_selectbox", + label_visibility="collapsed", + ) + if selected != active: + _switch_session(selected) + st.rerun() + + with st.expander("New session", expanded=False): + with st.form("new_session_form", clear_on_submit=True): + new_name = st.text_input( + "Session name", + placeholder="e.g. work, research, personal", + key="new_session_name_input", + ) + created = st.form_submit_button("Create") + + if created and new_name.strip(): + name = re.sub(r"[^a-zA-Z0-9_-]", "_", new_name.strip())[:40] + if name and name not in sessions: + register_session(store, name, label=new_name.strip()) + asyncio.run(_upload_memstore()) + _switch_session(name) + st.rerun() + elif name in sessions: + st.warning(f"Session '{name}' already exists.") + else: + st.warning("Invalid session name.") + + +def _render_sleep_section() -> None: + import streamlit as st + + st.subheader("🌙 Sleep Cycle") + + sleep_run = st.session_state.get("sleep_run") + sleep_url = st.session_state.get("sleep_run_url", "") + + if sleep_run or sleep_url: + # IMPORTANT: Do NOT auto-poll with run.sync() on every rerun. + try: + phase = getattr(getattr(sleep_run, "phase", None), "name", "") if sleep_run else "" + except Exception: + phase = "" # UNSPECIFIED phase raises ValueError before the run is scheduled + + url = getattr(sleep_run, "url", "") if sleep_run else sleep_url + if url: + st.info("Sleep cycle running") + st.code(url, language="text") + else: + st.info("Sleep running…") + + col1, col2 = st.columns([1, 1]) + if col1.button("Refresh status", use_container_width=True, disabled=not bool(sleep_run)): + try: + done, phase, _ = _try_finish_run(sleep_run) + if done: + st.session_state.sleep_run = None + st.session_state.sleep_run_id = "" + st.session_state.sleep_run_url = "" + if phase == "SUCCEEDED": + st.success("Sleep cycle complete — memory consolidated") + asyncio.run(_download_shared_state()) + else: + st.error(f"Sleep cycle ended: {phase}") + st.rerun() + else: + st.toast(f"Sleep still running: {phase or 'RUNNING'}") + except Exception as e: + st.warning(f"Could not refresh: {e}") + + if col2.button("Clear", use_container_width=True): + st.session_state.sleep_run = None + st.session_state.sleep_run_id = "" + st.session_state.sleep_run_url = "" + st.rerun() + + if not sleep_run: + st.caption("Run handle not available in this session; open the Union UI link to monitor status.") + + else: + st.caption("Runs every 6 hours while app is active · auto-promotes staged memories") + if st.button("Trigger sleep now", use_container_width=True): + try: + run = flyte.run(sleep_cycle) + run_url = str(getattr(run, "url", "")) + st.session_state.sleep_run = run + st.session_state.sleep_run_id = str(getattr(run, "id", "")) + st.session_state.sleep_run_url = run_url + print(f"[sleep_cycle] started: {run_url}") + st.toast("Sleep cycle started", icon="🌙") + st.rerun() + except Exception as e: + st.error(f"Could not start sleep cycle: {e}") + + +def _render_memory_viewer(store: MemoryStore) -> None: + import streamlit as st + + with st.expander("📋 Audit log (tail)", expanded=False): + for ev in store.audit_tail(30): + st.code( + f"{ev.get('ts')} {ev.get('op')} {ev.get('path')} actor={ev.get('actor')}", + language="text", + ) + + topic_index = load_topic_index(store) + with st.expander(f"📚 Topic knowledge base ({len(topic_index)} topics)", expanded=False): + if not topic_index: + st.caption("No topics yet — seed a URL above.") + for slug, entry in sorted(topic_index.items()): + label = entry.get("label", slug) + sources = entry.get("sources", []) + last_updated = entry.get("last_updated", "") + with st.container(border=True): + st.markdown(f"**{slug}** — {label}") + st.caption(f"Updated: {last_updated} | Sources: {len(sources)}") + for src in sources[:3]: + st.caption(f" {src}") + topic_paths = store.list_paths(f"memory/{slug}") + st.caption(f"{len(topic_paths)} file(s) stored") + for p in topic_paths[:50]: + content = store.read_text(p) + size_kb = len(content.encode()) / 1024 + st.caption(f"`{p}` ({size_kb:.1f} KB)") + st.code(content[:400] + ("…" if len(content) > 400 else ""), language="text") + + active_session = st.session_state.get("active_session", "default") + session_mem_prefix = session_memories_prefix(active_session) + user_paths = store.list_paths(session_mem_prefix) + # Exclude internal topic map from the count/display + user_paths = [p for p in user_paths if not Path(p).name.startswith("_")] + + with st.expander(f"Promoted memories — {active_session!r} ({len(user_paths)})", expanded=False): + if not user_paths: + st.caption("No promoted memories for this session yet.") + for p in user_paths[:30]: + st.markdown(f"**{p}**") + st.code(store.read_text(p)[:2000]) + + def _is_archived(proposal_id: str) -> bool: + for decision in ("approved", "rejected", "vetoed", "error", "needs_review"): + if store.exists(f"staging/sessions/{active_session}/archive/{decision}/{proposal_id}.json"): + return True + if store.exists(f"staging/archive/{decision}/{proposal_id}.json"): + return True + return False + + staged = list_staged_proposals(store, session=active_session) + user_staged_all = [p for p in staged if p.target == "user"] + user_staged_pending = [p for p in user_staged_all if not _is_archived(p.id)] + + with st.expander( + f"Staging inbox — {active_session!r} (pending {len(user_staged_pending)} · processed {len(user_staged_all) - len(user_staged_pending)})", + expanded=False, + ): + if not user_staged_pending: + st.caption("No pending staged proposals for this session.") + for prop in user_staged_pending[:10]: + st.markdown(f"`{prop.path}`") + st.caption(prop.reason or "(no reason)") + st.code(prop.content[:400]) + + +def _render_preferences(store: MemoryStore) -> None: + import streamlit as st + + st.subheader("Preferences") + prefs = _load_prefs(store) + + with st.form("prefs_form"): + tone = st.selectbox( + "Tone", + ["concise", "normal", "detailed"], + index=["concise", "normal", "detailed"].index(prefs.get("tone", "concise")) + if prefs.get("tone") in ("concise", "normal", "detailed") else 0, + ) + fmt = st.selectbox( + "Format", ["markdown", "plain"], + index=0 if prefs.get("format", "markdown") == "markdown" else 1, + ) + name = st.text_input("Name (optional)", value=str(prefs.get("name", ""))) + save = st.form_submit_button("Save preferences") + + if save: + new_prefs = dict(prefs) + new_prefs["tone"] = tone + new_prefs["format"] = fmt + if name.strip(): + new_prefs["name"] = name.strip() + else: + new_prefs.pop("name", None) + + _save_prefs(store, new_prefs, actor="streamlit", reason="manual-pref") + st.toast("Saved") + st.rerun() + + with st.expander("Advanced: stage raw proposal", expanded=False): + with st.form("proposal_form"): + path = st.text_input("Path", value="user/notes.txt") + content = st.text_area("Content", height=120) + reason = st.text_input("Reason", value="user requested") + author = st.text_input("Author", value="streamlit") + submitted = st.form_submit_button("Stage proposal") + + if submitted and content.strip(): + _api_key = os.environ.get("ANTHROPIC_API_KEY", "") + _active_sess = st.session_state.get("active_session", "default") + _topic_slug = classify_proposal_topic( + content=content, + source_question="", + topic_index=load_topic_index(store), + api_key=_api_key, + ) + prop = MemoryWriteProposal( + target="user", + path=path.strip(), + content=content, + author=author, + reason=reason, + topic_slug=_topic_slug, + session=_active_sess, + ) + stage_proposal(store, prop) + asyncio.run(_upload_memstore()) + st.success(f"Staged in session '{_active_sess}' — auto-promoted on next sleep cycle") + st.rerun() + + +def _clear_all_memory(store: MemoryStore) -> None: + import shutil + + # Read topic slugs before clearing so we can drop each remote per-topic cognee DB + topic_slugs = list(load_topic_index(store).keys()) + + for subdir in ("user", "memory", Path("staging") / "inbox"): + p = LOCAL_MEMSTORE_ROOT / subdir + if p.exists(): + shutil.rmtree(p) + + if LOCAL_COGNEE_ROOT.exists(): + shutil.rmtree(LOCAL_COGNEE_ROOT) + LOCAL_COGNEE_ROOT.mkdir(parents=True, exist_ok=True) + + store.ensure_layout() + prefs_obj = {"tone": "concise", "format": "markdown"} + store.write_json(PREFS_JSON_PATH, prefs_obj, actor="clear", reason="reset", op="create") + store.write_text( + PREFS_TXT_PATH, + "\n".join(f"{k}={v}" for k, v in sorted(prefs_obj.items())) + "\n", + actor="clear", reason="reset", op="create", + ) + + asyncio.run(_upload_memstore()) + # Upload empty local cognee root to each known topic's remote path (clears remote state) + for slug in topic_slugs: + asyncio.run(Dir.from_local(str(LOCAL_COGNEE_ROOT), remote_destination=_topic_db_path(slug))) + + +def _render_danger_zone(store: MemoryStore) -> None: + import streamlit as st + + with st.expander("⚠️ Danger Zone", expanded=False): + st.caption("Permanently deletes all memories and resets the Cognee knowledge graph. Reference docs are kept.") + confirm = st.checkbox("I understand this cannot be undone", key="confirm_clear") + if st.button("Clear all memory", type="primary", disabled=not confirm, use_container_width=True): + _clear_all_memory(store) + st.session_state.memory_proposals = {} + st.toast("All memory cleared", icon="🗑️") + st.rerun() + + +def _render_knowledge_seeding() -> None: + import streamlit as st + + st.subheader("🌐 Seed Knowledge from URL") + st.caption("Scrape a URL — Claude classifies it into a topic cluster and makes it retrievable by semantic search.") + + ingest_run = st.session_state.get("ingest_run") + ingest_url_val = st.session_state.get("ingest_run_url", "") + + if ingest_run or ingest_url_val: + url_display = getattr(ingest_run, "url", "") if ingest_run else ingest_url_val + if url_display: + st.info("Ingest running") + st.code(url_display, language="text") + else: + st.info("Ingesting…") + + col1, col2 = st.columns([1, 1]) + if col1.button("Refresh status", use_container_width=True, key="refresh_ingest", disabled=not bool(ingest_run)): + try: + done, phase, _ = _try_finish_run(ingest_run) + if done: + st.session_state.ingest_run = None + st.session_state.ingest_run_url = "" + if phase == "SUCCEEDED": + asyncio.run(_download_shared_state()) + index = load_topic_index(MemoryStore(LOCAL_MEMSTORE_ROOT)) + topic_summary = ", ".join(f"`{s}`" for s in sorted(index)[:8]) if index else "none yet" + st.success(f"URL ingested — topics: {topic_summary}") + else: + st.error(f"Ingest ended: {phase}") + st.rerun() + else: + st.toast(f"Ingest still running: {phase or 'RUNNING'}") + except Exception as e: + st.warning(f"Could not refresh: {e}") + + if col2.button("Clear", use_container_width=True, key="clear_ingest"): + st.session_state.ingest_run = None + st.session_state.ingest_run_url = "" + st.rerun() + + if not ingest_run: + st.caption("Run handle not available in this session; open the Union UI link to monitor status.") + else: + url_input = st.text_input( + "Seed URL", + placeholder="https://docs.union.ai/v2/union/user-guide/", + key="ingest_url_input", + ) + max_pages = st.slider("Max pages to crawl", min_value=1, max_value=50, value=10, key="ingest_max_pages") + st.caption("Crawls all linked subpages within the same domain and path prefix.") + if st.button("Ingest URL", use_container_width=True, disabled=not bool((url_input or "").strip())): + url = (url_input or "").strip() + try: + run = flyte.run(ingest_url, url=url, max_pages=max_pages) + run_url = str(getattr(run, "url", "")) + st.session_state.ingest_run = run + st.session_state.ingest_run_url = run_url + print(f"[ingest_url] started for {url!r}: {run_url}") + st.toast(f"Ingesting {url}", icon="🌐") + st.rerun() + except Exception as e: + st.error(f"Could not start ingest: {e}") + + +def _render_sidebar(store: MemoryStore) -> None: + import streamlit as st + + st.header("🗂️ Memory Store") + + _render_session_selector(store) + st.divider() + _render_sleep_section() + st.divider() + _render_knowledge_seeding() + st.divider() + + # Chat continuity (durable transcript + optional summary) + active_session = st.session_state.get("active_session", "default") + session_id = st.session_state.get("chat_session_id", "") + with st.expander("💬 Chat continuity", expanded=False): + st.caption(f"Memory session: `{active_session}` · Chat ID: `{session_id}`") + if session_id: + st.caption(f"Transcript: `{_chat_transcript_path(session_id, active_session)}`") + st.caption(f"Summary: `{_chat_summary_path(session_id, active_session)}`") + current_summary = store.read_text(_chat_summary_path(session_id, active_session), default="").strip() + st.text_area("Current summary", value=current_summary, height=120, disabled=True) + + summary_run = st.session_state.get("summary_run") + summary_url = st.session_state.get("summary_run_url", "") + + if summary_run or summary_url: + url = getattr(summary_run, "url", "") if summary_run else summary_url + if url: + st.caption("Summary running") + st.code(url, language="text") + + if st.button( + "Refresh summary status", + use_container_width=True, + key="refresh_summary", + disabled=not bool(summary_run), + ): + try: + done, phase, _ = _try_finish_run(summary_run) + if done: + st.session_state.summary_run = None + st.session_state.summary_run_id = "" + st.session_state.summary_run_url = "" + if phase == "SUCCEEDED": + asyncio.run(_download_shared_state()) + st.toast("Summary updated") + else: + st.error(f"Summary run ended: {phase}") + st.rerun() + else: + st.toast(f"Summary still running: {phase}") + except Exception as e: + st.warning(f"Could not refresh: {e}") + + if st.button("Clear summary run", use_container_width=True, key="clear_summary_run"): + st.session_state.summary_run = None + st.session_state.summary_run_id = "" + st.session_state.summary_run_url = "" + st.rerun() + + if not summary_run: + st.caption("Run handle not available in this session; open the Union UI link to monitor status.") + + else: + if st.button("Update summary (Flyte)", use_container_width=True): + try: + run = flyte.run(summarize_chat_session, session_id=session_id, session=active_session) + run_url = str(getattr(run, "url", "")) + st.session_state.summary_run = run + st.session_state.summary_run_id = str(getattr(run, "id", "")) + st.session_state.summary_run_url = run_url + print(f"[summarize_chat_session] started for session {session_id!r} ({active_session!r}): {run_url}") + st.toast("Summary task started") + st.rerun() + except Exception as e: + st.error(f"Could not start summary task: {e}") + + _render_memory_viewer(store) + st.divider() + _render_preferences(store) + + if DEBUG: + st.divider() + with st.expander("Debug: last proposal detection", expanded=False): + st.caption("Raw extractor output") + st.code(st.session_state.get("last_proposal_raw", "")[:4000]) + st.caption("Parsed JSON (if any)") + st.json(st.session_state.get("last_proposal_parsed", None)) + err = st.session_state.get("last_proposal_error", "") + if err: + st.caption(f"Validation error: {err}") + + st.divider() + _render_danger_zone(store) + + +# --------------------------------------------------------------------------- +# Main app +# --------------------------------------------------------------------------- + +def _init_flyte() -> None: + """Initialize flyte once at startup (cached so it doesn't re-run on every Streamlit rerun).""" + import streamlit as st + + @st.cache_resource + def _do_init(): + flyte.init_from_config() + return True + + _do_init() + + +_SLEEP_INTERVAL_S = 6 * 3600 # 6 hours + + +def _start_sleep_scheduler() -> None: + """Start a daemon thread that fires sleep_cycle every 6 hours via flyte.run(). + + flyte.Trigger + flyte.Cron on a task is broken on the Union cluster: the cluster + never writes inputs.pb before starting the container, so every triggered execution + fails with READ_FAILED. flyte.run() works correctly because the SDK uploads + inputs.pb via the dataproxy service before creating the execution. + """ + import streamlit as st + + @st.cache_resource + def _start_once(): + def _loop() -> None: + while True: + time.sleep(_SLEEP_INTERVAL_S) + try: + run = flyte.run(sleep_cycle) + print(f"[scheduler] sleep_cycle triggered: {getattr(run, 'url', '')}") + except Exception as e: + print(f"[scheduler] sleep_cycle failed to start: {e}") + + t = threading.Thread(target=_loop, daemon=True, name="sleep-scheduler") + t.start() + return True + + _start_once() + + +def main() -> None: + import streamlit as st + + st.set_page_config(page_title="Cognee Memory Store", layout="wide") + + _init_flyte() + _start_sleep_scheduler() + + # _init_session must run before the seeded check so session_state exists + _init_session() + + if not st.session_state.get("_seeded"): + _ensure_seeded() + st.session_state["_seeded"] = True + + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + + with st.sidebar: + _render_sidebar(store) + + active_session = st.session_state.get("active_session", "default") + st.title("🧠 Cognee + Flyte Memory Store") + st.caption( + f"Session: **{active_session}** · " + "Sleep/wake architecture: Flyte consolidates memories every 6 hours. " + "After each answer, Claude suggests a memory to stage — accept, edit, or deny inline." + ) + + _drain_toasts() + + # Render chat history + any pending proposal cards + for i, msg in enumerate(st.session_state.messages): + with st.chat_message(msg["role"]): + st.markdown(msg["content"]) + if msg["role"] == "assistant": + proposal = st.session_state.memory_proposals.get(i) + if proposal: + _render_proposal_card(i, proposal, store) + + _maybe_open_preference_dialog(store) + + if user_input := st.chat_input("Ask a question…"): + import concurrent.futures + + active_session = st.session_state.get("active_session", "default") + current_prefs = _load_prefs(store) + + st.session_state.messages.append({"role": "user", "content": user_input}) + with st.chat_message("user"): + st.markdown(user_input) + + # Build system prompt + prefs_text = _prefs_to_text(current_prefs).strip() + + # Retrieve context then run answer + proposal detection in parallel + t0 = time.perf_counter() + context = _retrieve_context(user_input, session=active_session) + t_retrieve = time.perf_counter() - t0 + + session_id = st.session_state.get("chat_session_id", "") + chat_summary = store.read_text(_chat_summary_path(session_id, active_session), default="") if session_id else "" + + system = ( + "You are an assistant. Prefer correctness over verbosity.\n" + "Treat the user's latest messages as authoritative for newly introduced facts.\n" + "Treat [preferences] as requirements.\n" + "- If preferences include name=, address the user by that name in every response.\n" + "- If preferences include tone/format, comply.\n" + "- For other preference keys, interpret them as user directives and follow them as best you can.\n" + "When [retrieved] is non-empty, treat it as the authoritative source for the topic and answer " + "primarily from it. If it contradicts your prior knowledge (e.g. an older framework version), " + "follow [retrieved]. Quote API names, decorators, and code samples exactly as they appear there. " + "If [retrieved] does not contain enough to answer, say so explicitly rather than guessing.\n\n" + f"[preferences]\n{prefs_text or '<>'}\n\n" + f"[chat_summary]\n{chat_summary.strip() or '<>'}\n\n" + f"[retrieved]\n{context or '<>'}\n" + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + t1 = time.perf_counter() + llm_messages = _build_llm_messages(st.session_state.messages, CHAT_CONTEXT_MESSAGES) + answer_future = pool.submit(_call_llm, system, llm_messages, 30.0) + proposal_future = pool.submit(_extract_proposal_from_message, user_input) + answer = answer_future.result() + proposal = proposal_future.result() + t_answer = time.perf_counter() - t1 + + with st.chat_message("assistant"): + st.markdown(answer) + if DEBUG: + st.caption( + f"retrieve={t_retrieve:.2f}s answer+proposal={t_answer:.2f}s ctx_chars={len(context)}" + ) + + st.session_state.messages.append({"role": "assistant", "content": answer}) + assistant_msg_idx = len(st.session_state.messages) - 1 + st.session_state.last_answer = answer + + # Persist transcript (durable continuity across restarts). + session_id = st.session_state.get("chat_session_id", "") + if session_id: + now = datetime.now(timezone.utc).isoformat(timespec="seconds") + _append_transcript( + store, + session_id, + [ + {"ts": now, "role": "user", "content": user_input}, + {"ts": now, "role": "assistant", "content": answer}, + ], + actor="chat", + reason="chat-transcript", + session=active_session, + ) + asyncio.run(_upload_memstore()) + + if proposal: + # Avoid spamming proposals that don't change anything. + if proposal.get("type") == "preference" and isinstance(proposal.get("updates"), dict): + updates = {k: v for k, v in proposal["updates"].items() if current_prefs.get(k) != v} + if not updates: + proposal = None + else: + proposal["updates"] = updates + + if proposal: + proposal["status"] = "pending" + proposal["source_question"] = user_input + st.session_state.memory_proposals[assistant_msg_idx] = proposal + + if proposal.get("type") == "preference": + st.session_state.pending_pref_dialog = assistant_msg_idx + _queue_toast("Preference detected — please review", icon="⚙️") + else: + _queue_toast("Memory suggestion ready to review", icon="💡") + + st.rerun() + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +def _looks_like_repo_test_runner() -> bool: + if os.environ.get("SELF_CHECK") == "true": + return True + cfg = os.environ.get("FLYTECTL_CONFIG", "") + return cfg.endswith("/test/config.flyte.yaml") or cfg.endswith("\\test\\config.flyte.yaml") + + +if __name__ == "__main__": + if _looks_like_repo_test_runner(): + from memory_store import _self_check + _self_check() + print("app self-check: ok") + raise SystemExit(0) + + if "--server" in sys.argv: + # Union container entrypoint — started by app_env with --server flag + main() + else: + # Deploy: register sleep schedule + serve app on Union + flyte.init_from_config() + from workflow import env + flyte.deploy(env) + print("Sleep cycle schedule registered (every 6 hours).") + app = flyte.serve(app_env) + print(f"App URL: {app.url}") \ No newline at end of file diff --git a/v2/tutorials/cognee_memory_store/memory_store.py b/v2/tutorials/cognee_memory_store/memory_store.py new file mode 100644 index 00000000..cb0fddb9 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/memory_store.py @@ -0,0 +1,502 @@ +"""Claude-inspired memory store for Flyte tutorials. + +This module implements an *auditable, versioned, file-based* memory store that is +meant to be synced via Flyte object storage (flyte.io.Dir). + +Design goals (inspired by Claude Managed Agents memory stores): +- Many small focused files ("memories") addressed by path. +- Access modes by prefix (e.g. reference/ is read-only). +- Immutable version history for every mutation. +- Append-only audit log. +- Optimistic concurrency via expected sha256 preconditions. + +This is intentionally plain-files + JSON metadata so it stays transparent and +teachable. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Optional + + +class MemoryStoreError(RuntimeError): + pass + + +class AccessDenied(MemoryStoreError): + pass + + +class ConcurrencyError(MemoryStoreError): + def __init__(self, path: str, expected_sha: str, actual_sha: str): + super().__init__( + f"ConcurrencyError for {path!r}: expected_sha={expected_sha} actual_sha={actual_sha}" + ) + self.path = path + self.expected_sha = expected_sha + self.actual_sha = actual_sha + + +def _utc_ts() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +# --------------------------------------------------------------------------- +# Topic index helpers +# --------------------------------------------------------------------------- + +TOPIC_INDEX_PATH = "memory/_index.json" +TOPIC_MAP_PATH = "user/_topic_map.json" + +# --------------------------------------------------------------------------- +# Session helpers +# --------------------------------------------------------------------------- + +SESSION_REGISTRY_PATH = "user/sessions/_registry.json" + + +def session_memories_prefix(session: str) -> str: + return f"user/sessions/{session}/memories" + + +def session_topic_map_path(session: str) -> str: + return f"user/sessions/{session}/memories/_topic_map.json" + + +def session_staging_inbox_prefix(session: str) -> str: + return f"staging/sessions/{session}/inbox" + + +def session_staging_archive_prefix(session: str) -> str: + return f"staging/sessions/{session}/archive" + + +def register_session(store: "MemoryStore", name: str, label: str = "") -> None: + """Idempotently register a named session in the session registry.""" + registry = store.read_json(SESSION_REGISTRY_PATH, default={}) + if name not in registry: + registry[name] = {"created_at_s": time.time(), "label": label or name} + store.write_json(SESSION_REGISTRY_PATH, registry, actor="system", reason="register-session") + + +def list_sessions(store: "MemoryStore") -> list[str]: + """Return sorted list of registered session names.""" + registry = store.read_json(SESSION_REGISTRY_PATH, default={}) + return sorted(registry.keys()) + + +def read_topic_map(store: "MemoryStore", topic_map_path: str = TOPIC_MAP_PATH) -> dict: + """Return {user_rel_path: topic_slug} from the given topic map path.""" + return store.read_json(topic_map_path, default={}) + + +def upsert_topic_map(store: "MemoryStore", rel_path: str, slug: Optional[str], *, topic_map_path: str = TOPIC_MAP_PATH) -> None: + """Set or clear the topic association for a promoted memory file. + + Writes directly to the filesystem, bypassing MemoryStore access control, + because _topic_map.json is machine-managed system metadata. + """ + m = read_topic_map(store, topic_map_path) + if slug: + m[rel_path] = slug + else: + m.pop(rel_path, None) + p = store.root / Path(topic_map_path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(m, indent=2, sort_keys=True), encoding="utf-8") + + +def load_topic_index(store: "MemoryStore") -> dict: + """Return {slug: {label, sources, last_updated}} from memory/_index.json.""" + return store.read_json(TOPIC_INDEX_PATH, default={}) + + +def upsert_topic_index( + store: "MemoryStore", + slug: str, + *, + label: str, + source_url: Optional[str] = None, + actor: str = "system", +) -> None: + """Idempotently add or update one slug entry in memory/_index.json. + + Writes directly to the filesystem, bypassing MemoryStore access control, + because memory/ is machine-managed and the index is its own metadata. + """ + index = load_topic_index(store) + entry = index.get(slug, {"label": label, "sources": [], "last_updated": ""}) + entry["label"] = label + if source_url and source_url not in entry["sources"]: + entry["sources"].append(source_url) + entry["last_updated"] = _utc_ts() + index[slug] = entry + index_path = store.root / Path(TOPIC_INDEX_PATH) + index_path.parent.mkdir(parents=True, exist_ok=True) + index_path.write_text(json.dumps(index, indent=2, sort_keys=True), encoding="utf-8") + + +# --------------------------------------------------------------------------- +# JSON parsing helpers (shared by workflow.py and app.py) +# --------------------------------------------------------------------------- + +def _parse_json_object(text: str) -> "dict | None": + """Parse a JSON object from an LLM response, tolerating code fences and preamble.""" + if not text: + return None + t = text.strip() + if not t or t.lower() == "null": + return None + if t.startswith("```"): + t = re.sub(r"^```(?:json)?\s*", "", t, flags=re.IGNORECASE) + t = re.sub(r"\s*```\s*$", "", t).strip() + if t.lower() == "null": + return None + start = t.find("{") + if start < 0: + return None + try: + obj, _ = json.JSONDecoder().raw_decode(t[start:]) + except Exception: + return None + return obj if isinstance(obj, dict) else None + + +def _parse_json_array(text: str) -> list: + """Parse a JSON array from an LLM response, tolerating code fences.""" + if not text: + return [] + t = text.strip() + if t.startswith("```"): + t = re.sub(r"^```(?:json)?\s*", "", t, flags=re.IGNORECASE) + t = re.sub(r"\s*```\s*$", "", t).strip() + start = t.find("[") + if start < 0: + return [] + try: + obj, _ = json.JSONDecoder().raw_decode(t[start:]) + return obj if isinstance(obj, list) else [] + except Exception: + return [] + + +def _sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _sha256_text(text: str) -> str: + return _sha256_bytes(text.encode("utf-8")) + + +def _ensure_relative_posix(path: str) -> str: + """Normalize and validate a memory path. + + - Must be relative (no leading '/'). + - Must not contain '..'. + - Uses POSIX separators regardless of OS. + """ + p = Path(path) + if p.is_absolute() or str(path).startswith("/"): + raise MemoryStoreError(f"Path must be relative, got {path!r}") + + parts: list[str] = [] + for part in p.parts: + if part in ("", "."): + continue + if part == "..": + raise MemoryStoreError(f"Path traversal is not allowed, got {path!r}") + parts.append(part) + + if not parts: + raise MemoryStoreError("Empty path is not allowed") + + return "/".join(parts) + + +@dataclass(frozen=True) +class MemoryMeta: + path: str + sha256: str + updated_at: str + updated_by: str + reason: str + bytes: int + + +class MemoryStore: + """A directory-backed memory store with audit + versioning.""" + + # Prefix policy (inspired by Claude's read_only vs read_write). + # memory/ is machine-managed (written only by ingest_url and sleep_cycle); + # user proposals must land in user/ only. + READ_ONLY_PREFIXES = ("memory/",) + + def __init__(self, root: Path): + self.root = Path(root) + self._audit_path = self.root / "audit" / "log.jsonl" + self._meta_root = self.root / "meta" + self._versions_root = self.root / "versions" + + def ensure_layout(self) -> None: + (self.root / "audit").mkdir(parents=True, exist_ok=True) + (self.root / "memory").mkdir(parents=True, exist_ok=True) + (self.root / "user").mkdir(parents=True, exist_ok=True) + (self.root / "staging" / "inbox").mkdir(parents=True, exist_ok=True) + self._meta_root.mkdir(parents=True, exist_ok=True) + self._versions_root.mkdir(parents=True, exist_ok=True) + + # --------------------------------------------------------------------- + # Core path helpers + # --------------------------------------------------------------------- + + def _abs_memory_path(self, rel_path: str) -> Path: + rel = _ensure_relative_posix(rel_path) + return self.root / Path(rel) + + def _abs_meta_path(self, rel_path: str) -> Path: + rel = _ensure_relative_posix(rel_path) + return self._meta_root / (rel.replace("/", "__") + ".json") + + def _abs_versions_dir(self, rel_path: str) -> Path: + rel = _ensure_relative_posix(rel_path) + return self._versions_root / rel.replace("/", "__") + + def _assert_can_write(self, rel_path: str) -> None: + rel = _ensure_relative_posix(rel_path) + if any(rel.startswith(p) for p in self.READ_ONLY_PREFIXES): + raise AccessDenied(f"Writes to {rel!r} are not allowed (read-only prefix)") + + # --------------------------------------------------------------------- + # Read / list + # --------------------------------------------------------------------- + + def exists(self, rel_path: str) -> bool: + return self._abs_memory_path(rel_path).exists() + + def read_text(self, rel_path: str, default: str = "") -> str: + p = self._abs_memory_path(rel_path) + try: + return p.read_text(encoding="utf-8") + except FileNotFoundError: + return default + + def read_json(self, rel_path: str, default: Any = None) -> Any: + text = self.read_text(rel_path, default="") + if not text.strip(): + return default + return json.loads(text) + + def list_paths(self, prefix: str = "") -> list[str]: + """List memory file paths under a prefix (relative POSIX paths).""" + prefix_norm = "" if not prefix else _ensure_relative_posix(prefix) + base = self.root / Path(prefix_norm) + if not base.exists(): + return [] + + out: list[str] = [] + for p in base.rglob("*"): + if p.is_dir(): + continue + + rel = p.relative_to(self.root).as_posix() + # Exclude internal bookkeeping. + if rel.startswith("audit/") or rel.startswith("meta/") or rel.startswith("versions/"): + continue + out.append(rel) + + out.sort() + return out + + # --------------------------------------------------------------------- + # Metadata + # --------------------------------------------------------------------- + + def get_meta(self, rel_path: str) -> Optional[MemoryMeta]: + mp = self._abs_meta_path(rel_path) + if not mp.exists(): + return None + try: + raw = json.loads(mp.read_text(encoding="utf-8")) + return MemoryMeta(**raw) + except Exception as e: + raise MemoryStoreError(f"Failed to read meta for {rel_path!r}: {e}") + + def current_sha(self, rel_path: str) -> str: + meta = self.get_meta(rel_path) + if meta is not None: + return meta.sha256 + if not self.exists(rel_path): + return "" + return _sha256_bytes(self._abs_memory_path(rel_path).read_bytes()) + + # --------------------------------------------------------------------- + # Writes (versioned + audited) + # --------------------------------------------------------------------- + + def write_text( + self, + rel_path: str, + content: str, + *, + actor: str = "system", + reason: str = "", + expected_sha: Optional[str] = None, + op: str = "update", + extra_audit: Optional[dict[str, Any]] = None, + ) -> MemoryMeta: + """Write a memory file with optimistic concurrency + audit/versioning. + + expected_sha: if provided, the write is applied only if the current sha matches. + """ + self.ensure_layout() + + rel = _ensure_relative_posix(rel_path) + self._assert_can_write(rel) + + p = self._abs_memory_path(rel) + old_sha = self.current_sha(rel) + if expected_sha is not None and expected_sha != old_sha: + raise ConcurrencyError(rel, expected_sha=expected_sha, actual_sha=old_sha) + + p.parent.mkdir(parents=True, exist_ok=True) + + new_sha = _sha256_text(content) + p.write_text(content, encoding="utf-8") + + # Immutable version snapshot. + versions_dir = self._abs_versions_dir(rel) + versions_dir.mkdir(parents=True, exist_ok=True) + ts = _utc_ts().replace(":", "-") + version_path = versions_dir / f"{ts}_{new_sha}.txt" + # Avoid accidental overwrite when timestamps collide. + if version_path.exists(): + salt = _sha256_bytes(os.urandom(8))[:8] + version_path = versions_dir / f"{ts}_{new_sha}_{salt}.txt" + version_path.write_text(content, encoding="utf-8") + + meta = MemoryMeta( + path=rel, + sha256=new_sha, + updated_at=_utc_ts(), + updated_by=actor, + reason=reason, + bytes=len(content.encode("utf-8")), + ) + self._abs_meta_path(rel).parent.mkdir(parents=True, exist_ok=True) + self._abs_meta_path(rel).write_text(json.dumps(meta.__dict__, indent=2), encoding="utf-8") + + self._append_audit( + { + "ts": meta.updated_at, + "op": op, + "path": rel, + "old_sha": old_sha, + "new_sha": new_sha, + "actor": actor, + "reason": reason, + "version_file": version_path.relative_to(self.root).as_posix(), + **(extra_audit or {}), + } + ) + + return meta + + def write_json( + self, + rel_path: str, + obj: Any, + *, + actor: str = "system", + reason: str = "", + expected_sha: Optional[str] = None, + op: str = "update", + extra_audit: Optional[dict[str, Any]] = None, + ) -> MemoryMeta: + content = json.dumps(obj, indent=2, sort_keys=True) + return self.write_text( + rel_path, + content, + actor=actor, + reason=reason, + expected_sha=expected_sha, + op=op, + extra_audit=extra_audit, + ) + + # --------------------------------------------------------------------- + # Audit + # --------------------------------------------------------------------- + + def _append_audit(self, event: dict[str, Any]) -> None: + self._audit_path.parent.mkdir(parents=True, exist_ok=True) + with self._audit_path.open("a", encoding="utf-8") as f: + f.write(json.dumps(event, sort_keys=True) + "\n") + + def audit_tail(self, n: int = 20) -> list[dict[str, Any]]: + if not self._audit_path.exists(): + return [] + lines = self._audit_path.read_text(encoding="utf-8").splitlines() + tail = lines[-n:] if n > 0 else lines + out: list[dict[str, Any]] = [] + for line in tail: + line = line.strip() + if not line: + continue + try: + out.append(json.loads(line)) + except Exception: + continue + return out + + +def _self_check() -> None: + """Fast local self-check used by CI-friendly paths in tutorial scripts.""" + import tempfile + + with tempfile.TemporaryDirectory() as td: + store = MemoryStore(Path(td)) + store.ensure_layout() + + # Write user memory + m1 = store.write_text( + "user/preferences.txt", + "pref=on", + actor="self-check", + reason="unit", + ) + assert m1.sha256 == store.current_sha("user/preferences.txt") + + # Concurrency + try: + store.write_text( + "user/preferences.txt", + "pref=off", + actor="self-check", + reason="unit", + expected_sha="deadbeef", + ) + raise AssertionError("Expected ConcurrencyError") + except ConcurrencyError: + pass + + # Access control — memory/ is machine-managed and read-only via the store + try: + store.write_text("memory/docs.txt", "nope", actor="x", reason="x") + raise AssertionError("Expected AccessDenied") + except AccessDenied: + pass + + assert store.audit_tail(5), "Expected audit events" + + +if __name__ == "__main__": + _self_check() + print("memory_store self-check: ok") diff --git a/v2/tutorials/cognee_memory_store/workflow.py b/v2/tutorials/cognee_memory_store/workflow.py new file mode 100644 index 00000000..99b030c8 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -0,0 +1,1563 @@ +# /// script +# requires-python = "==3.13" +# dependencies = [ +# "flyte==2.1.5", +# "cognee==1.0.7", # 1.0.3-1.0.6 had broken/inconsistent handler defaults; 1.0.9+ adds a subprocess worker that times out in this image +# "pydantic>=2.11.0", +# "litellm>=1.83.0", +# "anthropic>=0.40.0", +# "fastembed>=0.3.0", +# ] +# main = "main" +# /// + +"""Cognee + Flyte memory stores — sleep/wake workflow. + +Architecture +------------ +Knowledge is organised into per-topic cognee datasets ("topic_"). +URL ingestion classifies content into a topic and cognifies only that dataset. +At query time a cheap Claude classifier routes the question to 1-2 relevant +datasets for targeted retrieval — no reference material injected into every prompt. + +Sleep cycle (autonomous, every 6 h via flyte.Cron): + 1. Download latest state from shared object storage + 2. Auto-promote user/ staged proposals — the validator is the only gate + 3. Cluster related user/ memories by topic prefix + 4. Consolidate each cluster in parallel via flyte.map.aio (Claude merges them) + 5. Per-topic rebuild: empty_dataset → re-add → cognify (background) → memify (background) + 6. Upload updated state; stream live HTML report to Union UI + + Flyte features in play: + - app.py background thread → calls flyte.run(sleep_cycle) every 6 h + - flyte.map.aio → parallel cluster consolidation across pods + - cache="auto" → consolidate_cluster is idempotent on retry + - retries=2 → transient failures (network, cognify) auto-retried + - report=True → live HTML progress streamed to Union UI dashboard + - flyte.group() → per-phase spans visible in execution timeline + +Wake cycle (on-demand, triggered per question): + - Downloads latest consolidated state + - Claude classifier routes question to relevant topic dataset(s) + - Targeted cognee.search(datasets=[slugs]) for retrieved context + - Assembles memory-augmented prompt (preferences + retrieved) + - Calls Claude, returns answer + timing metrics + +Deployment +---------- +Register the sleep schedule (once per cluster): + python workflow.py --deploy + +Run the app: + python app.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import re +import shutil +import tempfile +import time +import urllib.request +from datetime import datetime, timedelta, timezone +from html.parser import HTMLParser +from pathlib import Path +from typing import Optional +from urllib.parse import urlparse + +import flyte +from flyte.io import Dir + +from memory_store import ( + MemoryStore, + SESSION_REGISTRY_PATH, + TOPIC_INDEX_PATH, + TOPIC_MAP_PATH, + load_topic_index, + list_sessions, + register_session, + session_memories_prefix, + session_topic_map_path, + upsert_topic_index, + read_topic_map, + upsert_topic_map, + _parse_json_object, + _parse_json_array, +) + +# Shared object-storage root for cross-run persistence. +# Override via env var to target a different cluster's bucket. +# Bare relative paths (e.g. "cognee-memory-store/memstore") are NOT supported by +# flyte.io.Dir — without a scheme `Dir.download()` treats the path as local and +# `Dir.from_local(remote_destination=...)` uploads to local pod disk. +SHARED_REMOTE_ROOT = os.environ.get( + "COGNEE_MEMORY_STORE_REMOTE_ROOT", + "s3://union-oc-production-persistent/cognee-memory-store", +) +SHARED_MEMSTORE_PATH = f"{SHARED_REMOTE_ROOT}/memstore" +SHARED_COGNEE_DB_PREFIX = f"{SHARED_REMOTE_ROOT}/cognee_db" +LOCAL_MEMSTORE_ROOT = Path("/tmp/memory_store") +LOCAL_COGNEE_ROOT = Path("/tmp/cognee_db") + + +def _topic_db_path(slug: str) -> str: + """Remote object-storage path for a topic's isolated Cognee DB.""" + return f"{SHARED_COGNEE_DB_PREFIX}/{slug}" + +GENERAL_TOPIC_SLUG = "topic_user_general" + +DEFAULT_MODEL = os.environ.get("AI_MEMORY_STORE_MODEL", "claude-haiku-4-5-20251001") +# Cognee entity extraction needs a model with high output capacity (graph JSON can be large). +# Haiku has an 8192-token output ceiling; Sonnet handles denser knowledge graphs. +COGNEE_LLM_MODEL = os.environ.get("COGNEE_LLM_MODEL", "claude-sonnet-4-6") +THIS_DIR = Path(__file__).resolve().parent + + +_IMAGE = ( + flyte.Image.from_uv_script(__file__, name="cognee-memory-store", pre=True) + .with_source_file(THIS_DIR / "memory_store.py", "/root") + .with_source_file(THIS_DIR / "agent.py", "/root") +) + +env = flyte.TaskEnvironment( + name="cognee-memory-store", + secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")], + image=_IMAGE, + resources=flyte.Resources(cpu=2, memory="4Gi"), +) + + +# --------------------------------------------------------------------------- +# URL scraping helpers +# --------------------------------------------------------------------------- + +class _TextExtractor(HTMLParser): + """Minimal HTML → plain-text stripper using only stdlib.""" + + _SKIP = {"script", "style", "head", "meta", "link", "noscript", "nav", "footer"} + + def __init__(self) -> None: + super().__init__() + self._parts: list[str] = [] + self._skip_depth = 0 + + def handle_starttag(self, tag: str, attrs: list) -> None: + if tag.lower() in self._SKIP: + self._skip_depth += 1 + + def handle_endtag(self, tag: str) -> None: + if tag.lower() in self._SKIP and self._skip_depth: + self._skip_depth -= 1 + + def handle_data(self, data: str) -> None: + if not self._skip_depth: + stripped = data.strip() + if stripped: + self._parts.append(stripped) + + def get_text(self) -> str: + return "\n".join(self._parts) + + +def _fetch_url_text(url: str, max_bytes: int = 500_000) -> str: + """Fetch a URL and return plain text. + + Strategy 1: Jina Reader (r.jina.ai) — handles JS-rendered pages and SPAs, + returns clean markdown. No auth needed for public URLs. + Strategy 2: Direct HTTP fetch with HTML stripping — fallback for sites where + Jina is unavailable or times out. + """ + # Strategy 1: Jina Reader + try: + jina_req = urllib.request.Request( + f"https://r.jina.ai/{url}", + headers={"Accept": "text/plain", "User-Agent": "Mozilla/5.0"}, + ) + with urllib.request.urlopen(jina_req, timeout=30) as resp: + jina_text = resp.read(max_bytes).decode("utf-8", errors="replace").strip() + if len(jina_text) > 200: + print(f"[ingest] fetch strategy: jina-reader ({len(jina_text):,} chars)") + return re.sub(r'\n{3,}', '\n\n', jina_text) + except Exception: + pass + + print("[ingest] fetch strategy: direct-http (jina unavailable or too short)") + # Strategy 2: Direct fetch + HTML stripping + direct_req = urllib.request.Request( + url, + headers={ + "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + "Accept-Language": "en-US,en;q=0.9", + }, + ) + with urllib.request.urlopen(direct_req, timeout=30) as resp: + raw = resp.read(max_bytes) + content_type = resp.headers.get("Content-Type", "") + + charset = "utf-8" + if "charset=" in content_type: + charset = content_type.split("charset=")[-1].split(";")[0].strip() or "utf-8" + text = raw.decode(charset, errors="replace") + + if " None: + super().__init__() + self.links: list[str] = [] + + def handle_starttag(self, tag: str, attrs: list) -> None: + if tag.lower() == "a": + for name, val in attrs: + if name == "href" and val: + self.links.append(val) + + +def _extract_links(html: str, base_url: str) -> list[str]: + """Return deduplicated absolute URLs found in anchor tags.""" + from urllib.parse import urljoin + + extractor = _LinkExtractor() + extractor.feed(html) + seen: set[str] = set() + result: list[str] = [] + for raw in extractor.links: + absolute = urljoin(base_url, raw).split("#")[0].rstrip("/") + if absolute and absolute not in seen: + seen.add(absolute) + result.append(absolute) + return result + + +def _fetch_raw_html(url: str, max_bytes: int = 300_000) -> str: + """Lightweight raw HTML fetch for link discovery (no Jina, no stripping).""" + try: + req = urllib.request.Request( + url, + headers={ + "User-Agent": "Mozilla/5.0 (compatible; CogneeBot/1.0)", + "Accept": "text/html,application/xhtml+xml;q=0.9,*/*;q=0.8", + }, + ) + with urllib.request.urlopen(req, timeout=15) as resp: + return resp.read(max_bytes).decode("utf-8", errors="replace") + except Exception: + return "" + + +def _strip_jina_header(text: str) -> str: + """Strip Jina Reader preamble and page navigation chrome before actual content. + + Jina returns pages in two formats: + 1. With metadata block: "Title: ...\nURL Source: ...\nMarkdown Content:\n..." + 2. Without metadata: raw markdown starting directly with nav chrome + + Either way, navigation chrome (logo, nav links, banners, search bar) always + appears before the first real section heading or code block. We scan forward + to find where real content begins. + + A line is nav chrome if it has fewer than 3 prose words after stripping + markdown links and images. Real content starts at the first heading + (## or deeper, or # without "|"), code fence, or prose-dense line. + """ + # Format 1: strip Jina metadata preamble when present + marker = "\nMarkdown Content:\n" + idx = text.find(marker) + if idx >= 0: + text = text[idx + len(marker):] + + def _is_nav_chrome(line: str) -> bool: + s = line.strip() + if not s: + return True + # Our own ingest-written comments + if s.startswith("# Source:") or s.startswith("# (truncated"): + return True + # Jina's "# Page Title|Site Name" heading is always site chrome + if s.startswith("# ") and "|" in s: + return True + # Any other heading or code fence marks real content — stop here + if s.startswith("#") or s.startswith("```"): + return False + # Lines whose non-link text has fewer than 3 four-letter prose words are nav chrome. + # Using {4,} avoids short UI labels ("OSS", "v2") being counted as prose. + # URLs are stripped explicitly so path components ("union", "docs") don't + # inflate the count — dangling ](url) left by nested image-link removal also + # contains URL words and would otherwise pass through as "real content". + cleaned = re.sub(r"!?\[[^\]]*\]\([^)]*\)", "", s) # remove links/images + cleaned = re.sub(r"https?://\S+", "", cleaned) # remove bare URLs + cleaned = re.sub(r"`[^`]*`", "", cleaned) # remove inline code + return len(re.findall(r"\b[a-zA-Z]{4,}\b", cleaned)) < 3 + + lines = text.splitlines() + i = 0 + # Scan forward until _is_nav_chrome returns False (real prose or heading found). + # No line cap — Mintlify/React sidebars can have 500+ nav lines before content. + while i < len(lines) and _is_nav_chrome(lines[i]): + i += 1 + + return "\n".join(lines[i:]).strip() + + +def _extract_links_from_markdown(markdown: str, base_url: str) -> list[str]: + """Extract absolute URLs from Jina Reader's markdown output (inline links only).""" + from urllib.parse import urljoin + + seen: set[str] = set() + result: list[str] = [] + for raw in re.findall(r'\]\(([^)\s]+)\)', markdown): + raw = raw.split("#")[0].rstrip("/") + if not raw: + continue + if raw.startswith(("http://", "https://")): + absolute = raw + else: + absolute = urljoin(base_url, raw).split("#")[0].rstrip("/") + if absolute and absolute not in seen: + seen.add(absolute) + result.append(absolute) + return result + + +def _crawl_site(seed_url: str, max_pages: int = 50) -> list[str]: + """BFS crawl from seed_url, staying within the same domain and path prefix. + + Returns an ordered list of discovered page URLs (seed first). + Scope is limited to URLs whose path starts with the seed's path prefix so that + e.g. https://docs.union.ai/v2/union/ only crawls /v2/union/* pages. + + For JS-rendered sites (e.g. Mintlify, Docusaurus) raw HTML contains no + navigation tags. In those cases we fall back to Jina Reader and parse markdown + links from the rendered output. + """ + parsed_seed = urlparse(seed_url) + base_domain = parsed_seed.netloc + seed_path = parsed_seed.path + path_prefix = seed_path[: seed_path.rfind("/") + 1] if "/" in seed_path else "/" + + visited: set[str] = set() + queue: list[str] = [seed_url.rstrip("/")] + discovered: list[str] = [] + + while queue and len(discovered) < max_pages: + url = queue.pop(0) + clean = url.rstrip("/") + if clean in visited: + continue + visited.add(clean) + + html = _fetch_raw_html(clean) + if not html: + continue + + discovered.append(clean) + print(f"[crawl] discovered ({len(discovered)}/{max_pages}): {clean}") + + candidate_links = _extract_links(html, clean) + + # JS-rendered sites (Mintlify, Docusaurus, etc.) put navigation in React + # bundles — raw HTML has header/footer tags but none within the crawl + # scope. Check in-scope count (not total) before deciding to fall back. + in_scope_raw = [ + lnk for lnk in candidate_links + if urlparse(lnk).netloc == base_domain + and urlparse(lnk).path.startswith(path_prefix) + and urlparse(lnk).scheme in ("http", "https") + and lnk.rstrip("/") != clean # exclude self + ] + if len(in_scope_raw) < 3: + try: + jina_req = urllib.request.Request( + f"https://r.jina.ai/{clean}", + headers={"Accept": "text/plain", "User-Agent": "Mozilla/5.0"}, + ) + with urllib.request.urlopen(jina_req, timeout=30) as resp: + jina_md = resp.read(300_000).decode("utf-8", errors="replace") + candidate_links = _extract_links_from_markdown(jina_md, clean) + print(f"[crawl] jina link fallback: {len(candidate_links)} link(s) on {clean}") + except Exception as e: + print(f"[crawl] jina link fallback failed: {e}") + + for link in candidate_links: + p = urlparse(link) + if ( + p.netloc == base_domain + and p.path.startswith(path_prefix) + and p.scheme in ("http", "https") + and link.rstrip("/") not in visited + ): + queue.append(link) + + print(f"[crawl] total pages discovered: {len(discovered)}") + return discovered + + +def _classify_topic( + content_preview: str, + url_path: str, + existing_slugs: dict[str, str], + api_key: str, +) -> dict: + """Classify scraped content into a topic slug using Claude. + + Returns {"slug": "topic_x", "label": "Human Label", "is_new": bool}. + Falls back to a slug derived from the URL path on parse failure. + """ + from anthropic import Anthropic + + existing_lines = "\n".join(f" {s}: {l}" for s, l in list(existing_slugs.items())[:20]) + existing_block = f"Existing topics:\n{existing_lines}" if existing_slugs else "No topics exist yet — create a new one." + + prompt = ( + f"{existing_block}\n\n" + f"URL path: {url_path}\n\n" + f"Content preview (first 2000 chars):\n{content_preview[:2000]}\n\n" + "Return a JSON object with:\n" + ' {"slug": "topic_", "label": "Human Readable Label", "is_new": true|false}\n' + "Assign to an existing slug if the content clearly belongs there, otherwise create a new one.\n" + "slug must start with topic_ and use only lowercase letters, digits, and underscores.\n" + "Return JSON only, no explanation." + ) + + try: + client = Anthropic(api_key=api_key, timeout=20.0) + msg = client.messages.create( + model=DEFAULT_MODEL, + max_tokens=120, + temperature=0, + messages=[{"role": "user", "content": prompt}], + ) + result = _parse_json_object(msg.content[0].text) + if result and isinstance(result.get("slug"), str): + raw_slug = result["slug"] + slug = "topic_" + re.sub(r"[^a-z0-9]+", "_", raw_slug.lower().removeprefix("topic_")).strip("_")[:40] + label = str(result.get("label", slug)) + is_new = slug not in existing_slugs + return {"slug": slug, "label": label, "is_new": is_new} + except Exception: + pass + + # Fallback: derive slug from URL path + slug = "topic_" + re.sub(r"[^a-z0-9]+", "_", url_path.lower()).strip("_")[:40] + return {"slug": slug, "label": slug.replace("_", " ").title(), "is_new": slug not in existing_slugs} + + +def _route_query_to_topics( + question: str, + topic_index: dict, + api_key: str, +) -> list[str]: + """Return 0-2 dataset slugs most relevant to the question. + + Returns [] for general/cross-cutting questions — caller should then search + across all datasets without scoping. + """ + if not topic_index: + return [] + + from anthropic import Anthropic + + topic_lines = "\n".join(f" {s}: {v.get('label', s)}" for s, v in list(topic_index.items())[:20]) + prompt = ( + f"Topics:\n{topic_lines}\n\n" + f"Question: {question}\n\n" + "Return a JSON array of 0-2 topic slugs most relevant to this question.\n" + "Return [] if none match well or the question is general.\n" + "Return JSON only, no explanation." + ) + + try: + client = Anthropic(api_key=api_key, timeout=15.0) + msg = client.messages.create( + model=DEFAULT_MODEL, + max_tokens=80, + temperature=0, + messages=[{"role": "user", "content": prompt}], + ) + slugs = _parse_json_array(msg.content[0].text) + return [s for s in slugs if isinstance(s, str) and s in topic_index][:2] + except Exception: + return [] + + +def _url_to_filename(url: str) -> str: + """Convert a URL to a safe, readable filename stem.""" + parsed = urlparse(url) + domain = re.sub(r'[^a-zA-Z0-9]', '_', parsed.netloc) + path = re.sub(r'[^a-zA-Z0-9]', '_', parsed.path) + name = re.sub(r'_+', '_', f"{domain}{path}").strip("_")[:80] + return name or "scraped" + + +# --------------------------------------------------------------------------- +# Internal helpers (not Flyte tasks — called within task bodies) +# --------------------------------------------------------------------------- + +def _setup_cognee_env(local_cognee_root: Path = LOCAL_COGNEE_ROOT) -> None: + """Set env vars that cognee reads at first import (before config is cached).""" + local_cognee_root.mkdir(parents=True, exist_ok=True) + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + os.environ["DATA_ROOT_DIRECTORY"] = str(local_cognee_root / "data") + os.environ["SYSTEM_ROOT_DIRECTORY"] = str(local_cognee_root / "system") + os.environ["LLM_PROVIDER"] = "anthropic" + os.environ["LLM_MODEL"] = COGNEE_LLM_MODEL + os.environ["LLM_API_KEY"] = api_key + os.environ["EMBEDDING_PROVIDER"] = "fastembed" + os.environ["EMBEDDING_MODEL"] = "BAAI/bge-small-en-v1.5" + os.environ.setdefault("COGNEE_SKIP_CONNECTION_TEST", "true") + # Pin the graph DB provider/handler. Some cognee builds default to "ladybug" + # which isn't in the supported_dataset_database_handlers registry, causing + # `KeyError: 'ladybug'` during cognify. Kuzu is the embedded default that + # actually ships in the wheel. + os.environ["GRAPH_DATABASE_PROVIDER"] = "kuzu" + os.environ["GRAPH_DATASET_DATABASE_HANDLER"] = "kuzu" + # Flyte pods set LOG_LEVEL to a numeric string (e.g. "30") which cognee's + # setup_logging() can't translate — it indexes a name→int dict with the raw + # value and crashes with KeyError. Override with a level name cognee accepts. + log_level_raw = os.environ.get("LOG_LEVEL", "INFO") + if log_level_raw.upper() not in {"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"}: + os.environ["LOG_LEVEL"] = "INFO" + + +def _configure_cognee_runtime(cognee_module, local_cognee_root: Path) -> None: + """Override cognee's cached config via its Python API after import. + + get_llm_config() and get_graph_config() are cached — env var changes after + first import are invisible to them. This function pushes values directly into + the live config objects, which is the only reliable way to reconfigure cognee + when switching topic DBs within the same process (wake_cycle per-topic loop). + + llm_args passes max_tokens to work around a cognee bug: AnthropicAdapter accepts + max_completion_tokens in its constructor but never forwards it to messages.create(), + which requires max_tokens as a mandatory field. Without it every entity extraction + call fails with "Missing required arguments". + """ + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + cognee_module.config.set_llm_config({ + "llm_provider": "anthropic", + "llm_model": COGNEE_LLM_MODEL, + "llm_api_key": api_key, + # 8192 is the largest Anthropic accepts for *non-streaming* requests on + # Sonnet 4.6. Going higher trips "Streaming is required for operations + # that may take longer than 10 minutes" from inside cognee, which uses + # blocking calls. To keep outputs small, MAX_CHARS in ingest_url caps + # per-page input so the generated entity-extraction JSON fits. + "llm_args": {"max_tokens": 8192}, + }) + # system_root_directory() cascades to graph + vector + relational DB paths + cognee_module.config.system_root_directory(str(local_cognee_root / "system")) + cognee_module.config.set_embedding_config({ + "embedding_provider": "fastembed", + "embedding_model": "BAAI/bge-small-en-v1.5", + "embedding_dimensions": 384, + }) + + +async def _download_dir(d: Dir, local_path: Path) -> None: + local_path.mkdir(parents=True, exist_ok=True) + await d.download(local_path=str(local_path)) + + +async def _upload_dir(local_path: Path, remote_destination: str) -> Dir: + return await Dir.from_local(str(local_path), remote_destination=remote_destination) + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def _is_already_archived(store: MemoryStore, proposal_id: str, session: str = "default") -> bool: + """Return True if a proposal has already been processed (promoted, rejected, or vetoed).""" + for decision in ("approved", "rejected", "vetoed", "error", "needs_review"): + if store.exists(f"staging/sessions/{session}/archive/{decision}/{proposal_id}.json"): + return True + # Backward compat: check old-style paths for default session + if session == "default" and store.exists(f"staging/archive/{decision}/{proposal_id}.json"): + return True + return False + + +def _try_parse_iso(ts: str) -> Optional[datetime]: + try: + return datetime.fromisoformat(ts) + except Exception: + return None + + +async def _preserve_newest_preferences_before_upload() -> None: + """Avoid clobbering preferences due to long-running sleep cycles. + + sleep_cycle downloads remote state at the start, mutates locally, and uploads + at the end. If preferences were updated remotely during the run (e.g. a user + clicked 'Save' in the app), the final upload can overwrite them. + + This function re-downloads the remote memstore right before upload and keeps + the newer copy of user/preferences.json + user/preferences.txt. + """ + with tempfile.TemporaryDirectory() as td: + remote_root = Path(td) / "memstore_remote" + await _download_dir(Dir(path=SHARED_MEMSTORE_PATH), remote_root) + + remote_store = MemoryStore(remote_root) + local_store = MemoryStore(LOCAL_MEMSTORE_ROOT) + + pref_json = "user/preferences.json" + pref_txt = "user/preferences.txt" + + rmeta = remote_store.get_meta(pref_json) + lmeta = local_store.get_meta(pref_json) + + rts = _try_parse_iso(rmeta.updated_at) if rmeta else None + lts = _try_parse_iso(lmeta.updated_at) if lmeta else None + + remote_newer = False + if rts and lts: + remote_newer = rts > lts + elif rts and not lts: + remote_newer = True + + if not remote_newer: + return + + src_json = remote_root / pref_json + src_txt = remote_root / pref_txt + dst_json = LOCAL_MEMSTORE_ROOT / pref_json + dst_txt = LOCAL_MEMSTORE_ROOT / pref_txt + + if src_json.exists(): + dst_json.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src_json, dst_json) + if src_txt.exists(): + dst_txt.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src_txt, dst_txt) + + +def _cluster_user_memories(store: MemoryStore, session: str = "default") -> list[dict]: + """Group a session's promoted memories by topic prefix for parallel consolidation. + + Groups files whose stems share a common base (everything before the first '_'). + Skips JSON files (structured data). Only returns groups with 2+ members. + + Example: notes_flyte.txt + notes_tasks.txt → cluster "notes". + """ + prefix = session_memories_prefix(session) + paths = [p for p in store.list_paths(prefix) if not p.endswith(".json")] + groups: dict[str, list[dict]] = {} + for path in paths: + stem = Path(path).stem + label = stem.split("_")[0] + content = store.read_text(path, default="") + if content.strip(): + groups.setdefault(label, []).append({"path": path, "content": content}) + + return [ + {"label": label, "memories": mems} + for label, mems in groups.items() + if len(mems) >= 2 + ] + + +# --------------------------------------------------------------------------- +# Init task +# --------------------------------------------------------------------------- + +@env.task +async def init_memory_store() -> Dir: + """Seed memory store with default preferences and empty topic index, upload to shared storage.""" + with flyte.group("init:setup"): + LOCAL_MEMSTORE_ROOT.mkdir(parents=True, exist_ok=True) + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + store.ensure_layout() + + # Seed empty topic index — write directly since memory/ is read-only in the store + index_path = LOCAL_MEMSTORE_ROOT / Path(TOPIC_INDEX_PATH) + index_path.parent.mkdir(parents=True, exist_ok=True) + if not index_path.exists(): + index_path.write_text("{}", encoding="utf-8") + + prefs_obj = {"tone": "concise", "format": "markdown"} + store.write_json( + "user/preferences.json", prefs_obj, + actor="init_memory_store", reason="seed", op="create", + ) + store.write_text( + "user/preferences.txt", + "\n".join(f"{k}={v}" for k, v in sorted(prefs_obj.items())) + "\n", + actor="init_memory_store", reason="seed", op="create", + ) + register_session(store, "default", label="Default Session") + + with flyte.group("init:upload"): + memstore_dir = await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) + # No cognee_db upload — per-topic DBs are created on first ingest + + return memstore_dir + + +# --------------------------------------------------------------------------- +# URL ingestion task — on-demand, triggered per URL +# --------------------------------------------------------------------------- + +@env.task(retries=1, timeout=timedelta(minutes=30)) +async def ingest_url(url: str, max_pages: int = 10) -> tuple[Dir, Dir]: + """Crawl a URL and all linked subpages, classify into a topic cluster, index in cognee. + + Pipeline: + 1. Crawl: BFS from seed URL, staying within the same domain + path prefix, + up to max_pages pages. Link discovery uses raw HTML; content fetching + uses Jina Reader (handles JS-rendered SPAs) with direct-HTTP fallback. + 2. Classify: Claude assigns or creates a topic slug from the seed page content. + 3. Write: each page → memory/topic_/.txt; update _index.json. + 4. Index: cognee.add() per page, then cognify(background) + memify(background). + 5. Upload: pod stays alive through upload, giving background cognee tasks time + to process. + + Flyte features: + retries=1 handles transient network / cognee failures + timeout=30 min allows crawling up to 50 pages via Jina Reader + """ + with flyte.group("ingest:download"): + await _download_dir(Dir(path=SHARED_MEMSTORE_PATH), LOCAL_MEMSTORE_ROOT) + # Per-topic cognee DB downloaded after classification — slug is unknown until then + + LOCAL_MEMSTORE_ROOT.mkdir(parents=True, exist_ok=True) + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + store.ensure_layout() + + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + MAX_CHARS = 60_000 # captures full Jina-rendered docs page; cognify's chunk_size=512 splits it internally so each LLM call stays under the 8192-token non-streaming output ceiling + + with flyte.group("ingest:crawl"): + print(f"[ingest] Crawling from seed: {url} (max_pages={max_pages})") + all_urls = _crawl_site(url, max_pages=max_pages) + if not all_urls: + raise ValueError(f"No pages reachable from {url!r}") + print(f"[ingest] Crawl complete — {len(all_urls)} page(s) to ingest") + + with flyte.group("ingest:classify"): + # Classify topic using the seed page content + print(f"[ingest] Fetching seed page for classification: {all_urls[0]}") + seed_text = _fetch_url_text(all_urls[0]) + topic_index = load_topic_index(store) + existing_slugs = {s: e.get("label", s) for s, e in topic_index.items()} + parsed_seed = urlparse(url) + url_path = parsed_seed.netloc + parsed_seed.path + classification = _classify_topic(seed_text[:2000], url_path, existing_slugs, api_key) + slug = classification["slug"] + label = classification["label"] + print(f"[ingest] classified → {slug!r} ({'new' if classification['is_new'] else 'existing'}): {label}") + upsert_topic_index(store, slug, label=label, source_url=url, actor="ingest_url") + + # Initialize Cognee with a fresh per-topic dir. We intentionally do NOT + # re-download the prior cognee_db from remote — when the cognee library + # version differs from what wrote those files (e.g. Kuzu↔Ladybug schema + # rename), downloads cause cryptic "version_code" / KeyError crashes during + # cognify. Each ingest re-builds the graph from the source documents. + local_cognee = Path(f"/tmp/cognee_db_{slug}") + if local_cognee.exists(): + shutil.rmtree(local_cognee, ignore_errors=True) + local_cognee.mkdir(parents=True, exist_ok=True) + _setup_cognee_env(local_cognee) + import cognee + _configure_cognee_runtime(cognee, local_cognee) + + with flyte.group("ingest:fetch_and_write"): + topic_dir = LOCAL_MEMSTORE_ROOT / "memory" / slug + topic_dir.mkdir(parents=True, exist_ok=True) + total_chars = 0 + + for page_url in all_urls: + try: + text = _fetch_url_text(page_url) if page_url != all_urls[0] else seed_text + if not text.strip(): + print(f"[ingest] skip (empty): {page_url}") + continue + + text = _strip_jina_header(text) + truncated = len(text) > MAX_CHARS + text = text[:MAX_CHARS] + content = ( + f"# Source: {page_url}\n" + + (f"# (truncated to {MAX_CHARS} chars)\n" if truncated else "") + + f"\n{text}\n" + ) + filename = _url_to_filename(page_url) + file_path = topic_dir / f"{filename}.txt" + file_path.write_text(content, encoding="utf-8") + total_chars += len(content) + print(f"[ingest] written memory/{slug}/{filename}.txt ({len(content):,} chars)") + + await cognee.add(f"[REFERENCE]\n{content}", dataset_name=slug) + except Exception as e: + print(f"[ingest] error on {page_url}: {type(e).__name__}: {e}") + + print(f"[ingest] total written: {total_chars:,} chars across {len(all_urls)} page(s)") + + # Upload memstore first so scraped files and topic index are persisted even if + # cognify times out. With large ingestions (50 pages), Cognee's non-streaming + # entity-extraction calls can exceed Anthropic's timeout ceiling, causing the + # whole task to fail and losing all the written files. Raw files are always + # useful for retrieval via the raw-file fallback in _retrieve_context. + with flyte.group("ingest:upload_memstore"): + print("[ingest] Uploading memstore (before cognify) ...") + memstore_dir = await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) + print("[ingest] Memstore upload complete.") + + with flyte.group("ingest:cognee"): + # chunk_size caps tokens per chunk. cognee's default (~max_chunk_tokens + # of the embedding model) is large enough that the entity-extraction + # JSON for a dense docs page overflows the 8192-token non-streaming + # output ceiling. 512 keeps each call's output well under that. + print(f"[ingest] cognee.cognify(datasets=[{slug!r}], chunk_size=512) ...") + try: + await cognee.cognify(datasets=[slug], chunk_size=512) + print(f"[ingest] cognify complete for {slug!r}") + cognify_ok = True + except Exception as e: + print(f"[ingest] cognify failed (raw files still available for retrieval): {type(e).__name__}: {e}") + cognify_ok = False + + with flyte.group("ingest:upload_cognee"): + print("[ingest] Uploading cognee DB ...") + if cognify_ok: + cognee_dir = await _upload_dir(local_cognee, _topic_db_path(slug)) + else: + # Upload whatever partial state cognee produced — may be empty but + # avoids leaving stale state from a previous run at the remote path. + try: + cognee_dir = await _upload_dir(local_cognee, _topic_db_path(slug)) + except Exception: + cognee_dir = memstore_dir # fallback: return memstore dir as placeholder + print("[ingest] Upload complete.") + + return memstore_dir, cognee_dir + + +# --------------------------------------------------------------------------- +# Consolidation subtask — called via flyte.map.aio inside sleep_cycle +# --------------------------------------------------------------------------- + +@env.task( + cache="auto", + retries=1, + timeout=timedelta(minutes=5), +) +async def consolidate_cluster(cluster_json: str) -> str: + """Merge a cluster of related memories into one coherent summary using Claude. + + Accepts JSON: {"label": str, "memories": [{"path": str, "content": str}]} + Returns JSON: {"path": str, "content": str, "merged_from": [str]} + + cache="auto": if this cluster's content hasn't changed since the last sleep + cycle (e.g. after a crash/retry), the cached result is returned — no + redundant Claude call, no redundant Flyte pod spin-up. + """ + from anthropic import Anthropic + + cluster = json.loads(cluster_json) + memories: list[dict] = cluster["memories"] + + if len(memories) == 1: + m = memories[0] + return json.dumps({"path": m["path"], "content": m["content"], "merged_from": []}) + + label = cluster["label"] + combined = "\n\n".join(f"--- {m['path']} ---\n{m['content']}" for m in memories) + + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + if not api_key: + m = memories[0] + return json.dumps({"path": m["path"], "content": m["content"], "merged_from": []}) + + client = Anthropic(api_key=api_key, timeout=60.0) + msg = client.messages.create( + model=DEFAULT_MODEL, + max_tokens=1200, + system=( + "You consolidate related memory entries into a single coherent summary. " + "Preserve all distinct facts. Remove duplicates and contradictions (keep newest). " + "Be concise but complete. Return only the consolidated text, no preamble." + ), + messages=[{ + "role": "user", + "content": ( + f"Consolidate these {len(memories)} related memories about '{label}':\n\n" + f"{combined}" + ), + }], + temperature=0, + ) + content = msg.content[0].text + canonical_path = memories[0]["path"] + merged_from = [m["path"] for m in memories[1:]] + return json.dumps({"path": canonical_path, "content": content, "merged_from": merged_from}) + + +# --------------------------------------------------------------------------- +# Per-topic rebuild — fanned out via flyte.map.aio inside sleep_cycle +# --------------------------------------------------------------------------- + +@env.task(retries=1, timeout=timedelta(minutes=15)) +async def rebuild_topic_dataset(rebuild_json: str) -> str: + """Rebuild Cognee knowledge graph for one topic in an isolated pod. + + Accepts JSON: {"slug": str} + + Downloads the latest memstore (read-only) and the per-topic cognee DB, + clears stale nodes, re-adds all content with source tags, then fires + cognify + memify in the background before uploading the updated DB. + + Returns JSON: {"slug": str, "ref_docs": int, "user_docs": int} + + Using a separate pod per topic means all topics rebuild in parallel + (concurrency=3 in the calling flyte.map.aio), and each pod only touches + its own isolated cognee DB — no shared-DB write conflicts. + """ + data = json.loads(rebuild_json) + slug = data["slug"] + local_cognee = Path(f"/tmp/cognee_db_{slug}") + + with flyte.group(f"rebuild:{slug}:download"): + await _download_dir(Dir(path=SHARED_MEMSTORE_PATH), LOCAL_MEMSTORE_ROOT) + await _download_dir(Dir(path=_topic_db_path(slug)), local_cognee) + + _setup_cognee_env(local_cognee) + import cognee + _configure_cognee_runtime(cognee, local_cognee) + + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + ref_docs = 0 + + with flyte.group(f"rebuild:{slug}:index"): + datasets_list = await cognee.datasets.list_datasets() + ds_by_name = {ds.name: ds for ds in (datasets_list or [])} + if slug in ds_by_name: + await cognee.datasets.empty_dataset(ds_by_name[slug].id) + + # Reference content — authoritative ground truth from ingested URLs + topic_dir = LOCAL_MEMSTORE_ROOT / "memory" / slug + if topic_dir.exists(): + for fpath in sorted(topic_dir.glob("*.txt")): + fc = fpath.read_text(encoding="utf-8") + if fc.strip(): + await cognee.add(f"[REFERENCE]\n{fc}", dataset_name=slug) + ref_docs += 1 + + print(f"[rebuild] {slug}: {ref_docs} reference docs") + + with flyte.group(f"rebuild:{slug}:cognify"): + await cognee.cognify(datasets=[slug], chunk_size=512) + print(f"[rebuild] cognify complete for {slug!r}") + + with flyte.group(f"rebuild:{slug}:upload"): + await _upload_dir(local_cognee, _topic_db_path(slug)) + + return json.dumps({"slug": slug, "ref_docs": ref_docs}) + + +# --------------------------------------------------------------------------- +# Sleep cycle — autonomous, scheduled every 6 hours +# --------------------------------------------------------------------------- + +@env.task( + retries=2, + timeout=timedelta(minutes=45), + report=True, +) +async def sleep_cycle() -> dict: + """Autonomous memory consolidation — fired every 6 hours by the app-level scheduler. + + This task is the heart of the sleep/wake architecture. It runs with no human + interaction; the app.py background thread calls flyte.run(sleep_cycle) every 6 hours. + (flyte.Trigger + flyte.Cron on tasks is not used: the Union cluster does not write + inputs.pb for triggered task executions, causing READ_FAILED on every trigger fire.) + + Pipeline (each phase visible as a flyte.group span in the Union UI timeline): + 1. Download latest state from shared object storage + 2. Auto-promote user/ staged proposals (validator is the only gate) + 3. Cluster related user/ memories by topic prefix + 4. Consolidate each cluster in parallel via flyte.map.aio + → each cluster runs as an isolated Flyte pod + → Claude merges related memories into coherent summaries + → cache="auto" on consolidate_cluster skips unchanged clusters + 5. cognee.cognify() — rebuild the full knowledge graph + 6. Upload updated state to shared object storage + 7. Stream final HTML summary to Union UI report panel + + Flyte features: + flyte.map.aio parallel pods, one per memory cluster + retries=2 transient failures auto-retried + report=True live HTML progress in Union UI + flyte.group() per-phase spans in execution timeline + """ + from agent import ( + archive_proposal, + classify_proposal_topic, + list_staged_proposals, + promote_proposal, + validate_proposal, + ) + + ts = _utc_now() + summary: dict = { + "trigger_time": ts, + "promoted": 0, + "rejected": 0, + "clusters_found": 0, + "clusters_consolidated": 0, + "memories_merged": 0, + "topics_rebuilt": 0, + "cognify_ran": False, + "cognify_s": 0.0, + "errors": [], + "phase": "starting", + } + + # Phase 1: Download latest state + with flyte.group("sleep:download"): + summary["phase"] = "downloading" + await _emit_report(summary) + await _download_dir(Dir(path=SHARED_MEMSTORE_PATH), LOCAL_MEMSTORE_ROOT) + # No cognee_db download — each rebuild_topic_dataset pod handles its own topic DB + + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + store.ensure_layout() + + # Load topic index and api_key — used throughout the cycle + topic_index = load_topic_index(store) + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + changed_topics: set[str] = set() + + # Discover sessions — register "default" if none exist (first-time / migration) + sessions = list_sessions(store) + if not sessions: + register_session(store, "default", label="Default Session") + sessions = ["default"] + + # Phase 2: Auto-promote staged proposals for all sessions + with flyte.group("sleep:promote"): + summary["phase"] = "promoting" + await _emit_report(summary) + + for session in sessions: + staged = list_staged_proposals(store, limit=100, session=session) + staged = [p for p in staged if not _is_already_archived(store, p.id, session)] + user_proposals = [p for p in staged if p.target == "user"] + + for proposal in user_proposals: + decision = validate_proposal(store, proposal) + if decision.ok: + try: + promote_proposal( + store, proposal, + actor="sleep_cycle", + promotion_reason="auto-promoted by scheduled sleep cycle", + ) + archive_proposal( + store, proposal, + actor="sleep_cycle", decision="approved", + note="auto-promoted by sleep_cycle", + ) + summary["promoted"] += 1 + # Track topic for rebuild (reference datasets may need refreshing) + slug = proposal.topic_slug + if not slug: + slug = classify_proposal_topic( + proposal.content, proposal.source_question, topic_index, api_key + ) + if slug and slug in topic_index: + changed_topics.add(slug) + # Write session-scoped topic map entry + upsert_topic_map( + store, decision.normalized_path, slug, + topic_map_path=session_topic_map_path(session), + ) + except Exception as e: + summary["errors"].append(f"promote:{session}:{proposal.id[:8]}:{type(e).__name__}") + else: + archive_proposal( + store, proposal, + actor="sleep_cycle", decision="rejected", + note=decision.reason, + ) + summary["rejected"] += 1 + + # Phase 3: Cluster + consolidate memories across all sessions in parallel + with flyte.group("sleep:consolidate"): + summary["phase"] = "consolidating" + await _emit_report(summary) + + all_clusters = [] + for session in sessions: + session_clusters = _cluster_user_memories(store, session) + # Embed session in cluster dict so we know which session after consolidation + for c in session_clusters: + c["session"] = session + all_clusters.extend(session_clusters) + + clusters = all_clusters + summary["clusters_found"] = len(clusters) + + if clusters: + cluster_jsons = [json.dumps(c) for c in clusters] + consolidated: list[str] = [] + + # flyte.map.aio fans out consolidate_cluster as parallel Flyte pods. + # concurrency=3 caps simultaneous pods to avoid overwhelming the cluster. + async for result in flyte.map.aio( + consolidate_cluster, + cluster_jsons, + concurrency=3, + return_exceptions=True, + ): + if isinstance(result, Exception): + summary["errors"].append(f"consolidate:{type(result).__name__}:{result}") + else: + consolidated.append(result) + + for result_json in consolidated: + try: + result = json.loads(result_json) + merged_from = result.get("merged_from", []) + if merged_from: + store.write_text( + result["path"], + result["content"], + actor="sleep_cycle", + reason=f"consolidated {len(merged_from) + 1} memories", + op="consolidate", + ) + summary["clusters_consolidated"] += 1 + summary["memories_merged"] += len(merged_from) + # Resolve topic via the session-scoped topic map + result_session = result.get("session", "default") + topic_map = read_topic_map(store, topic_map_path=session_topic_map_path(result_session)) + slug = topic_map.get(result["path"]) + if slug and slug in topic_index: + changed_topics.add(slug) + except Exception as e: + summary["errors"].append(f"write_consolidated:{type(e).__name__}") + + # Phase 4: Early upload — push promoted + consolidated memstore to shared storage + # so that rebuild_topic_dataset pods download the latest state (not the pre-sleep snapshot). + with flyte.group("sleep:early_upload"): + summary["phase"] = "uploading_memstore" + await _emit_report(summary) + await _preserve_newest_preferences_before_upload() + await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) + + # Phase 5: Fan out per-topic Cognee rebuild as parallel Flyte pods. + # Each rebuild_topic_dataset pod downloads the fresh memstore + its own topic DB, + # tags content as [REFERENCE] or [USER_MEMORY], runs cognify+memify, and + # uploads its updated topic DB. Pods are isolated — no shared-DB write conflicts. + with flyte.group("sleep:cognify"): + summary["phase"] = "cognifying" + await _emit_report(summary) + t0 = time.perf_counter() + + rebuild_results: list[str] = [] + if changed_topics: + rebuild_jsons = [json.dumps({"slug": slug}) for slug in sorted(changed_topics)] + + async for result in flyte.map.aio( + rebuild_topic_dataset, + rebuild_jsons, + concurrency=3, + return_exceptions=True, + ): + if isinstance(result, Exception): + summary["errors"].append(f"rebuild:{type(result).__name__}:{result}") + else: + rebuild_results.append(result) + + summary["topics_rebuilt"] = len(rebuild_results) + summary["cognify_ran"] = bool(rebuild_results) + summary["cognify_s"] = round(time.perf_counter() - t0, 2) + + # Phase 6: Final memstore upload — catches any preference changes that arrived + # while rebuild pods were running (preference-race guard). + with flyte.group("sleep:upload"): + summary["phase"] = "uploading" + await _emit_report(summary) + + await _preserve_newest_preferences_before_upload() + await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) + # cognee_db not uploaded here — each rebuild_topic_dataset pod handled its own + + summary["phase"] = "complete" + await flyte.report.replace.aio(_build_sleep_report(summary), do_flush=True) + await flyte.report.flush.aio() + + return summary + + +# --------------------------------------------------------------------------- +# Chat summary — on-demand +# --------------------------------------------------------------------------- + +@env.task(retries=1, timeout=timedelta(minutes=3)) +async def summarize_chat_session( + session_id: str, + session: str = "default", + model: str = DEFAULT_MODEL, + max_lines: int = 200, +) -> str: + """Summarize a chat session transcript into a short running summary. + + Reads: + user/sessions//chat//transcript.jsonl + Writes: + user/sessions//chat//summary.txt + + This enables durable conversation continuity without stuffing the entire + transcript into every prompt. + """ + with flyte.group("summary:download"): + await _download_dir(Dir(path=SHARED_MEMSTORE_PATH), LOCAL_MEMSTORE_ROOT) + + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + store.ensure_layout() + + transcript_path = f"user/sessions/{session}/chat/{session_id}/transcript.jsonl" + summary_path = f"user/sessions/{session}/chat/{session_id}/summary.txt" + + transcript = store.read_text(transcript_path, default="").strip() + if not transcript: + return "" + + # Build a compact plain-text view for the summarizer. + lines = transcript.splitlines()[-max_lines:] + rendered: list[str] = [] + for ln in lines: + try: + obj = json.loads(ln) + role = str(obj.get("role", "")) + content = str(obj.get("content", "")).strip() + if role in ("user", "assistant") and content: + rendered.append(f"{role}: {content}") + except Exception: + continue + + excerpt = "\n".join(rendered).strip() + if not excerpt: + return "" + + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + if not api_key: + return "" + + from anthropic import Anthropic + + prev = store.read_text(summary_path, default="").strip() + + client = Anthropic(api_key=api_key, timeout=60.0) + msg = client.messages.create( + model=model, + max_tokens=500, + system=( + "You maintain a running summary of a chat between a user and an assistant. " + "Update the summary to reflect any new facts, decisions, and open questions. " + "Keep it short and concrete. Return only the summary text." + ), + messages=[{ + "role": "user", + "content": ( + f"Current summary (may be empty):\n{prev or '<>'}\n\n" + f"Recent transcript excerpt:\n{excerpt}\n" + ), + }], + temperature=0, + ) + summary = msg.content[0].text.strip() + + expected = store.current_sha(summary_path) or None + store.write_text( + summary_path, + (summary + "\n") if summary else "", + actor="summarize_chat_session", + reason="chat-summary", + expected_sha=expected, + op="summarize", + extra_audit={"chat_session_id": session_id, "session": session}, + ) + + with flyte.group("summary:upload"): + await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) + + return summary + + +# --------------------------------------------------------------------------- +# Wake cycle — on-demand per question +# --------------------------------------------------------------------------- + +@env.task(retries=1, timeout=timedelta(minutes=2)) +async def wake_cycle( + question: str, + session: str = "default", + model: str = DEFAULT_MODEL, + search_timeout_s: float = 60.0, + answer_timeout_s: float = 30.0, +) -> tuple[str, dict]: + """Answer a question using the latest consolidated memory + Cognee retrieval. + + Downloads the current shared state, runs a Cognee semantic search, assembles + a memory-augmented system prompt (preferences + reference docs + retrieved + context), and calls Claude for the answer. + + retries=1 handles transient Cognee or Anthropic API failures. + Each call always reads the latest state — sleep cycle consolidations are + immediately visible to the next wake call. + """ + with flyte.group("wake:download"): + await _download_dir(Dir(path=SHARED_MEMSTORE_PATH), LOCAL_MEMSTORE_ROOT) + # Per-topic cognee DBs downloaded after routing — only fetch what's needed + + store = MemoryStore(LOCAL_MEMSTORE_ROOT) + + with flyte.group("wake:retrieve"): + prefs_obj = store.read_json("user/preferences.json", default=None) + prefs = ( + "\n".join(f"{k}={v}" for k, v in sorted(prefs_obj.items())) + if isinstance(prefs_obj, dict) and prefs_obj + else store.read_text("user/preferences.txt", default="") + ) + + api_key = os.environ.get("ANTHROPIC_API_KEY", "") + topic_index = load_topic_index(store) + target_slugs = _route_query_to_topics(question, topic_index, api_key) + if not target_slugs: + target_slugs = list(topic_index.keys()) + print(f"[wake] routing → {target_slugs if target_slugs else 'no topics available'}") + + # Inject raw session user memories as additional context + session_mem_dir = LOCAL_MEMSTORE_ROOT / "user" / "sessions" / session / "memories" + user_memory_parts: list[str] = [] + if session_mem_dir.exists(): + for fpath in sorted(session_mem_dir.glob("*.txt")): + if fpath.name.startswith("_"): + continue + try: + uc = fpath.read_text(encoding="utf-8") + if uc.strip(): + user_memory_parts.append(f"[USER_MEMORY]\n{uc[:5000]}") + except Exception: + pass + print(f"[wake] session {session!r}: {len(user_memory_parts)} user memory file(s)") + + t0 = time.perf_counter() + all_results: list = [] + for slug in target_slugs: + local_cognee = Path(f"/tmp/cognee_db_{slug}") + try: + await _download_dir(Dir(path=_topic_db_path(slug)), local_cognee) + except Exception as e: + print(f"[wake] skip {slug!r}: DB not found ({type(e).__name__})") + continue + _setup_cognee_env(local_cognee) + import cognee + _configure_cognee_runtime(cognee, local_cognee) + try: + results = await asyncio.wait_for( + cognee.search( + query_text=question, + datasets=[slug], + ), + timeout=search_timeout_s, + ) + all_results.extend(results or []) + print(f"[wake] {slug!r}: {len(results or [])} result(s)") + except asyncio.TimeoutError: + print(f"[wake] {slug!r}: search timed out after {search_timeout_s}s") + except Exception as e: + print(f"[wake] {slug!r}: search error: {type(e).__name__}: {e}") + + def _extract_result_text(r) -> str: + # SearchResult.search_result holds the actual content + sr = getattr(r, "search_result", None) + if sr is not None: + if isinstance(sr, str): + return sr.strip() + if isinstance(sr, dict): + return str(sr).strip() + return str(sr).strip() + # Fallback: check other common attrs then raw repr + for attr in ("text", "content", "payload", "value"): + val = getattr(r, attr, None) + if isinstance(val, str) and val.strip(): + return val + s = str(r) + return "" if (s.startswith("<") and s.endswith(">")) else s + + context_parts = [_extract_result_text(r) for r in all_results[:5]] + cognee_ctx = "\n\n".join(p for p in context_parts if p) + user_mem_ctx = "\n\n".join(user_memory_parts[:10]) + context = "\n\n".join(p for p in (cognee_ctx, user_mem_ctx) if p) + print(f"[wake] context: {len(context)} chars from {len(all_results)} result(s)") + if context: + print(f"[wake] context preview: {context[:300]!r}") + retrieve_s = time.perf_counter() - t0 + + with flyte.group("wake:answer"): + from anthropic import Anthropic, APITimeoutError + + if not api_key: + return "[error] ANTHROPIC_API_KEY not set", {} + + system = ( + "You are an assistant with access to two types of retrieved memory:\n" + "- [REFERENCE]: authoritative ground truth from ingested documents. Treat as fact.\n" + "- [USER_MEMORY]: user-specific overrides. These take precedence over [REFERENCE] " + "content for how this user wants things done.\n" + "If context is empty, answer from general knowledge and say so.\n" + "Treat the user's latest messages as authoritative for newly introduced facts.\n" + "Treat [preferences] as requirements.\n" + "- If preferences include name=, address the user by that name in every response.\n" + "- If preferences include tone/format, comply.\n" + "- For other preference keys, interpret them as user directives and follow them as best you can.\n\n" + f"[preferences]\n{prefs or '(none)'}\n\n" + f"[retrieved]\n{context or '<>'}\n" + ) + client = Anthropic(api_key=api_key, timeout=answer_timeout_s) + t0 = time.perf_counter() + try: + msg = client.messages.create( + model=model, + max_tokens=900, + system=system, + messages=[{"role": "user", "content": question}], + temperature=0, + ) + answer = msg.content[0].text + except APITimeoutError: + answer = f"[error] timed out after {answer_timeout_s}s" + answer_s = time.perf_counter() - t0 + + return answer, { + "retrieve_s": round(retrieve_s, 2), + "answer_s": round(answer_s, 2), + "ctx_chars": len(context), + } + + +# --------------------------------------------------------------------------- +# HTML report for sleep_cycle (streamed live to Union UI) +# --------------------------------------------------------------------------- + +async def _emit_report(summary: dict) -> None: + await flyte.report.replace.aio(_build_sleep_report(summary), do_flush=True) + + +def _build_sleep_report(summary: dict) -> str: + phase = summary.get("phase", "starting") + phases = ["downloading", "promoting", "consolidating", "uploading_memstore", "cognifying", "uploading", "complete"] + + steps_html = "" + for step in phases: + done = phases.index(step) < phases.index(phase) or phase == "complete" + current = step == phase and phase != "complete" + if done: + icon, color = "✅", "#2ecc71" + elif current: + icon, color = "⏳", "#f39c12" + else: + icon, color = "○", "#aaa" + steps_html += f'
{icon} {step.capitalize()}
' + + errors_html = "" + if summary.get("errors"): + items = "".join(f"
  • {e}
    • " for e in summary["errors"]) + errors_html = f"

    Errors ({len(summary['errors'])})

      {items}
    " + + cognify_val = ( + f'{summary.get("cognify_s", 0):.1f}s' + if summary.get("cognify_ran") else "—" + ) + + stats = [ + (summary["promoted"], "proposals auto-promoted"), + (summary["rejected"], "proposals rejected"), + (summary.get("clusters_found", 0), "memory clusters found"), + (summary["clusters_consolidated"], "clusters consolidated"), + (summary["memories_merged"], "memories merged"), + (summary.get("topics_rebuilt", 0), "topic datasets rebuilt"), + (cognify_val, "cognify+memify fired"), + ] + stats_html = "".join( + f'
    {v}
    {l}
    ' + for v, l in stats + ) + + status = "✅ Complete" if phase == "complete" else "⏳ Running…" + return f""" + + + + + +

    🌙 Sleep Cycle — {status}

    +

    Triggered: {summary.get("trigger_time", "—")}

    +
    {steps_html}
    +
    {stats_html}
    +{errors_html} +
    +

    + Cognee Memory Store · Flyte sleep/wake architecture
    + Schedule: every 6 hours, managed by app.py background thread
    + Deploy: python app.py +

    + +""" + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +def main() -> None: + if os.environ.get("SELF_CHECK") == "true": + from memory_store import _self_check + _self_check() + print("workflow self-check: ok") + return + + print( + "Primary entrypoint: python app.py\n" + "To register the 6-hour sleep schedule on Union: python workflow.py --deploy\n" + ) + + +if __name__ == "__main__": + import sys + + try: + flyte.init_from_config() + except Exception: + pass + + if "--deploy" in sys.argv: + flyte.deploy(env) + print("Sleep cycle trigger registered — fires every 6 hours.") + else: + main() \ No newline at end of file