app/rag/context_composer.py

"""Compose source evidence and the two memory scopes without conflating them."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable

from app.observability.operation import Operation, observe_operation, observe_stage
from app.rag.models import EvidenceBundle, EvidenceRequest, MemoryContextItem
from app.rag.project_memory import ProjectMemoryStore
from app.rag.retrieval_service import PaperEvidenceService


StudentMemoryReader = Callable[[str], Awaitable[str]]


class ContextComposer:
    """Single context-transfer boundary for chat, wiki, pair buddy, and agents."""

    def __init__(
        self,
        evidence: PaperEvidenceService | None = None,
        project_memory: ProjectMemoryStore | None = None,
        student_memory_reader: StudentMemoryReader | None = None,
    ) -> None:
        self.evidence = evidence or PaperEvidenceService()
        self.project_memory = project_memory or ProjectMemoryStore()
        self._student_memory_reader = student_memory_reader

    async def compose(self, request: EvidenceRequest) -> EvidenceBundle:
        consumer = request.consumer.value
        with observe_operation("context.compose", subsystem="context", consumer=consumer) as op:
            with observe_stage(op, "source_context", subsystem="context", consumer=consumer) as stage:
                source = self.evidence.retrieve(request)
                stage.add_count("source_evidence_count", len(source.evidence))
                stage.add_count("source_context_chars", source.diagnostics.source_context_chars)
                stage.add_count("context_truncated", int(source.diagnostics.context_truncated))
            project_items, student_items = await self.recall_memory(request, observer=op)
            with observe_stage(
                op,
                "context_composition",
                subsystem="context",
                consumer=consumer,
            ):
                diagnostics = source.diagnostics.model_copy(deep=True)
                diagnostics.project_memory_chars = sum(len(item.statement) for item in project_items)
                diagnostics.student_memory_chars = sum(len(item.statement) for item in student_items)
                bundle = EvidenceBundle(
                    source_evidence=source.evidence,
                    project_memory=project_items,
                    student_memory=student_items,
                    citations=source.citations,
                    diagnostics=diagnostics,
                )
            op.add_count("source_evidence_count", len(source.evidence))
            op.add_count("project_memory_count", len(project_items))
            op.add_count("student_memory_count", len(student_items))
            op.add_count("source_context_chars", diagnostics.source_context_chars)
            op.add_count("project_memory_chars", diagnostics.project_memory_chars)
            op.add_count("student_memory_chars", diagnostics.student_memory_chars)
            op.add_count("context_truncated", int(diagnostics.context_truncated))
            op.set("project_memory_policy", request.project_memory_policy)
            op.set("student_memory_policy", request.student_memory_policy)
            return bundle

    async def recall_memory(
        self,
        request: EvidenceRequest,
        *,
        observer: Operation | None = None,
    ) -> tuple[list[MemoryContextItem], list[MemoryContextItem]]:
        if observer is None:
            consumer = request.consumer.value
            with observe_operation(
                "context.memory_recall",
                subsystem="context",
                consumer=consumer,
            ) as operation:
                return await self._recall_memory(request, operation)
        return await self._recall_memory(request, observer)

    async def _recall_memory(
        self,
        request: EvidenceRequest,
        observer: Operation,
    ) -> tuple[list[MemoryContextItem], list[MemoryContextItem]]:
        consumer = request.consumer.value
        project_items: list[MemoryContextItem] = []
        with observe_stage(
            observer,
            "project_memory_recall",
            subsystem="memory",
            consumer=consumer,
        ) as stage:
            if request.project_memory_policy == "relevant":
                try:
                    project_items = await asyncio.to_thread(
                        self.project_memory.search,
                        request.project_id,
                        request.query,
                        6,
                        consumer=consumer,
                    )
                except Exception as exc:  # project memory must not block source RAG
                    stage.record_exception(
                        exc,
                        escaped=False,
                        stage="project_memory_recall",
                        category="project_memory_recall_fallback",
                    )
                    stage.mark_terminal("degraded")
                    observer.mark_terminal("degraded")
            stage.add_count("project_memory_count", len(project_items))
            stage.add_count(
                "project_memory_chars",
                sum(len(item.statement) for item in project_items),
            )
            stage.set("memory_policy", request.project_memory_policy)

        student_items: list[MemoryContextItem] = []
        with observe_stage(
            observer,
            "student_memory_recall",
            subsystem="memory",
            consumer=consumer,
        ) as stage:
            if request.student_memory_policy != "none":
                try:
                    statement = await self._read_student_memory(
                        request.query,
                        policy=request.student_memory_policy,
                    )
                except Exception as exc:  # defensive boundary for custom readers
                    stage.record_exception(
                        exc,
                        escaped=False,
                        stage="student_memory_recall",
                        category="student_memory_recall_fallback",
                    )
                    stage.mark_terminal("degraded")
                    observer.mark_terminal("degraded")
                    statement = ""
                if statement:
                    student_items.append(
                        MemoryContextItem(
                            memory_id="student_profile_recall",
                            source="student_memory",
                            statement=statement,
                            kind="student_profile",
                            metadata={"policy": request.student_memory_policy},
                        )
                    )
            stage.add_count("student_memory_count", len(student_items))
            stage.add_count(
                "student_memory_chars",
                sum(len(item.statement) for item in student_items),
            )
            stage.set("memory_policy", request.student_memory_policy)
        return project_items, student_items

    async def _read_student_memory(self, query: str, *, policy: str) -> str:
        if self._student_memory_reader is not None:
            return await self._student_memory_reader(query)
        from app.services.student_memory import StudentMemoryService

        # Cognee owns cross-project student identity/preferences only. Project
        # research state is intentionally retrieved from Chroma above.
        memory = StudentMemoryService()
        if policy == "cached":
            return await memory.query_prior_knowledge(
                query,
                mode="profile",
                cache_only=True,
            )
        return await memory.query_prior_knowledge(
            query,
            mode="profile",
            bypass_cache=True,
        )


def format_evidence_bundle(bundle: EvidenceBundle) -> str:
    """Render explicit prompt channels; never merge memory into source facts."""

    source_lines = []
    for index, item in enumerate(bundle.source_evidence, 1):
        unit = item.evidence
        location = f"p.{unit.page_start}" if unit.page_start == unit.page_end else f"pp.{unit.page_start}-{unit.page_end}"
        section = " > ".join(unit.section_path)
        source_lines.append(
            f"[E{index} | {location} | {unit.element_type.value}"
            f"{f' | {section}' if section else ''}]\n{unit.index_text}"
        )
    project_lines = [f"- {item.statement}" for item in bundle.project_memory]
    student_lines = [f"- {item.statement}" for item in bundle.student_memory]
    return (
        "<SOURCE_EVIDENCE>\n"
        + ("\n\n".join(source_lines) or "No source evidence retrieved.")
        + "\n</SOURCE_EVIDENCE>\n\n"
        + "<PROJECT_MEMORY>\n"
        + ("\n".join(project_lines) or "No relevant project memory.")
        + "\n</PROJECT_MEMORY>\n\n"
        + "<STUDENT_MEMORY>\n"
        + ("\n".join(student_lines) or "No relevant student memory.")
        + "\n</STUDENT_MEMORY>"
    )