Cognee Improve Sessions

Session memory and improve()

cognee has two kinds of memory:

• Session memory: a fast cache of conversation turns, agent traces, and feedback, keyed by session_id. Writing is instant, with no LLM extraction.
• The permanent graph: what remember() builds without a session.

improve() connects them. It bridges session content into the graph and enriches the graph itself. remember() calls it automatically, so most users never call it directly.

pythonCopy code
import cognee

# Session write: returns immediately; improve() bridges it in the background
await cognee.remember("User prefers dark mode.", session_id="chat_1")

# Session-aware query: session cache first, then the graph
results = await cognee.recall("What does the user prefer?", session_id="chat_1")

# Bridge sessions into a dataset's graph explicitly
result = await cognee.improve(dataset="main_dataset", session_ids=["chat_1"])
print(result.status, result.stage_summary())

await cognee.wait_for_background_tasks()  # before a script exits

Use it

Writing session memory

Requirements: CACHING=true (default). The cache backend is CACHE_BACKEND, one of sqlite (default), postgres, redis, fs, tapes. Sessions expire after SESSION_TTL_SECONDS (default 7 days).

What improve() does: nine stages, in order

Every run goes through the same ordered stages (cognee/modules/improve/registry.py). Each stage checks a gate before it spends any LLM or embedding cost, and reports one StageResult.

Stages 1–7 need session_ids; 8 and 9 work on the graph alone. The order matters: 4 feeds 5, 5 feeds 7, and 7 runs before 8.

Reading the result

improve() returns an ImproveResult, and so do POST /api/v1/improve, the CLI, and RememberResult.improve.

• result.status: completed, errored (any stage errored), skipped (every stage skipped), or running (background, not finished).
• result.stages: one StageResult per stage, with stage, status (completed / already_completed / skipped / errored), reason, error, counts, duration_ms.
• result.stage("distill_sessions"), result.stage_summary(), result.lock_held, result.rerun_requested, result.rerun_passes.
• await result.wait() finishes a background run (no-op otherwise).

Skip and no-op reasons:

How remember() triggers improve

• Without a session: add, then cognify, then a foreground improve(). Its outcome is on result.improve / result.improve_error. A failed improve never marks the remember as errored.
• With session_id: the text is cached, then a background improve(dataset, session_ids=[session_id]) starts if the debounce allows it. result.improve fills in only after await result or wait_for_background_tasks().
• self_improvement=False turns it off per call; IMPROVE_AUTO_ENABLED=false turns it off everywhere and overrides self_improvement=True.
• An application embedding cognee can decline one automatic improve before it starts: cognee.modules.improve.register_auto_improve_admission(check) registers one async check that remember() awaits on both paths. Returning a reason string skips the improve and sets result.improve_skipped (result.improve stays None); the data is stored either way. A check that raises allows the improve. Explicit improve() calls are never gated.

Settings (IMPROVE_*, cognee/modules/improve/config.py)

There is no debounce timer: held-back entries wait for the next remember() with that session, or an explicit improve().

Pitfalls

• A plain improve(dataset) often does nothing. Without session_ids, only stages 8 and 9 can run, and both are off by default. Result: every stage skipped. That is expected, not an error.
• Typed entries do not auto-improve. remember(QAEntry/TraceEntry/ FeedbackEntry, session_id=...) stores the entry but never starts an improve. Call improve(session_ids=[...]) yourself.
• lock_held does not wait. Improves for the same dataset or session run one at a time: a second call returns at once with every stage skipped: lock_held. If it shares a session with the running one, it sets rerun_requested=True and the holder runs up to 2 extra passes (3 passes in total; see rerun_passes). The bound is IMPROVE_MAX_RERUN_PASSES, a constant in cognee/api/v1/improve/improve.py, not an env var. The lock is per process only; multiple API workers do not share it.
• Sessions bridge once. Q&A and trace persistence are tracked per user and session, not per dataset, so bridging a session into dataset A and then into dataset B persists no new Q&A/traces into B. Distillation is tracked per (session, dataset) and still runs into B.
• IMPROVE_STAGES_DISABLED is validated. An unknown stage name, or persist_session_qa (fatal, cannot be disabled), raises ValueError, and the API server refuses to start.
• Session writes with the cache off. remember(text, session_id=...) with CACHING=false only logs a warning and stores nothing; typed entries raise RuntimeError. An explicit improve(session_ids=...) then fails in the fatal stage.
• Fatal stage failure. If stage 2 errors, improve() raises (HTTP 409) and the error carries .improve_result. In background mode it is recorded on result.error instead. Other stages failing only mark themselves errored; HTTP still returns 200, so check status.
• Remote mode. After cognee.serve(url), a background improve is fire-and-forget: status stays running and there is no polling.
• cognee/modules/session_bridge/ is gone (only stale __pycache__ may be left). The bridging now lives in the improve stages; some test names still say "session_bridge".

How it works

improve() resolves the dataset once (creating it if the name is new), claims the improve lock for dataset:<id> plus every session:<user>:<session_id>, probes the graph adapter's capabilities, and runs the stages with one frozen ImproveRunInputs. Each stage is a gate plus a call into existing pipelines plus a result mapping. Stages never own retries or ordering.

Watermarks keep repeat runs cheap:

• Session stages track how many entries were already persisted, per user and session.
• Stage 8 compares the last completed enrichment against later write pipelines for the dataset. A node_name or custom-task run bypasses that check.

The improve operation row is written when the run finishes: failed if any stage errored (a retry is never gated off), noop if nothing ran.

• Orchestrator: cognee/api/v1/improve/improve.py (HTTP: routers/get_improve_router.py; CLI: cognee/cli/commands/improve_command.py)
• Stages, registry, results: cognee/modules/improve/ (stages.py, registry.py, stage.py, result.py, inputs.py, capabilities.py, config.py, graph_changes.py, constants.py)
• Lock: cognee/infrastructure/locks/session_lock.py
• Session store and watermarks: cognee/infrastructure/session/ (session_manager.py, session_persist_watermark.py, feedback_detection.py); cache backends in cognee/infrastructure/databases/cache/
• Auto-improve from remember: cognee/api/v1/remember/remember.py, cognee/api/v1/remember/auto_improve_debounce.py
• Entry types: cognee/memory/entries.py
• Background tasks: cognee/infrastructure/background_tasks.py

Examples: examples/guides/improve_quickstart.py, sessions.py, session_distillation.py, global_context_index.py, agent_memory_quickstart.py, and examples/advanced_guides/remember_recall_improve_example.py.

Extending it

Adding a stage:

1. Subclass BaseStage in cognee/modules/improve/stages.py. Set name, needs_sessions, and fatal (leave it False; exactly one fatal stage is enforced at import). Implement gate() (return a skip reason constant, or None, before any LLM/embedding cost) and run() (call existing pipeline code, return a StageResult).
2. Insert it into DEFAULT_STAGES in registry.py at the right position. The order is pinned by cognee/tests/unit/modules/improve/test_registry_order.py; update it deliberately.
3. If it can re-run cheaply, give it a watermark so a repeat run reports already_completed.
4. If it writes graph data under a new pipeline name, add that name to WRITE_PIPELINE_NAMES in graph_changes.py, or stage 8 will not notice the change.
5. Tests: cognee/tests/unit/modules/improve/ (gates, results, config) and cognee/tests/unit/api/v1/improve/ (orchestration, rerun, router).