from uuid import UUID from typing import Union, Optional, List, Type, Any try: from typing import Unpack except ImportError: from typing_extensions import Unpack from typing_extensions import TypedDict from cognee.shared.logging_utils import get_logger from cognee.modules.observability import ( new_span, COGNEE_DATASET_NAME, COGNEE_SESSION_ID, COGNEE_IMPROVE_STAGES, ) logger = get_logger("improve") class ImproveKwargs(TypedDict, total=False): """Power-user overrides for improve(). Most users never need these.""" extraction_tasks: list enrichment_tasks: list data: Any node_type: Type user: object vector_db_config: dict graph_db_config: dict feedback_alpha: float async def improve( dataset: Union[str, UUID] = "main_dataset", *, run_in_background: bool = False, node_name: Optional[List[str]] = None, session_ids: Optional[List[str]] = None, build_global_context_index: bool = False, build_truth_subspace: bool = False, **kwargs: Unpack[ImproveKwargs], ): """Enrich an existing knowledge graph with additional context and rules. When ``session_ids`` is provided, the improvement pipeline runs four stages: 1. **Apply feedback weights** -- session entries with feedback scores update ``feedback_weight`` on the graph nodes/edges that were used to produce those answers. Higher-rated answers boost their source nodes; lower-rated answers decrease them. 2. **Persist session Q&A** -- the question/answer text from those sessions is cognified into the permanent graph, tagged with ``node_set="user_sessions_from_cache"``. 2c. **Distill sessions** -- each session's gated active-guidance entries are curated into entity-anchored lessons and add+cognified into the graph (tagged ``session_learnings``). Sessions with no gated guidance produce nothing. This is what lets ``remember(session, self_improvement=True)`` cover session distillation without an explicit ``distill_session`` call. 3. **Default enrichment** -- triplet embeddings are extracted and indexed (same as calling ``improve()`` without sessions). 4. **Global context index** -- when ``build_global_context_index=True``, builds retrieval-ready bucket and root summaries over the graph's text summaries. Without ``session_ids``, only stage 3 runs by default. Args: dataset: Dataset name or UUID to process. run_in_background: Run processing asynchronously. node_name: Filter graph to specific named entities. session_ids: Session IDs whose feedback and Q&A content should be bridged into the permanent graph. build_global_context_index: Opt-in flag for building the global context index after default enrichment. Skipped in background mode because ordered background pipeline chaining is not supported yet. build_truth_subspace: Opt-in flag (default ``False``) for building the truth subspace from distilled session learnings after distillation and before enrichment. Only runs when ``session_ids`` is provided. Off by default = no behaviour change. **kwargs: Additional options -- see ``ImproveKwargs``. Returns: Pipeline run info (same as ``cognee.memify()``). Example:: # Enrich graph + bridge session feedback and content await cognee.improve(dataset="docs", session_ids=["chat_1", "chat_2"]) # Enrich graph only (no session bridging) await cognee.improve(dataset="docs") """ from cognee.shared.utils import send_telemetry from cognee import __version__ as cognee_version stages_run = [] send_telemetry( "cognee.improve", kwargs.get("user", "sdk"), additional_properties={ "dataset": str(dataset), "session_count": len(session_ids) if session_ids else 0, "session_ids": ",".join(session_ids) if session_ids else "", "run_in_background": run_in_background, "build_global_context_index": build_global_context_index, "build_truth_subspace": build_truth_subspace, "cognee_version": cognee_version, }, ) with new_span("cognee.api.improve") as span: span.set_attribute(COGNEE_DATASET_NAME, str(dataset)) if session_ids: span.set_attribute(COGNEE_SESSION_ID, ",".join(session_ids)) from cognee.api.v1.serve.state import get_remote_client client = get_remote_client() if client is not None: return await client.improve(dataset, node_name=node_name, **kwargs) from cognee.modules.users.methods import get_default_user user = kwargs.pop("user", None) if user is None: user = await get_default_user() feedback_alpha = kwargs.pop("feedback_alpha", 0.1) # Mutex: single-session improves serialize on the session's # lock so auto-improve + idle-watcher + SessionEnd don't # duplicate work. Multi-session improves skip the lock — the # pattern is rare and locking N sessions at once is messy. acquired_lock_for: Optional[str] = None if session_ids and len(session_ids) == 1: from cognee.infrastructure.locks import ( release_improve_lock, try_acquire_improve_lock, ) sole_session = session_ids[0] if not await try_acquire_improve_lock(sole_session): logger.info( "improve: session '%s' already being improved, skipping", sole_session, ) return {} acquired_lock_for = sole_session try: # Stage 1 & 2: bridge sessions into the permanent graph if session_ids: await _bridge_sessions( dataset=dataset, session_ids=session_ids, user=user, feedback_alpha=feedback_alpha, run_in_background=run_in_background, ) stages_run.extend(["feedback_weights", "persist_sessions"]) # Stage 2b: persist agent trace steps (tool calls with # per-step feedback) into the graph. Without this, the # plugin's trace activity never reaches permanent # memory — only QA entries do. await _persist_session_traces( dataset=dataset, session_ids=session_ids, user=user, run_in_background=run_in_background, ) stages_run.append("persist_trace_steps") # Stage 2b2: distill each session's agent traces into agent-profile # session-context lessons (the LLM batch pass) before distillation, so # those lessons are available as gated guidance for stage 2c. if await _extract_agent_context(session_ids=session_ids, user=user): stages_run.append("extract_agent_context") # Stage 2c: distill each session's gated guidance into curated, # entity-anchored lessons and add+cognify them into the graph. # This is what lets remember(session, self_improvement=True) # cover session distillation without an explicit # cognee.session.distill_session call. distilled = await _distill_sessions( dataset=dataset, session_ids=session_ids, user=user, ) if distilled: stages_run.append("distill_sessions") # Stage 2d: build the truth subspace from distilled session # learnings (opt-in, default OFF). Runs after distillation so # freshly accepted lessons are available as anchors, and before # enrichment. Non-fatal — never blocks the rest of improve(). if build_truth_subspace: try: from cognee.modules.truth_subspace.build import ( build_truth_subspace as _build_truth_subspace, ) result_ts = await _build_truth_subspace( dataset=dataset, session_ids=session_ids, user=user, ) logger.info("improve: truth subspace built -> %s", result_ts) stages_run.append("build_truth_subspace") except Exception as e: logger.warning("improve: truth subspace build failed (non-fatal): %s", e) # Stage 3: default enrichment (triplet embeddings) from cognee.modules.memify import memify if "node_type" not in kwargs or kwargs.get("node_type") is None: from cognee.modules.engine.models.node_set import NodeSet kwargs["node_type"] = NodeSet result = await memify( dataset=dataset, node_name=node_name, user=user, run_in_background=run_in_background, **kwargs, ) stages_run.append("memify_enrichment") if build_global_context_index: if run_in_background: logger.warning( "improve: global context index skipped in background mode " "because ordered background pipeline chaining is not supported" ) else: global_context_index_updated = await _build_global_context_index( dataset=dataset, user=user, ) if global_context_index_updated: stages_run.append("global_context_index") span.set_attribute(COGNEE_IMPROVE_STAGES, ",".join(stages_run)) return result finally: if acquired_lock_for: from cognee.infrastructure.locks import release_improve_lock await release_improve_lock(acquired_lock_for) async def _build_global_context_index( dataset: Union[str, UUID], user, ) -> bool: from cognee.memify_pipelines.global_context_index import global_context_index_pipeline try: await global_context_index_pipeline( user=user, dataset=dataset, run_in_background=False, bucketing_strategy="graph", max_bucket_size=4, ) logger.info("improve: global context index updated") return True except Exception as e: logger.warning("improve: global context index update failed (non-fatal): %s", e) return False async def _resolve_dataset_name(dataset: Union[str, UUID], user) -> str: """Resolve a dataset reference to its name string.""" if isinstance(dataset, str): return dataset from cognee.modules.data.methods.get_authorized_dataset import get_authorized_dataset ds = await get_authorized_dataset(user, dataset, "write") return ds.name if ds else "main_dataset" async def _bridge_sessions( dataset: Union[str, UUID], session_ids: List[str], user, feedback_alpha: float, run_in_background: bool, ): """Run feedback weights and session persistence pipelines. Stage 1 (feedback weights): Updates ``feedback_weight`` on graph nodes and edges that were *used during retrieval* in session Q&A entries. Only elements referenced in ``used_graph_element_ids`` are affected. If no retrieval has occurred in these sessions, no weights are updated. Stage 2 (persist Q&A): Cognifies the actual question/answer text from sessions into the permanent graph, tagged with ``node_set="user_sessions_from_cache"``. This persists the Q&A content itself, not serialized graph edges. """ # Stage 1: apply feedback weights from session retrieval traces from cognee.memify_pipelines.apply_feedback_weights import apply_feedback_weights_pipeline dataset_name = await _resolve_dataset_name(dataset, user) try: await apply_feedback_weights_pipeline( user=user, session_ids=session_ids, dataset=dataset_name, alpha=feedback_alpha, run_in_background=run_in_background, ) logger.info("improve: feedback weights applied from %d session(s)", len(session_ids)) except Exception as e: logger.warning("improve: feedback weights failed (non-fatal): %s", e) # Stage 2: persist session Q&A into permanent graph from cognee.memify_pipelines.persist_sessions_in_knowledge_graph import ( persist_sessions_in_knowledge_graph_pipeline, ) try: await persist_sessions_in_knowledge_graph_pipeline( user=user, session_ids=session_ids, dataset=dataset_name, run_in_background=run_in_background, ) logger.info("improve: session Q&A persisted from %d session(s)", len(session_ids)) except Exception as e: logger.warning("improve: session persistence failed (non-fatal): %s", e) async def _extract_agent_context( session_ids: List[str], user, ) -> int: """Flush pending trace windows into agent-profile lessons before distillation. Delegates to ``agent_context_extraction.extract_pending_agent_context`` per session, which shares the same watermark used by mid-session trace extraction. ``min_new_traces=1`` makes improve/session-end flush any remaining unprocessed traces before distillation. Gated on automatic session context and best-effort/fail-open: an error on one session never blocks the others or the rest of ``improve()``. Returns the number of lessons created/linked. """ from cognee.infrastructure.session.agent_context_extraction import ( extract_pending_agent_context, ) from cognee.infrastructure.session.get_session_manager import get_session_manager session_manager = get_session_manager() if not session_manager.is_available or not session_manager.is_auto_feedback_enabled(): return 0 user_id = str(user.id) touched = 0 for session_id in session_ids: try: ids = await extract_pending_agent_context( session_manager=session_manager, user_id=user_id, session_id=session_id, min_new_traces=1, ) touched += len(ids) except Exception as e: logger.warning( "improve: agent-context extraction failed for '%s' (non-fatal): %s", session_id, e, ) return touched async def _distill_sessions( dataset: Union[str, UUID], session_ids: List[str], user, ) -> int: """Distill each session's gated learnings into curated lessons in the graph. Delegates to ``session_distillation.distill_session`` per session: it loads the session's gated active-guidance entries, curates them into proposed lessons, writes/rejects each with entity anchoring, and add+cognifies the accepted lessons into ``dataset`` (tagged ``session_learnings``). Best-effort and fail-open: a session with no gated guidance simply yields no lessons (status ``no_gated_entries``), and an error on one session never blocks the others or the rest of ``improve()``. Returns the total number of lesson documents written across all sessions. Note: ``distill_session`` runs its own ``add``/``cognify`` (it does not call ``improve``), so there is no recursion back into this function. """ from cognee.modules.session_distillation import distill_session distilled = 0 for session_id in session_ids: try: result = await distill_session(session_id, dataset=dataset, user=user) distilled += len(result.documents) logger.info( "improve: distilled session '%s' -> status=%s documents=%d", session_id, result.status, len(result.documents), ) except Exception as e: logger.warning( "improve: session distillation failed for '%s' (non-fatal): %s", session_id, e, ) return distilled async def _persist_session_traces( dataset: Union[str, UUID], session_ids: List[str], user, run_in_background: bool, ): """Cognify per-step agent trace feedbacks into the knowledge graph. Without this step, the Claude Code plugin's tool-call activity (the bulk of session data — hundreds of Bash/Edit/Read/Write trace steps per session) never makes it into permanent memory. Only QA entries do via ``persist_sessions_in_knowledge_graph_pipeline``. Runs the dedicated ``persist_agent_trace_feedbacks_in_knowledge_graph_pipeline`` that extracts per-step ``session_feedback`` from the cache and cognifies it into the ``agent_trace_feedbacks`` node-set. """ dataset_name = await _resolve_dataset_name(dataset, user) try: from cognee.memify_pipelines.persist_agent_trace_feedbacks_in_knowledge_graph import ( persist_agent_trace_feedbacks_in_knowledge_graph_pipeline, ) await persist_agent_trace_feedbacks_in_knowledge_graph_pipeline( user=user, session_ids=session_ids, dataset=dataset_name, node_set_name="agent_trace_feedbacks", raw_trace_content=False, last_n_steps=None, # persist all stored steps on demand run_in_background=run_in_background, ) logger.info( "improve: agent trace steps persisted from %d session(s)", len(session_ids), ) except Exception as e: logger.warning("improve: trace persistence failed (non-fatal): %s", e)