From ef072543d4e4673f6d04ffb7d0f32931789bac4a Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Wed, 13 May 2026 08:54:51 -0700 Subject: [PATCH 1/8] init work Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- CLAUDE.md | 108 ++ v2/tutorials/cognee_memory_store/README.md | 146 ++ v2/tutorials/cognee_memory_store/agent.py | 261 +++ v2/tutorials/cognee_memory_store/app.py | 1354 ++++++++++++++++ .../cognee_memory_store/memory_store.py | 465 ++++++ v2/tutorials/cognee_memory_store/workflow.py | 1401 +++++++++++++++++ 6 files changed, 3735 insertions(+) create mode 100644 CLAUDE.md create mode 100644 v2/tutorials/cognee_memory_store/README.md create mode 100644 v2/tutorials/cognee_memory_store/agent.py create mode 100644 v2/tutorials/cognee_memory_store/app.py create mode 100644 v2/tutorials/cognee_memory_store/memory_store.py create mode 100644 v2/tutorials/cognee_memory_store/workflow.py diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..7998362d --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,108 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Commands + +```bash +# Setup +make setup-venv # creates ~/.venv with uv + Python 3.12 +source ~/.venv/bin/activate +make update-flyte # installs latest flyte pre-release into active venv + +# Testing (venv must be active) +make test-preview # discover scripts without running +make test FILE=v2/tutorials/foo/main.py # run one script on Union cloud +make test-local FILE=v2/tutorials/foo/main.py # run one script locally (isolated venv per script) +make test FILTER=cognee # run all scripts whose path matches "cognee" +make test-local VERBOSE=vv FILE=... # with flyte verbosity (-v / -vv / -vvv) +make clean # remove test/reports/, test/logs/, test/venvs/ + +# Direct runner (same flags) +python3 test/test_runner.py --preview +python3 test/test_runner.py --local --file "v2/tutorials/foo/main.py" +``` + +Cloud tests need `FLYTE_CLIENT_SECRET` set; they target `playground.canary.unionai.cloud` (project `docs-examples`, domain `development`) as configured in `test/config.flyte.yaml`. + +## Repository layout + +``` +v2/ Modern Flyte 2.x examples (only these are tested) + tutorials/ Full end-to-end tutorial projects (multi-file) + user-guide/ Short single-file illustrative snippets +v1/ Legacy Flyte 1.x examples (not tested) +_blogs/ Blog-post companion code +test/ Test runner + CI config +``` + +## Flyte 2.x script conventions + +Every runnable script uses **PEP 723 inline metadata** at the top — this is how `uv run` discovers dependencies and how the test runner installs them per-script in isolation: + +```python +# /// script +# requires-python = "==3.13" +# dependencies = [ +# "flyte==2.1.5", +# "some-lib>=1.0", +# ] +# main = "main" ← entrypoint task name +# /// +``` + +### Core Flyte 2.x patterns + +**TaskEnvironment** groups config shared across tasks (image, secrets, resources): +```python +env = flyte.TaskEnvironment( + name="my-env", + image=flyte.Image.from_uv_script(__file__, name="my-image"), + secrets=[flyte.Secret(key="my-secret", as_env_var="MY_ENV_VAR")], + resources=flyte.Resources(cpu=2, memory="4Gi"), +) + +@env.task +async def my_task(...): ... +``` + +**`flyte.Image.from_uv_script(__file__)`** builds a container image from the PEP 723 deps automatically. Additional source files needed in the container are added with `.with_source_file(local_path, container_path)`. + +**`flyte.io.Dir`** is used to hand off directory state between tasks running in isolated pods — download with `await d.download(local_path=...)`, upload with `await Dir.from_local(local_path, remote_destination=...)`. + +**`flyte.map.aio(task_fn, inputs, concurrency=N)`** fans out a task as parallel pods and async-iterates results. + +**Scheduling**: `flyte.Trigger` + `flyte.Cron` registers a schedule on the cluster without needing a LaunchPlan. + +**`flyte.group("label")`** creates a named span visible in the Union UI execution timeline. + +**`report=True`** on a task enables live HTML streaming to the Union UI dashboard. + +**`cache="auto"`** keys caching on task inputs + source code hash — useful for idempotent subtasks. + +### Deployment + +```bash +# Register schedules and deploy apps in one shot +python workflow.py --deploy +python app.py # also calls flyte.deploy() then flyte.serve() +``` + +## Active tutorial: cognee_memory_store + +The most complex active tutorial. It implements a sleep/wake memory architecture using Cognee (knowledge graph) + Flyte. + +**Files:** +- `memory_store.py` — audited, versioned file-based memory store synced via `flyte.io.Dir`. Prefixes enforce access: `reference/` is read-only, `user/` is read-write. +- `agent.py` — staged proposal pipeline: untrusted writes land in `staging/inbox/`, validated, then promoted to `user/`. +- `workflow.py` — all Flyte tasks: `init_reference`, `ingest_url`, `consolidate_cluster`, `sleep_cycle`, `wake_cycle`, `summarize_chat_session`. +- `app.py` — Streamlit UI served via `flyte.app.AppEnvironment`. + +**Sleep/wake cycle:** +- **Sleep** (every 6h via Cron): download state → promote staged proposals → cluster + consolidate `user/` memories via `flyte.map.aio` → full graph rebuild (per-topic `empty_dataset` + `cognify`) → upload. +- **Wake** (per question): download state → classify question to topic slugs → `cognee.search(datasets=[slugs])` → Claude answer. +- **`ingest_url`**: scrapes URL via Jina Reader (`r.jina.ai/`) with direct-HTTP fallback, stores in `memory/topic_/`, `cognify` + `memify` for enrichment. + +**Shared object storage paths:** +- `cognee-memory-store/memstore` → memstore files +- `cognee-memory-store/cognee_db` → Cognee SQLite state \ No newline at end of file diff --git a/v2/tutorials/cognee_memory_store/README.md b/v2/tutorials/cognee_memory_store/README.md new file mode 100644 index 00000000..bb53c464 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/README.md @@ -0,0 +1,146 @@ +# Cognee + Flyte “Memory Store” + +This tutorial shows how to build **excellent, durable memory** for an agent by combining: + +- **Cognee** → semantic memory (RAG / KG / embeddings): *retrieve relevant context* +- **Flyte v2** → orchestration + observability: *promote, cognify, persist, time, cache* +- **A file-based “memory store”** (this tutorial) → **auditable, versioned, access-controlled** durable memory + +It takes inspiration from Claude’s Managed Agents “memory stores” concepts: +- many small memories (files) +- multiple stores / mounts with different access rules +- immutable version history (audit trail) +- safe edits (optimistic concurrency) +- staging vs trusted memory (prompt-injection defense) + +## Why this design +If you rely only on semantic search (embeddings), it’s hard to enforce: +- “don’t overwrite shared memory accidentally” +- “don’t let untrusted inputs write persistent instructions” +- “show me exactly what changed and when” + +So we add a **curated memory layer** that’s transparent and governable. + +## Architecture + +``` +User message + | + | (1) Retrieve + | - read user/ preferences + | - run cognee.search(query) + v +LLM answer (fast path runs inside app container) + | + | (2) Optional: stage proposals (untrusted) + v +staging/inbox/*.json (NOT trusted) + | + | (3) Flyte task: review_and_promote + | - validate (anti-poisoning rules) + | - write into user/ or shared/ + | - version snapshot + audit log + | - cognee.add(promoted memory) + v +Trusted memory (user/, shared/) + audit/ + versions/ + | + | (4) Flyte task: cognify_if_needed (cached) + v +Cognee KG / indices updated +``` + +## Memory store layout +Everything lives under a directory that is synced via `flyte.io.Dir` to object storage. + +- `reference/` **read-only** (curated docs/conventions) +- `user/` **read-write** (per-user preferences, notes) +- `shared/` **promotion-only** (team/org conventions; requires explicit approval) +- `staging/inbox/` **untrusted proposals** +- `staging/archive/` archived proposals + decisions +- `versions/` immutable snapshots of every write +- `audit/log.jsonl` append-only log of memory mutations +- `meta/` current sha256 + last update info per memory + +### Access modes (Claude-inspired) +- `reference/` is **read-only** (writes rejected) +- `shared/` is **write-behind-review** (direct writes rejected; only allowed via promotion) + +### Optimistic concurrency (safe edits) +Writes can include an `expected_sha256` precondition. If the on-disk sha does not match, +we fail fast to avoid clobbering another writer. + +## Files +- `memory_store.py` — the memory store implementation (audit/versioning/access/concurrency) +- `agent.py` — proposal schema + staging + validation + promotion helpers +- `workflow.py` — Flyte tasks (init, answer, promote, cognify) +- `app.py` — Streamlit app (chat + memory sidebar; triggers Flyte tasks) + +## Run (Union / Flyte App) + +### 1) Ensure Anthropic key exists +This example expects a Union secret named `internal-anthropic-api-key` injected as `ANTHROPIC_API_KEY`. + +### 2) Configure GHCR +This tutorial always publishes images to **GHCR**. + +Default registry is `ghcr.io/flyteorg`. If you want to use your own org/user, set: + +```bash +export AI_MEMORY_STORE_IMAGE_REGISTRY="ghcr.io/" +# Optional: Union secret key that contains registry credentials (for private registries) +export AI_MEMORY_STORE_IMAGE_REGISTRY_SECRET="" +``` + +### 3) Serve the app + +```bash +# Default behavior prefers local image builds if Docker is available. +# To force remote builds: +# export AI_MEMORY_STORE_IMAGE_BUILDER=remote +uv run v2/tutorials/cognee_memory_store/app.py +``` + +Open the printed URL. + +### 3) Try the workflow +1. Ask a question. +2. Change preferences either: + - manually in the **Preferences** panel, or + - by saying something like “please be concise” / “use markdown” / “use my name Adil every time you answer”. + You’ll get a **preference approval popup**. +3. Confirm by clicking **Save preference** in the popup (or use **Save preferences** in the sidebar). +4. (Optional) Use **Advanced: stage raw proposal** + **Run promotion task** to exercise staged→promoted memory. +5. Click **Run cognify task**. +6. Ask a similar question again and observe retrieval improving. + +### Debug +Set: + +```bash +export AI_MEMORY_STORE_DEBUG=1 +export AI_MEMORY_STORE_MODEL=claude-haiku-4-5-20251001 +``` + +The chat will print retrieval/answer timings when debug is enabled. + +## Run (local dev) +Local mode does not require Flyte/Union credentials; persistence is local-only. + +```bash +streamlit run v2/tutorials/cognee_memory_store/app.py -- --server +``` + +You can still stage proposals and inspect audit/versioning locally, but promotion/cognify +Flyte tasks require remote object storage. + +## What to look for +- **Auditability:** expand “Audit log (tail)” and see every stage/promote/archive event. +- **Versioning:** each promoted write creates a `versions//*` immutable snapshot. +- **Safety:** staged proposals are not trusted until validated and promoted. +- **Separation of concerns:** Cognee retrieves; the memory store governs what is persisted. + +## Notes / extensions +- Add an LLM-based “memory gate” validator (still stage first; only promote after approval). +- Split into multiple stores (per-user store vs shared org store) the same way Claude supports + multiple stores with different access rules. +- Add redaction tools over `audit/` + `versions/` for compliance workflows. diff --git a/v2/tutorials/cognee_memory_store/agent.py b/v2/tutorials/cognee_memory_store/agent.py new file mode 100644 index 00000000..dd6910b2 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/agent.py @@ -0,0 +1,261 @@ +"""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 + + +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 + + +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) -> str: + return f"staging/inbox/{proposal_id}.json" + + +def proposal_archive_path(proposal_id: str, decision: str) -> str: + safe = re.sub(r"[^a-zA-Z0-9_-]", "_", decision)[:24] or "unknown" + return f"staging/archive/{safe}/{proposal_id}.json" + + +def stage_proposal(store: MemoryStore, proposal: MemoryWriteProposal) -> MemoryMeta: + """Write a proposal into staging/inbox (untrusted).""" + return store.write_json( + proposal_inbox_path(proposal.id), + proposal.model_dump(), + actor=proposal.author, + reason=proposal.reason or "stage-proposal", + op="stage", + ) + + +def list_staged_proposals(store: MemoryStore, limit: int = 50) -> list[MemoryWriteProposal]: + paths = store.list_paths("staging/inbox") + paths = 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: + if not proposal.path.startswith("user/"): + return "user/" + proposal.path.lstrip("/") + return proposal.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) + 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..038fc6e2 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/app.py @@ -0,0 +1,1354 @@ +# /// 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, + _setup_cognee_env, + _route_query_to_topics, + _topic_db_path, + init_memory_store, + ingest_url, + sleep_cycle, + summarize_chat_session, +) +from memory_store import ( + MemoryStore, + TOPIC_INDEX_PATH, + load_topic_index, + 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, 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()) + if GENERAL_TOPIC_SLUG not in target_slugs: + target_slugs.append(GENERAL_TOPIC_SLUG) + + 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) + 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 + return "\n".join(str(r) for r in all_results[:5]) + + 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 = 6000 + 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")): + try: + text = fpath.read_text(encoding="utf-8", errors="replace") + except Exception: + continue + if not text.strip(): + continue + raw_parts.append(f"--- {slug}/{fpath.name} ---\n{text[:per_slug_cap]}") + if sum(len(p) for p in raw_parts) > 20_000: + break + if sum(len(p) for p in raw_parts) > 20_000: + break + + parts = [p for p in (cognee_ctx, "\n\n".join(raw_parts)) 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) -> str: + return f"user/chat_sessions/{session_id}/transcript.jsonl" + + +def _chat_summary_path(session_id: str) -> str: + return f"user/chat_sessions/{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) -> 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) + 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 = { + "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 _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"} + try: + run.sync() + except Exception: + pass + try: + phase = getattr(run.phase, "name", str(run.phase)) + except Exception: + return False, "RUNNING", [] + 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", "") + _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, + ) + stage_proposal(store, prop) + asyncio.run(_upload_memstore()) + st.toast("Memory staged — 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_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. + phase = ( + getattr(getattr(sleep_run, "phase", None), "name", "") + if sleep_run else "" + ) + + 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}") + for p in topic_paths[:5]: + 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") + + user_paths = store.list_paths("user") + + with st.expander(f"Promoted memories ({len(user_paths)})", expanded=False): + 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/archive/{decision}/{proposal_id}.json"): + return True + return False + + staged = list_staged_proposals(store) + 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 — user/ (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 user/ proposals.") + 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", "") + _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, + ) + stage_proposal(store, prop) + asyncio.run(_upload_memstore()) + st.success("Staged — 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_sleep_section() + st.divider() + _render_knowledge_seeding() + st.divider() + + # Chat continuity (durable transcript + optional summary) + session_id = st.session_state.get("chat_session_id", "") + with st.expander("💬 Chat continuity", expanded=False): + st.caption(f"Session: `{session_id}`") + if session_id: + st.caption(f"Transcript: `{_chat_transcript_path(session_id)}`") + st.caption(f"Summary: `{_chat_summary_path(session_id)}`") + current_summary = store.read_text(_chat_summary_path(session_id), 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) + 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}: {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) + + st.title("🧠 Cognee + Flyte Memory Store") + st.caption( + "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 + + 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) + 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), 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", + ) + 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..720ec3b4 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/memory_store.py @@ -0,0 +1,465 @@ +"""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 +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" + + +def read_topic_map(store: "MemoryStore") -> dict: + """Return {user_rel_path: topic_slug} from user/_topic_map.json.""" + return store.read_json(TOPIC_MAP_PATH, default={}) + + +def upsert_topic_map(store: "MemoryStore", rel_path: str, slug: Optional[str]) -> None: + """Set or clear the topic association for a promoted user/ file. + + Writes directly to the filesystem, bypassing MemoryStore access control, + because _topic_map.json is machine-managed system metadata. + """ + m = read_topic_map(store) + 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..4a7fa6e3 --- /dev/null +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -0,0 +1,1401 @@ +# /// 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, + TOPIC_INDEX_PATH, + TOPIC_MAP_PATH, + load_topic_index, + 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 _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. + """ + 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}") + + for link in _extract_links(html, clean): + 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) -> 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/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) -> list[dict]: + """Group user/ text 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: user/notes_flyte.txt + user/notes_tasks.txt → cluster "notes". + """ + paths = [p for p in store.list_paths("user") 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", + ) + + 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 = 4_000 # keeps cognify's extraction JSON under the 8192-token non-streaming 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 + + 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)") + + 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) ...") + await cognee.cognify(datasets=[slug], chunk_size=512) + print(f"[ingest] cognify complete for {slug!r}") + + with flyte.group("ingest:upload"): + print("[ingest] Uploading (pod stays alive for background cognee tasks) ...") + memstore_dir = await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) + cognee_dir = await _upload_dir(local_cognee, _topic_db_path(slug)) + 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) + topic_map = read_topic_map(store) + ref_docs = 0 + user_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 + + # User preference content — personal overrides, classified to this topic + user_files = [p for p, t in topic_map.items() if t == slug and p.endswith(".txt")] + for rel_path in user_files: + uc = store.read_text(rel_path) + if uc.strip(): + await cognee.add(f"[USER_MEMORY]\n{uc}", dataset_name=slug) + user_docs += 1 + + print(f"[rebuild] {slug}: {ref_docs} reference docs, {user_docs} user preference 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, "user_docs": user_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() + + # Phase 2: Auto-promote user/ proposals (no human needed) + with flyte.group("sleep:promote"): + summary["phase"] = "promoting" + await _emit_report(summary) + staged = list_staged_proposals(store, limit=100) + + # Skip proposals already processed in a prior sleep cycle + staged = [p for p in staged if not _is_already_archived(store, p.id)] + 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 + # Use the proposal's classified topic slug, or fall back to + # classifying now for proposals staged before this feature + slug = proposal.topic_slug + if not slug: + slug = classify_proposal_topic( + proposal.content, proposal.source_question, topic_index, api_key + ) + effective_slug = slug if (slug and slug in topic_index) else GENERAL_TOPIC_SLUG + upsert_topic_map(store, decision.normalized_path, effective_slug) + changed_topics.add(effective_slug) + except Exception as e: + summary["errors"].append(f"promote:{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 user/ memories in parallel + with flyte.group("sleep:consolidate"): + summary["phase"] = "consolidating" + await _emit_report(summary) + clusters = _cluster_user_memories(store) + 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 the topic for this consolidated path via the topic_map + topic_map = read_topic_map(store) + 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, + model: str = DEFAULT_MODEL, + max_lines: int = 200, +) -> str: + """Summarize a chat session transcript into a short running summary. + + Reads: + user/chat_sessions//transcript.jsonl + Writes: + user/chat_sessions//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/chat_sessions/{session_id}/transcript.jsonl" + summary_path = f"user/chat_sessions/{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}, + ) + + 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, + 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()) + if GENERAL_TOPIC_SLUG not in target_slugs: + target_slugs.append(GENERAL_TOPIC_SLUG) + print(f"[wake] routing → {target_slugs if target_slugs else 'no topics available'}") + + 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]] + context = "\n\n".join(p for p in context_parts 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 From 83fe37e461857c313409095e713e29f04ce62395 Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Wed, 13 May 2026 10:12:54 -0700 Subject: [PATCH 2/8] incr max char Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- v2/tutorials/cognee_memory_store/workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/v2/tutorials/cognee_memory_store/workflow.py b/v2/tutorials/cognee_memory_store/workflow.py index 4a7fa6e3..eb918839 100644 --- a/v2/tutorials/cognee_memory_store/workflow.py +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -624,7 +624,7 @@ async def ingest_url(url: str, max_pages: int = 10) -> tuple[Dir, Dir]: store.ensure_layout() api_key = os.environ.get("ANTHROPIC_API_KEY", "") - MAX_CHARS = 4_000 # keeps cognify's extraction JSON under the 8192-token non-streaming ceiling + 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})") From ed8728ff72d0fde0d989a08c6c9429143e16c172 Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Mon, 25 May 2026 20:42:12 -0700 Subject: [PATCH 3/8] wip Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- v2/tutorials/cognee_memory_store/app.py | 48 ++++++--- v2/tutorials/cognee_memory_store/workflow.py | 101 ++++++++++++++++++- 2 files changed, 135 insertions(+), 14 deletions(-) diff --git a/v2/tutorials/cognee_memory_store/app.py b/v2/tutorials/cognee_memory_store/app.py index 038fc6e2..798931c4 100644 --- a/v2/tutorials/cognee_memory_store/app.py +++ b/v2/tutorials/cognee_memory_store/app.py @@ -56,6 +56,7 @@ GENERAL_TOPIC_SLUG, _setup_cognee_env, _route_query_to_topics, + _strip_jina_header, _topic_db_path, init_memory_store, ingest_url, @@ -418,22 +419,34 @@ async def _search() -> str: # short. Cap per-slug bytes to keep the prompt bounded. raw_parts: list[str] = [] if len(cognee_ctx) < 400: - per_slug_cap = 6000 + 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")): + 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) > 20_000: + if sum(len(p) for p in raw_parts) > 40_000: break - if sum(len(p) for p in raw_parts) > 20_000: + if sum(len(p) for p in raw_parts) > 40_000: break parts = [p for p in (cognee_ctx, "\n\n".join(raw_parts)) if p] @@ -656,14 +669,23 @@ def _dialog() -> None: 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"} - try: - run.sync() - except Exception: - pass + run.sync() # let exceptions propagate so callers can surface them via st.warning try: phase = getattr(run.phase, "name", str(run.phase)) except Exception: - return False, "RUNNING", [] + # 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": @@ -761,10 +783,10 @@ def _render_sleep_section() -> None: if sleep_run or sleep_url: # IMPORTANT: Do NOT auto-poll with run.sync() on every rerun. - phase = ( - getattr(getattr(sleep_run, "phase", None), "name", "") - if sleep_run else "" - ) + 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: diff --git a/v2/tutorials/cognee_memory_store/workflow.py b/v2/tutorials/cognee_memory_store/workflow.py index eb918839..c1352ee8 100644 --- a/v2/tutorials/cognee_memory_store/workflow.py +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -252,12 +252,85 @@ def _fetch_raw_html(url: str, max_bytes: int = 300_000) -> str: 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 < 3 prose words are nav chrome + cleaned = re.sub(r"!?\[[^\]]*\]\([^)]*\)", "", s) # remove links/images + cleaned = re.sub(r"`[^`]*`", "", cleaned) # remove inline code + return len(re.findall(r"\b[a-zA-Z]{3,}\b", cleaned)) < 3 + + lines = text.splitlines() + i = 0 + # Never skip more than the first third of the file (safety bound) + max_skip = min(len(lines), max(len(lines) // 3, 100)) + while i < max_skip 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 @@ -282,7 +355,32 @@ def _crawl_site(seed_url: str, max_pages: int = 50) -> list[str]: discovered.append(clean) print(f"[crawl] discovered ({len(discovered)}/{max_pages}): {clean}") - for link in _extract_links(html, 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 @@ -672,6 +770,7 @@ async def ingest_url(url: str, max_pages: int = 10) -> tuple[Dir, Dir]: print(f"[ingest] skip (empty): {page_url}") continue + text = _strip_jina_header(text) truncated = len(text) > MAX_CHARS text = text[:MAX_CHARS] content = ( From adbac79058a359f5e1da15fb3f6957420f31469b Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Mon, 25 May 2026 21:24:23 -0700 Subject: [PATCH 4/8] add: multi-session support Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- v2/tutorials/cognee_memory_store/agent.py | 48 +++-- v2/tutorials/cognee_memory_store/app.py | 166 ++++++++++++++--- .../cognee_memory_store/memory_store.py | 51 +++++- v2/tutorials/cognee_memory_store/workflow.py | 170 +++++++++++------- 4 files changed, 322 insertions(+), 113 deletions(-) diff --git a/v2/tutorials/cognee_memory_store/agent.py b/v2/tutorials/cognee_memory_store/agent.py index dd6910b2..43e0feaa 100644 --- a/v2/tutorials/cognee_memory_store/agent.py +++ b/v2/tutorials/cognee_memory_store/agent.py @@ -20,7 +20,14 @@ from pydantic import BaseModel, Field -from memory_store import AccessDenied, ConcurrencyError, MemoryMeta, MemoryStore +from memory_store import ( + AccessDenied, + ConcurrencyError, + MemoryMeta, + MemoryStore, + session_staging_inbox_prefix, + session_staging_archive_prefix, +) ProposalTarget = Literal["user"] @@ -59,6 +66,9 @@ class MemoryWriteProposal(BaseModel): # 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 @@ -112,19 +122,19 @@ def classify_proposal_topic( return None -def proposal_inbox_path(proposal_id: str) -> str: - return f"staging/inbox/{proposal_id}.json" +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) -> str: +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"staging/archive/{safe}/{proposal_id}.json" + return f"{session_staging_archive_prefix(session)}/{safe}/{proposal_id}.json" def stage_proposal(store: MemoryStore, proposal: MemoryWriteProposal) -> MemoryMeta: - """Write a proposal into staging/inbox (untrusted).""" + """Write a proposal into the session's staging inbox (untrusted).""" return store.write_json( - proposal_inbox_path(proposal.id), + proposal_inbox_path(proposal.id, proposal.session), proposal.model_dump(), actor=proposal.author, reason=proposal.reason or "stage-proposal", @@ -132,9 +142,14 @@ def stage_proposal(store: MemoryStore, proposal: MemoryWriteProposal) -> MemoryM ) -def list_staged_proposals(store: MemoryStore, limit: int = 50) -> list[MemoryWriteProposal]: - paths = store.list_paths("staging/inbox") - paths = paths[:limit] +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: @@ -150,9 +165,14 @@ def list_staged_proposals(store: MemoryStore, limit: int = 50) -> list[MemoryWri def _normalize_target_path(proposal: MemoryWriteProposal) -> str: - if not proposal.path.startswith("user/"): - return "user/" + proposal.path.lstrip("/") - return proposal.path + 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( @@ -244,7 +264,7 @@ def archive_proposal( 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) + archive_path = proposal_archive_path(proposal.id, decision, proposal.session) store.write_json( archive_path, { diff --git a/v2/tutorials/cognee_memory_store/app.py b/v2/tutorials/cognee_memory_store/app.py index 798931c4..e50cc846 100644 --- a/v2/tutorials/cognee_memory_store/app.py +++ b/v2/tutorials/cognee_memory_store/app.py @@ -65,8 +65,14 @@ ) 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, @@ -378,15 +384,13 @@ def _extract_proposal_from_message(user_message: str) -> dict | None: return None -def _retrieve_context(question: str, timeout_s: float = 15.0) -> str: +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()) - if GENERAL_TOPIC_SLUG not in target_slugs: - target_slugs.append(GENERAL_TOPIC_SLUG) async def _search() -> str: import cognee @@ -449,7 +453,21 @@ def _file_relevance(p: Path) -> int: if sum(len(p) for p in raw_parts) > 40_000: break - parts = [p for p in (cognee_ctx, "\n\n".join(raw_parts)) if p] + # 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) @@ -477,12 +495,12 @@ def _prefs_to_text(prefs: dict) -> str: return "\n".join(lines).strip() + ("\n" if lines else "") -def _chat_transcript_path(session_id: str) -> str: - return f"user/chat_sessions/{session_id}/transcript.jsonl" +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) -> str: - return f"user/chat_sessions/{session_id}/summary.txt" +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: @@ -500,13 +518,13 @@ def _build_llm_messages(history: list[dict], max_messages: int) -> list[dict]: return out -def _append_transcript(store: MemoryStore, session_id: str, entries: list[dict], *, actor: str, reason: str) -> None: +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) + path = _chat_transcript_path(session_id, session) existing = store.read_text(path, default="") lines = existing.splitlines() if existing.strip() else [] @@ -545,6 +563,8 @@ def _init_session() -> None: import streamlit as st defaults: dict = { + "active_session": "default", + "_last_active_session": "", "messages": [], "last_answer": "", "sleep_run": None, @@ -575,6 +595,34 @@ def _init_session() -> None: 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 @@ -740,6 +788,7 @@ def _render_proposal_card(msg_idx: int, proposal: dict, store: MemoryStore) -> N 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", ""), @@ -754,10 +803,11 @@ def _render_proposal_card(msg_idx: int, proposal: dict, store: MemoryStore) -> N 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("Memory staged — auto-promoted on next sleep cycle", icon="💡") + 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 @@ -773,6 +823,55 @@ def _render_proposal_card(msg_idx: int, proposal: dict, store: MemoryStore) -> N # 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 @@ -869,29 +968,37 @@ def _render_memory_viewer(store: MemoryStore) -> None: st.caption(f"`{p}` ({size_kb:.1f} KB)") st.code(content[:400] + ("…" if len(content) > 400 else ""), language="text") - user_paths = store.list_paths("user") + 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 ({len(user_paths)})", expanded=False): + 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) + 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 — user/ (pending {len(user_staged_pending)} · processed {len(user_staged_all) - len(user_staged_pending)})", + 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 user/ proposals.") + 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)") @@ -941,6 +1048,7 @@ def _render_preferences(store: MemoryStore) -> None: 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="", @@ -954,10 +1062,11 @@ def _render_preferences(store: MemoryStore) -> None: author=author, reason=reason, topic_slug=_topic_slug, + session=_active_sess, ) stage_proposal(store, prop) asyncio.run(_upload_memstore()) - st.success("Staged — auto-promoted on next sleep cycle") + st.success(f"Staged in session '{_active_sess}' — auto-promoted on next sleep cycle") st.rerun() @@ -1075,19 +1184,22 @@ def _render_sidebar(store: MemoryStore) -> None: 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"Session: `{session_id}`") + st.caption(f"Memory session: `{active_session}` · Chat ID: `{session_id}`") if session_id: - st.caption(f"Transcript: `{_chat_transcript_path(session_id)}`") - st.caption(f"Summary: `{_chat_summary_path(session_id)}`") - current_summary = store.read_text(_chat_summary_path(session_id), default="").strip() + 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") @@ -1134,12 +1246,12 @@ def _render_sidebar(store: MemoryStore) -> None: else: if st.button("Update summary (Flyte)", use_container_width=True): try: - run = flyte.run(summarize_chat_session, session_id=session_id) + 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}: {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: @@ -1231,8 +1343,10 @@ def main() -> None: 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." ) @@ -1253,6 +1367,7 @@ def main() -> None: 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}) @@ -1264,11 +1379,11 @@ def main() -> None: # Retrieve context then run answer + proposal detection in parallel t0 = time.perf_counter() - context = _retrieve_context(user_input) + 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), default="") if session_id else "" + 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" @@ -1319,6 +1434,7 @@ def main() -> None: ], actor="chat", reason="chat-transcript", + session=active_session, ) asyncio.run(_upload_memstore()) diff --git a/v2/tutorials/cognee_memory_store/memory_store.py b/v2/tutorials/cognee_memory_store/memory_store.py index 720ec3b4..cb0fddb9 100644 --- a/v2/tutorials/cognee_memory_store/memory_store.py +++ b/v2/tutorials/cognee_memory_store/memory_store.py @@ -20,6 +20,7 @@ import json import os import re +import time from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -55,24 +56,60 @@ def _utc_ts() -> str: 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") -> dict: - """Return {user_rel_path: topic_slug} from user/_topic_map.json.""" - return store.read_json(TOPIC_MAP_PATH, default={}) +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]) -> None: - """Set or clear the topic association for a promoted user/ file. +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) + 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 = 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") diff --git a/v2/tutorials/cognee_memory_store/workflow.py b/v2/tutorials/cognee_memory_store/workflow.py index c1352ee8..970486b1 100644 --- a/v2/tutorials/cognee_memory_store/workflow.py +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -73,9 +73,14 @@ 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, @@ -568,10 +573,13 @@ def _utc_now() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") -def _is_already_archived(store: MemoryStore, proposal_id: str) -> bool: +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/archive/{decision}/{proposal_id}.json"): + 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 @@ -631,15 +639,16 @@ async def _preserve_newest_preferences_before_upload() -> None: shutil.copy2(src_txt, dst_txt) -def _cluster_user_memories(store: MemoryStore) -> list[dict]: - """Group user/ text memories by topic prefix for parallel consolidation. +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: user/notes_flyte.txt + user/notes_tasks.txt → cluster "notes". + Example: notes_flyte.txt + notes_tasks.txt → cluster "notes". """ - paths = [p for p in store.list_paths("user") if not p.endswith(".json")] + 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 @@ -683,6 +692,7 @@ async def init_memory_store() -> Dir: "\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) @@ -901,9 +911,7 @@ async def rebuild_topic_dataset(rebuild_json: str) -> str: _configure_cognee_runtime(cognee, local_cognee) store = MemoryStore(LOCAL_MEMSTORE_ROOT) - topic_map = read_topic_map(store) ref_docs = 0 - user_docs = 0 with flyte.group(f"rebuild:{slug}:index"): datasets_list = await cognee.datasets.list_datasets() @@ -920,15 +928,7 @@ async def rebuild_topic_dataset(rebuild_json: str) -> str: await cognee.add(f"[REFERENCE]\n{fc}", dataset_name=slug) ref_docs += 1 - # User preference content — personal overrides, classified to this topic - user_files = [p for p, t in topic_map.items() if t == slug and p.endswith(".txt")] - for rel_path in user_files: - uc = store.read_text(rel_path) - if uc.strip(): - await cognee.add(f"[USER_MEMORY]\n{uc}", dataset_name=slug) - user_docs += 1 - - print(f"[rebuild] {slug}: {ref_docs} reference docs, {user_docs} user preference docs") + print(f"[rebuild] {slug}: {ref_docs} reference docs") with flyte.group(f"rebuild:{slug}:cognify"): await cognee.cognify(datasets=[slug], chunk_size=512) @@ -937,7 +937,7 @@ async def rebuild_topic_dataset(rebuild_json: str) -> str: 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, "user_docs": user_docs}) + return json.dumps({"slug": slug, "ref_docs": ref_docs}) # --------------------------------------------------------------------------- @@ -1013,56 +1013,74 @@ async def sleep_cycle() -> dict: api_key = os.environ.get("ANTHROPIC_API_KEY", "") changed_topics: set[str] = set() - # Phase 2: Auto-promote user/ proposals (no human needed) + # 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) - staged = list_staged_proposals(store, limit=100) - # Skip proposals already processed in a prior sleep cycle - staged = [p for p in staged if not _is_already_archived(store, p.id)] - 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", - ) + 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="approved", - note="auto-promoted by sleep_cycle", + actor="sleep_cycle", decision="rejected", + note=decision.reason, ) - summary["promoted"] += 1 - # Use the proposal's classified topic slug, or fall back to - # classifying now for proposals staged before this feature - slug = proposal.topic_slug - if not slug: - slug = classify_proposal_topic( - proposal.content, proposal.source_question, topic_index, api_key - ) - effective_slug = slug if (slug and slug in topic_index) else GENERAL_TOPIC_SLUG - upsert_topic_map(store, decision.normalized_path, effective_slug) - changed_topics.add(effective_slug) - except Exception as e: - summary["errors"].append(f"promote:{proposal.id[:8]}:{type(e).__name__}") - else: - archive_proposal( - store, proposal, - actor="sleep_cycle", decision="rejected", - note=decision.reason, - ) - summary["rejected"] += 1 + summary["rejected"] += 1 - # Phase 3: Cluster + consolidate user/ memories in parallel + # Phase 3: Cluster + consolidate memories across all sessions in parallel with flyte.group("sleep:consolidate"): summary["phase"] = "consolidating" await _emit_report(summary) - clusters = _cluster_user_memories(store) + + 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: @@ -1096,8 +1114,9 @@ async def sleep_cycle() -> dict: ) summary["clusters_consolidated"] += 1 summary["memories_merged"] += len(merged_from) - # Resolve the topic for this consolidated path via the topic_map - topic_map = read_topic_map(store) + # 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) @@ -1164,15 +1183,16 @@ async def sleep_cycle() -> dict: @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/chat_sessions//transcript.jsonl + user/sessions//chat//transcript.jsonl Writes: - user/chat_sessions//summary.txt + user/sessions//chat//summary.txt This enables durable conversation continuity without stuffing the entire transcript into every prompt. @@ -1183,8 +1203,8 @@ async def summarize_chat_session( store = MemoryStore(LOCAL_MEMSTORE_ROOT) store.ensure_layout() - transcript_path = f"user/chat_sessions/{session_id}/transcript.jsonl" - summary_path = f"user/chat_sessions/{session_id}/summary.txt" + 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: @@ -1243,7 +1263,7 @@ async def summarize_chat_session( reason="chat-summary", expected_sha=expected, op="summarize", - extra_audit={"chat_session_id": session_id}, + extra_audit={"chat_session_id": session_id, "session": session}, ) with flyte.group("summary:upload"): @@ -1259,6 +1279,7 @@ async def summarize_chat_session( @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, @@ -1292,10 +1313,23 @@ async def wake_cycle( target_slugs = _route_query_to_topics(question, topic_index, api_key) if not target_slugs: target_slugs = list(topic_index.keys()) - if GENERAL_TOPIC_SLUG not in target_slugs: - target_slugs.append(GENERAL_TOPIC_SLUG) 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: @@ -1341,7 +1375,9 @@ def _extract_result_text(r) -> str: return "" if (s.startswith("<") and s.endswith(">")) else s context_parts = [_extract_result_text(r) for r in all_results[:5]] - context = "\n\n".join(p for p in context_parts if p) + 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}") From 55f575f8ab15a5c977e0c0f6c530a99b802112a2 Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Mon, 25 May 2026 21:37:41 -0700 Subject: [PATCH 5/8] fix Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- v2/tutorials/cognee_memory_store/workflow.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/v2/tutorials/cognee_memory_store/workflow.py b/v2/tutorials/cognee_memory_store/workflow.py index 970486b1..400f16db 100644 --- a/v2/tutorials/cognee_memory_store/workflow.py +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -291,10 +291,15 @@ def _is_nav_chrome(line: str) -> bool: # 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 < 3 prose words are nav chrome + # 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]{3,}\b", cleaned)) < 3 + return len(re.findall(r"\b[a-zA-Z]{4,}\b", cleaned)) < 3 lines = text.splitlines() i = 0 From b518bdf452e7cc77f06642f1d6e3b6fd3a4f4e89 Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Tue, 26 May 2026 19:07:47 -0700 Subject: [PATCH 6/8] add: multisession support Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- v2/tutorials/cognee_memory_store/app.py | 20 +++++++++- v2/tutorials/cognee_memory_store/workflow.py | 42 +++++++++++++++----- 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/v2/tutorials/cognee_memory_store/app.py b/v2/tutorials/cognee_memory_store/app.py index e50cc846..047927cb 100644 --- a/v2/tutorials/cognee_memory_store/app.py +++ b/v2/tutorials/cognee_memory_store/app.py @@ -54,6 +54,7 @@ SHARED_COGNEE_DB_PREFIX, SHARED_MEMSTORE_PATH, GENERAL_TOPIC_SLUG, + _configure_cognee_runtime, _setup_cognee_env, _route_query_to_topics, _strip_jina_header, @@ -400,6 +401,7 @@ async def _search() -> str: 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]), @@ -408,7 +410,20 @@ async def _search() -> str: all_results.extend(results or []) except (asyncio.TimeoutError, Exception): pass - return "\n".join(str(r) for r in all_results[:5]) + + 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() @@ -962,7 +977,8 @@ def _render_memory_viewer(store: MemoryStore) -> None: for src in sources[:3]: st.caption(f" {src}") topic_paths = store.list_paths(f"memory/{slug}") - for p in topic_paths[:5]: + 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)") diff --git a/v2/tutorials/cognee_memory_store/workflow.py b/v2/tutorials/cognee_memory_store/workflow.py index 400f16db..99b030c8 100644 --- a/v2/tutorials/cognee_memory_store/workflow.py +++ b/v2/tutorials/cognee_memory_store/workflow.py @@ -303,9 +303,9 @@ def _is_nav_chrome(line: str) -> bool: lines = text.splitlines() i = 0 - # Never skip more than the first third of the file (safety bound) - max_skip = min(len(lines), max(len(lines) // 3, 100)) - while i < max_skip and _is_nav_chrome(lines[i]): + # 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() @@ -805,19 +805,41 @@ async def ingest_url(url: str, max_pages: int = 10) -> tuple[Dir, Dir]: 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) ...") - await cognee.cognify(datasets=[slug], chunk_size=512) - print(f"[ingest] cognify complete for {slug!r}") - - with flyte.group("ingest:upload"): - print("[ingest] Uploading (pod stays alive for background cognee tasks) ...") - memstore_dir = await _upload_dir(LOCAL_MEMSTORE_ROOT, SHARED_MEMSTORE_PATH) - cognee_dir = await _upload_dir(local_cognee, _topic_db_path(slug)) + 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 From cb7d165fdfe9f324024852da29c89b92c9f7b24f Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Tue, 26 May 2026 19:11:10 -0700 Subject: [PATCH 7/8] update readme Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- v2/tutorials/cognee_memory_store/README.md | 250 ++++++++++++--------- 1 file changed, 144 insertions(+), 106 deletions(-) diff --git a/v2/tutorials/cognee_memory_store/README.md b/v2/tutorials/cognee_memory_store/README.md index bb53c464..96cfbc19 100644 --- a/v2/tutorials/cognee_memory_store/README.md +++ b/v2/tutorials/cognee_memory_store/README.md @@ -1,146 +1,184 @@ -# Cognee + Flyte “Memory Store” +# Cognee + Flyte Memory Store -This tutorial shows how to build **excellent, durable memory** for an agent by combining: +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/). -- **Cognee** → semantic memory (RAG / KG / embeddings): *retrieve relevant context* -- **Flyte v2** → orchestration + observability: *promote, cognify, persist, time, cache* -- **A file-based “memory store”** (this tutorial) → **auditable, versioned, access-controlled** durable memory +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) -It takes inspiration from Claude’s Managed Agents “memory stores” concepts: -- many small memories (files) -- multiple stores / mounts with different access rules -- immutable version history (audit trail) -- safe edits (optimistic concurrency) -- staging vs trusted memory (prompt-injection defense) +## Architecture -## Why this design -If you rely only on semantic search (embeddings), it’s hard to enforce: -- “don’t overwrite shared memory accidentally” -- “don’t let untrusted inputs write persistent instructions” -- “show me exactly what changed and when” +``` +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 +``` -So we add a **curated memory layer** that’s transparent and governable. +**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 | -## Architecture +## Storage layout ``` -User message - | - | (1) Retrieve - | - read user/ preferences - | - run cognee.search(query) - v -LLM answer (fast path runs inside app container) - | - | (2) Optional: stage proposals (untrusted) - v -staging/inbox/*.json (NOT trusted) - | - | (3) Flyte task: review_and_promote - | - validate (anti-poisoning rules) - | - write into user/ or shared/ - | - version snapshot + audit log - | - cognee.add(promoted memory) - v -Trusted memory (user/, shared/) + audit/ + versions/ - | - | (4) Flyte task: cognify_if_needed (cached) - v -Cognee KG / indices updated +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 ``` -## Memory store layout -Everything lives under a directory that is synced via `flyte.io.Dir` to object storage. +## Files -- `reference/` **read-only** (curated docs/conventions) -- `user/` **read-write** (per-user preferences, notes) -- `shared/` **promotion-only** (team/org conventions; requires explicit approval) -- `staging/inbox/` **untrusted proposals** -- `staging/archive/` archived proposals + decisions -- `versions/` immutable snapshots of every write -- `audit/log.jsonl` append-only log of memory mutations -- `meta/` current sha256 + last update info per memory +| 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 | -### Access modes (Claude-inspired) -- `reference/` is **read-only** (writes rejected) -- `shared/` is **write-behind-review** (direct writes rejected; only allowed via promotion) +## Prerequisites -### Optimistic concurrency (safe edits) -Writes can include an `expected_sha256` precondition. If the on-disk sha does not match, -we fail fast to avoid clobbering another writer. +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) -## Files -- `memory_store.py` — the memory store implementation (audit/versioning/access/concurrency) -- `agent.py` — proposal schema + staging + validation + promotion helpers -- `workflow.py` — Flyte tasks (init, answer, promote, cognify) -- `app.py` — Streamlit app (chat + memory sidebar; triggers Flyte tasks) +## Quickstart (Union) -## Run (Union / Flyte App) +```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 +``` -### 1) Ensure Anthropic key exists -This example expects a Union secret named `internal-anthropic-api-key` injected as `ANTHROPIC_API_KEY`. +Open the printed app URL. -### 2) Configure GHCR -This tutorial always publishes images to **GHCR**. +### Optional: configure image registry -Default registry is `ghcr.io/flyteorg`. If you want to use your own org/user, set: +By default, images are pushed to `ghcr.io/flyteorg`. Override with: ```bash export AI_MEMORY_STORE_IMAGE_REGISTRY="ghcr.io/" -# Optional: Union secret key that contains registry credentials (for private registries) -export AI_MEMORY_STORE_IMAGE_REGISTRY_SECRET="" ``` -### 3) Serve the app +## Using the app -```bash -# Default behavior prefers local image builds if Docker is available. -# To force remote builds: -# export AI_MEMORY_STORE_IMAGE_BUILDER=remote -uv run v2/tutorials/cognee_memory_store/app.py -``` +### 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. -Open the printed URL. +### 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 -### 3) Try the workflow -1. Ask a question. -2. Change preferences either: - - manually in the **Preferences** panel, or - - by saying something like “please be concise” / “use markdown” / “use my name Adil every time you answer”. - You’ll get a **preference approval popup**. -3. Confirm by clicking **Save preference** in the popup (or use **Save preferences** in the sidebar). -4. (Optional) Use **Advanced: stage raw proposal** + **Run promotion task** to exercise staged→promoted memory. -5. Click **Run cognify task**. -6. Ask a similar question again and observe retrieval improving. +### 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 -Set: +## Debug mode ```bash export AI_MEMORY_STORE_DEBUG=1 -export AI_MEMORY_STORE_MODEL=claude-haiku-4-5-20251001 ``` -The chat will print retrieval/answer timings when debug is enabled. +Enables per-message retrieval/answer timing and proposal detection details in the UI. -## Run (local dev) -Local mode does not require Flyte/Union credentials; persistence is local-only. +## Local dev (no Union required) ```bash streamlit run v2/tutorials/cognee_memory_store/app.py -- --server ``` -You can still stage proposals and inspect audit/versioning locally, but promotion/cognify -Flyte tasks require remote object storage. +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 -- **Auditability:** expand “Audit log (tail)” and see every stage/promote/archive event. -- **Versioning:** each promoted write creates a `versions//*` immutable snapshot. -- **Safety:** staged proposals are not trusted until validated and promoted. -- **Separation of concerns:** Cognee retrieves; the memory store governs what is persisted. - -## Notes / extensions -- Add an LLM-based “memory gate” validator (still stage first; only promote after approval). -- Split into multiple stores (per-user store vs shared org store) the same way Claude supports - multiple stores with different access rules. -- Add redaction tools over `audit/` + `versions/` for compliance workflows. + +- **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 From 63d640cad2a7e4f2e9ca30be56735fd6a9aab2b7 Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Tue, 26 May 2026 19:15:11 -0700 Subject: [PATCH 8/8] cleanup Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- CLAUDE.md | 108 ------------------------------------------------------ 1 file changed, 108 deletions(-) delete mode 100644 CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index 7998362d..00000000 --- a/CLAUDE.md +++ /dev/null @@ -1,108 +0,0 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Commands - -```bash -# Setup -make setup-venv # creates ~/.venv with uv + Python 3.12 -source ~/.venv/bin/activate -make update-flyte # installs latest flyte pre-release into active venv - -# Testing (venv must be active) -make test-preview # discover scripts without running -make test FILE=v2/tutorials/foo/main.py # run one script on Union cloud -make test-local FILE=v2/tutorials/foo/main.py # run one script locally (isolated venv per script) -make test FILTER=cognee # run all scripts whose path matches "cognee" -make test-local VERBOSE=vv FILE=... # with flyte verbosity (-v / -vv / -vvv) -make clean # remove test/reports/, test/logs/, test/venvs/ - -# Direct runner (same flags) -python3 test/test_runner.py --preview -python3 test/test_runner.py --local --file "v2/tutorials/foo/main.py" -``` - -Cloud tests need `FLYTE_CLIENT_SECRET` set; they target `playground.canary.unionai.cloud` (project `docs-examples`, domain `development`) as configured in `test/config.flyte.yaml`. - -## Repository layout - -``` -v2/ Modern Flyte 2.x examples (only these are tested) - tutorials/ Full end-to-end tutorial projects (multi-file) - user-guide/ Short single-file illustrative snippets -v1/ Legacy Flyte 1.x examples (not tested) -_blogs/ Blog-post companion code -test/ Test runner + CI config -``` - -## Flyte 2.x script conventions - -Every runnable script uses **PEP 723 inline metadata** at the top — this is how `uv run` discovers dependencies and how the test runner installs them per-script in isolation: - -```python -# /// script -# requires-python = "==3.13" -# dependencies = [ -# "flyte==2.1.5", -# "some-lib>=1.0", -# ] -# main = "main" ← entrypoint task name -# /// -``` - -### Core Flyte 2.x patterns - -**TaskEnvironment** groups config shared across tasks (image, secrets, resources): -```python -env = flyte.TaskEnvironment( - name="my-env", - image=flyte.Image.from_uv_script(__file__, name="my-image"), - secrets=[flyte.Secret(key="my-secret", as_env_var="MY_ENV_VAR")], - resources=flyte.Resources(cpu=2, memory="4Gi"), -) - -@env.task -async def my_task(...): ... -``` - -**`flyte.Image.from_uv_script(__file__)`** builds a container image from the PEP 723 deps automatically. Additional source files needed in the container are added with `.with_source_file(local_path, container_path)`. - -**`flyte.io.Dir`** is used to hand off directory state between tasks running in isolated pods — download with `await d.download(local_path=...)`, upload with `await Dir.from_local(local_path, remote_destination=...)`. - -**`flyte.map.aio(task_fn, inputs, concurrency=N)`** fans out a task as parallel pods and async-iterates results. - -**Scheduling**: `flyte.Trigger` + `flyte.Cron` registers a schedule on the cluster without needing a LaunchPlan. - -**`flyte.group("label")`** creates a named span visible in the Union UI execution timeline. - -**`report=True`** on a task enables live HTML streaming to the Union UI dashboard. - -**`cache="auto"`** keys caching on task inputs + source code hash — useful for idempotent subtasks. - -### Deployment - -```bash -# Register schedules and deploy apps in one shot -python workflow.py --deploy -python app.py # also calls flyte.deploy() then flyte.serve() -``` - -## Active tutorial: cognee_memory_store - -The most complex active tutorial. It implements a sleep/wake memory architecture using Cognee (knowledge graph) + Flyte. - -**Files:** -- `memory_store.py` — audited, versioned file-based memory store synced via `flyte.io.Dir`. Prefixes enforce access: `reference/` is read-only, `user/` is read-write. -- `agent.py` — staged proposal pipeline: untrusted writes land in `staging/inbox/`, validated, then promoted to `user/`. -- `workflow.py` — all Flyte tasks: `init_reference`, `ingest_url`, `consolidate_cluster`, `sleep_cycle`, `wake_cycle`, `summarize_chat_session`. -- `app.py` — Streamlit UI served via `flyte.app.AppEnvironment`. - -**Sleep/wake cycle:** -- **Sleep** (every 6h via Cron): download state → promote staged proposals → cluster + consolidate `user/` memories via `flyte.map.aio` → full graph rebuild (per-topic `empty_dataset` + `cognify`) → upload. -- **Wake** (per question): download state → classify question to topic slugs → `cognee.search(datasets=[slugs])` → Claude answer. -- **`ingest_url`**: scrapes URL via Jina Reader (`r.jina.ai/`) with direct-HTTP fallback, stores in `memory/topic_/`, `cognify` + `memify` for enrichment. - -**Shared object storage paths:** -- `cognee-memory-store/memstore` → memstore files -- `cognee-memory-store/cognee_db` → Cognee SQLite state \ No newline at end of file