File size: 9,222 Bytes
2e818da
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""Orchestrates one Commit's adaptation review: reads unreviewed student
interactions and already-approved staged profile candidates, makes one
Cerebras call, routes proposals through MemoryPromotionGate, persists
approved candidates, and updates the persona only after Cognee confirms
durability. See docs/commit-memory-adaptation-architecture.md for the full
design and crash-safety rationale.
"""
from __future__ import annotations

import asyncio
import uuid
from collections import defaultdict

from app.rag.project_memory import ProjectMemoryStore
from app.services import approved_profile_staging, pending_review_store, student_interaction_log
from app.services.agent_control import AgentControlService
from app.services.commit_candidate_adapters import candidate_to_project_memory_record, proposal_to_memory_candidate
from app.services.memory_promotion import MemoryPromotionGate
from app.services.student_memory import StudentMemoryService

MAX_INTERACTIONS_CHARS = 8_000
MAX_PROFILE_CONTEXT_CHARS = 6_000
MAX_PERSONA_CHARS = 3_000

# Module-level singletons so tests can monkeypatch them and so the
# per-project lock registry is shared across calls within one process.
agent_control = AgentControlService()
agent_control_persona = agent_control
student_memory = StudentMemoryService()
project_memory = ProjectMemoryStore()
memory_gate = MemoryPromotionGate()

_adaptation_locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)


def _clip_text(text: str, limit: int) -> str:
    if len(text) <= limit:
        return text
    return text[:limit].rstrip() + "\n...[truncated]"


def _format_interactions(interactions) -> str:
    return "\n".join(f"- [{i.interaction_id}] {i.text}" for i in interactions)


def adaptation_lock(project_id: str) -> asyncio.Lock:
    return _adaptation_locks[project_id]


def _gate_and_split(candidates):
    promoted_project, promoted_student = [], []
    for candidate in candidates:
        decision = (
            memory_gate.evaluate_project(candidate) if candidate.destination == "project"
            else memory_gate.evaluate_student(candidate)
        )
        if not decision.promote:
            continue
        (promoted_project if candidate.destination == "project" else promoted_student).append(candidate)
    return promoted_project, promoted_student


async def run_commit_adaptation(*, project_id: str) -> None:
    async with adaptation_lock(project_id):
        pending = pending_review_store.load_pending_review(project_id)

        if pending is not None and pending.status == "cognee_flush_started":
            # A prior run crashed (or the process restarted) between marking
            # this batch as "Cognee may be mutating state" and confirming
            # the outcome. We cannot prove Cognee's write was safe to retry
            # -- quarantine it and move on rather than risk a duplicate.
            pending_review_store.archive_as_uncertain(project_id, pending)
            student_interaction_log.advance_cursor(project_id, pending.interaction_through_offset)
            approved_profile_staging.advance_cursor(project_id, pending.profile_staging_through_offset)
            pending_review_store.delete_pending_review(project_id)
            return

        if pending is None:
            interactions, i_from, i_through = student_interaction_log.read_pending(project_id)
            staged, s_from, s_through = approved_profile_staging.read_pending(project_id)
            if not interactions and not staged:
                return

            current_profile = await student_memory.query_prior_knowledge(
                "student preferences, knowledge, strengths, recurring difficulties, and goals",
                project_id=project_id,
                mode="profile",
            )
            review = await agent_control.review_commit_with_llm(
                project_id=project_id,
                interactions_text=_clip_text(_format_interactions(interactions), MAX_INTERACTIONS_CHARS),
                current_profile=_clip_text(current_profile, MAX_PROFILE_CONTEXT_CHARS),
                current_persona=_clip_text(agent_control_persona.get_inferred_text(project_id), MAX_PERSONA_CHARS),
                manual_persona=_clip_text(agent_control_persona.get_manual_text(project_id), MAX_PERSONA_CHARS),
            )

            interaction_index = {i.interaction_id: i for i in interactions}
            review_candidates = [
                proposal_to_memory_candidate(p, project_id=project_id, interactions_by_id=interaction_index)
                for p in review.memory_proposals
            ]
            promoted_project, promoted_student_from_review = _gate_and_split(review_candidates)

            # Already-approved staged candidates bypass the gate a second
            # time -- they produced their one promotion-decision ledger row
            # at capture time (see explicit_profile_capture.py).
            promoted_student_by_id = {c.candidate_id: c for c in [*promoted_student_from_review, *staged]}
            promoted_student = list(promoted_student_by_id.values())

            pending_review_store.save_pending_review(
                project_id, status="ready_to_flush",
                candidate_ids=[c.candidate_id for c in [*promoted_project, *promoted_student]],
                interaction_from_offset=i_from, interaction_through_offset=i_through,
                profile_staging_from_offset=s_from, profile_staging_through_offset=s_through,
                review=review,
            )
            pending = pending_review_store.load_pending_review(project_id)
        else:
            # Retry: re-derive candidates from the exact stored ranges and
            # the exact stored review text -- never a freshly re-proposed
            # or silently-expanded batch. Chroma/Cognee IDs stay stable.
            interactions = student_interaction_log.read_range(
                project_id, pending.interaction_from_offset, pending.interaction_through_offset,
            )
            staged = approved_profile_staging.read_range(
                project_id, pending.profile_staging_from_offset, pending.profile_staging_through_offset,
            )
            interaction_index = {i.interaction_id: i for i in interactions}
            review_candidates = [
                proposal_to_memory_candidate(p, project_id=project_id, interactions_by_id=interaction_index)
                for p in pending.review.memory_proposals
            ]
            candidate_ids = set(pending.candidate_ids)
            promoted_project = [c for c in review_candidates if c.destination == "project" and c.candidate_id in candidate_ids]
            promoted_student_from_review = [
                c for c in review_candidates if c.destination == "student" and c.candidate_id in candidate_ids
            ]
            promoted_student_by_id = {c.candidate_id: c for c in [*promoted_student_from_review, *staged]}
            promoted_student = [c for c in promoted_student_by_id.values() if c.candidate_id in candidate_ids]

        pending.promoted_project = promoted_project
        pending.promoted_student = promoted_student

        loop = asyncio.get_running_loop()
        project_records = [candidate_to_project_memory_record(c) for c in pending.promoted_project]
        if project_records:
            try:
                written = await loop.run_in_executor(None, project_memory.upsert, project_records)
            except Exception:
                return
            if written != len(project_records):
                return

        if pending.promoted_student:
            pending_review_store.update_pending_status(
                project_id, status="cognee_flush_started", flush_attempt_id=str(uuid.uuid4()),
            )
            try:
                for candidate in pending.promoted_student:
                    if not await student_memory.stage_promoted_candidate(candidate):
                        pending_review_store.update_pending_status(project_id, status="cognee_flush_uncertain")
                        return
                flush_results = await student_memory.flush_project(project_id)
            except Exception:
                pending_review_store.update_pending_status(project_id, status="cognee_flush_uncertain")
                return
            # A clean False here is still a caught exception inside
            # StudentMemoryService's flush pipeline -- it cannot be
            # distinguished from "Cognee partially wrote before failing."
            # Treat identically to the exception branch above.
            if flush_results.get("research_profile") is not True:
                pending_review_store.update_pending_status(project_id, status="cognee_flush_uncertain")
                return
            pending_review_store.update_pending_status(project_id, status="cognee_flush_succeeded")

        agent_control_persona.set_inferred_text(project_id, pending.review.inferred_persona_text)
        student_interaction_log.advance_cursor(project_id, pending.interaction_through_offset)
        approved_profile_staging.advance_cursor(project_id, pending.profile_staging_through_offset)
        pending_review_store.delete_pending_review(project_id)