infinity-context 0.1.0__py3-none-any.whl

infinity_context_adapters/postgres/fact_repositories.py

@@ -0,0 +1,766 @@
+"""Postgres repositories for facts and fact relations."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from infinity_context_core.domain.entities import MemoryFact, MemoryFactRelation, SourceRef
+from infinity_context_core.domain.errors import MemoryConflictError, MemoryNotFoundError
+from infinity_context_core.ports.repositories import (
+    ActiveFactBatchRepositoryPort,
+    ActiveFactSearch,
+    FactRelationRepositoryPort,
+    FactRepositoryPort,
+)
+from sqlalchemy import delete, func, or_, select, union, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from infinity_context_adapters.postgres.mappers import (
+    fact_relation_row_to_domain,
+    fact_relation_to_row,
+    fact_row_to_domain,
+    source_ref_to_json,
+)
+from infinity_context_adapters.postgres.models import (
+    MemoryFactRelationRow,
+    MemoryFactRow,
+    MemoryFactVersionRow,
+    MemorySourceRefRow,
+)
+from infinity_context_adapters.postgres.repository_helpers import (
+    _not_expired,
+    _retrieval_candidate_limit,
+    _score,
+    _source_ref_points_to_deleted_document,
+    _tags_match,
+    _terms,
+)
+
+_MAX_FACT_HYDRATION_BINDS = 900
+
+
+class PostgresFactRepository(FactRepositoryPort, ActiveFactBatchRepositoryPort):
+    def __init__(self, session: AsyncSession, *, now: datetime | None = None) -> None:
+        self._session = session
+        self._now = now
+
+    async def create(self, fact: MemoryFact) -> MemoryFact:
+        row = MemoryFactRow(
+            id=str(fact.id),
+            space_id=str(fact.space_id),
+            memory_scope_id=str(fact.memory_scope_id),
+            thread_id=str(fact.thread_id) if fact.thread_id else None,
+            kind=fact.kind.value,
+            text=fact.text,
+            status=fact.status.value,
+            confidence=fact.confidence.value,
+            trust_level=fact.trust_level.value,
+            classification=fact.classification,
+            category=fact.category,
+            tags_json=list(fact.tags),
+            ttl_policy=fact.ttl_policy,
+            expires_at=fact.expires_at,
+            version=fact.version,
+            created_at=fact.created_at,
+            updated_at=fact.updated_at,
+        )
+        self._session.add(row)
+        await self._write_version(fact)
+        await self._replace_source_refs(fact)
+        return fact
+
+    async def get_by_id(self, fact_id: str) -> MemoryFact | None:
+        row = await self._session.get(MemoryFactRow, fact_id)
+        if row is None:
+            return None
+        refs = await self._load_source_refs(fact_id=fact_id, version=row.version)
+        return fact_row_to_domain(row, refs)
+
+    async def get_by_ids(self, fact_ids: tuple[str, ...]) -> list[MemoryFact]:
+        unique_ids = tuple(dict.fromkeys(fact_id for fact_id in fact_ids if fact_id.strip()))
+        if not unique_ids:
+            return []
+        rows = list(
+            (
+                await self._session.execute(
+                    select(MemoryFactRow).where(MemoryFactRow.id.in_(unique_ids))
+                )
+            ).scalars()
+        )
+        if not rows:
+            return []
+        rows_by_id = {row.id: row for row in rows}
+        ref_rows = list(
+            (
+                await self._session.execute(
+                    select(MemorySourceRefRow)
+                    .where(MemorySourceRefRow.fact_id.in_(tuple(rows_by_id)))
+                    .order_by(
+                        MemorySourceRefRow.fact_id,
+                        MemorySourceRefRow.fact_version,
+                        MemorySourceRefRow.id,
+                    )
+                )
+            ).scalars()
+        )
+        refs_by_fact_version: dict[tuple[str, int], list[MemorySourceRefRow]] = {}
+        for ref in ref_rows:
+            refs_by_fact_version.setdefault((ref.fact_id, ref.fact_version), []).append(ref)
+        return [
+            fact_row_to_domain(
+                row,
+                refs_by_fact_version.get((row.id, row.version), []),
+            )
+            for fact_id in unique_ids
+            if (row := rows_by_id.get(fact_id)) is not None
+        ]
+
+    async def get_for_update(self, fact_id: str) -> MemoryFact | None:
+        statement = select(MemoryFactRow).where(MemoryFactRow.id == fact_id).with_for_update()
+        row = (await self._session.execute(statement)).scalar_one_or_none()
+        if row is None:
+            return None
+        refs = await self._load_source_refs(fact_id=fact_id, version=row.version)
+        return fact_row_to_domain(row, refs)
+
+    async def save(self, fact: MemoryFact) -> MemoryFact:
+        expected_version = fact.version - 1
+        if expected_version < 1:
+            raise MemoryConflictError("Stale fact version")
+        result = await self._session.execute(
+            update(MemoryFactRow)
+            .where(
+                MemoryFactRow.id == str(fact.id),
+                MemoryFactRow.version == expected_version,
+            )
+            .values(
+                space_id=str(fact.space_id),
+                memory_scope_id=str(fact.memory_scope_id),
+                thread_id=str(fact.thread_id) if fact.thread_id else None,
+                kind=fact.kind.value,
+                text=fact.text,
+                status=fact.status.value,
+                confidence=fact.confidence.value,
+                trust_level=fact.trust_level.value,
+                classification=fact.classification,
+                category=fact.category,
+                tags_json=list(fact.tags),
+                ttl_policy=fact.ttl_policy,
+                expires_at=fact.expires_at,
+                version=fact.version,
+                created_at=fact.created_at,
+                updated_at=fact.updated_at,
+            )
+        )
+        if result.rowcount == 0:
+            exists = await self._session.get(MemoryFactRow, str(fact.id))
+            if exists is None:
+                msg = "Fact row missing during save"
+                raise RuntimeError(msg)
+            if (
+                exists.version == fact.version
+                and exists.status == fact.status.value
+                and exists.status == "deleted"
+            ):
+                return fact
+            raise MemoryConflictError("Stale fact version")
+        row = await self._session.get(MemoryFactRow, str(fact.id))
+        if row is not None:
+            self._session.expire(row)
+        else:
+            msg = "Fact row missing during save"
+            raise RuntimeError(msg)
+        await self._write_version(fact)
+        await self._replace_source_refs(fact)
+        return fact
+
+    async def list_versions(self, fact_id: str) -> list[MemoryFact]:
+        rows = (
+            await self._session.execute(
+                select(MemoryFactVersionRow)
+                .where(MemoryFactVersionRow.fact_id == fact_id)
+                .order_by(MemoryFactVersionRow.version)
+            )
+        ).scalars()
+        current = await self.get_by_id(fact_id)
+        if current is None:
+            return []
+        versions: list[MemoryFact] = []
+        for version_row in rows:
+            version_fact = MemoryFact(
+                id=current.id,
+                space_id=current.space_id,
+                memory_scope_id=current.memory_scope_id,
+                thread_id=current.thread_id,
+                text=version_row.text,
+                kind=current.kind,
+                source_refs=tuple(
+                    SourceRef(
+                        source_type=str(ref["source_type"]),
+                        source_id=str(ref["source_id"]),
+                        chunk_id=str(ref["chunk_id"]) if ref.get("chunk_id") else None,
+                        char_start=(
+                            int(ref["char_start"]) if ref.get("char_start") is not None else None
+                        ),
+                        char_end=int(ref["char_end"]) if ref.get("char_end") is not None else None,
+                        quote_preview=str(ref["quote_preview"])
+                        if ref.get("quote_preview") is not None
+                        else None,
+                    )
+                    for ref in version_row.source_refs_json
+                ),
+                status=current.status.__class__(version_row.status),
+                version=version_row.version,
+                confidence=current.confidence,
+                trust_level=current.trust_level,
+                classification=current.classification,
+                category=current.category,
+                tags=current.tags,
+                ttl_policy=current.ttl_policy,
+                expires_at=current.expires_at,
+                created_at=current.created_at,
+                updated_at=version_row.created_at,
+            )
+            versions.append(version_fact)
+        return versions
+
+    async def find_active(
+        self,
+        *,
+        space_id: str,
+        memory_scope_ids: tuple[str, ...],
+        thread_id: str | None,
+        query: str,
+        limit: int,
+        category: str | None = None,
+        tags_any: tuple[str, ...] = (),
+        tags_all: tuple[str, ...] = (),
+        tags_none: tuple[str, ...] = (),
+    ) -> list[MemoryFact]:
+        (facts,) = await self.find_active_many(
+            (
+                ActiveFactSearch(
+                    space_id=space_id,
+                    memory_scope_ids=memory_scope_ids,
+                    thread_id=thread_id,
+                    query=query,
+                    limit=limit,
+                    category=category,
+                    tags_any=tags_any,
+                    tags_all=tags_all,
+                    tags_none=tags_none,
+                ),
+            )
+        )
+        return facts
+
+    async def find_active_many(
+        self,
+        searches: tuple[ActiveFactSearch, ...],
+    ) -> list[list[MemoryFact]]:
+        if not searches:
+            return []
+        grouped: dict[tuple[object, ...], list[tuple[int, ActiveFactSearch]]] = {}
+        for index, search in enumerate(searches):
+            grouped.setdefault(_active_fact_search_group_key(search), []).append((index, search))
+        results: list[list[MemoryFact]] = [[] for _ in searches]
+        for indexed_searches in grouped.values():
+            max_candidate_limit = max(
+                (_retrieval_candidate_limit(search.limit) for _, search in indexed_searches),
+                default=0,
+            )
+            if max_candidate_limit <= 0:
+                continue
+            first = indexed_searches[0][1]
+            rows = list(
+                (
+                    await self._session.execute(
+                        select(MemoryFactRow)
+                        .where(*_active_fact_conditions(first, now=self._now))
+                        .order_by(MemoryFactRow.updated_at.desc())
+                        .limit(max_candidate_limit)
+                    )
+                ).scalars()
+            )
+            selected_ids: dict[int, tuple[str, ...]] = {}
+            hydration_ids: list[str] = []
+            for index, search in indexed_searches:
+                candidate_rows = rows[: _retrieval_candidate_limit(search.limit)]
+                if search.tags_any or search.tags_all or search.tags_none:
+                    candidate_rows = [
+                        row
+                        for row in candidate_rows
+                        if _tags_match(
+                            row.tags_json or [],
+                            tags_any=search.tags_any,
+                            tags_all=search.tags_all,
+                            tags_none=search.tags_none,
+                        )
+                    ]
+                ids = tuple(
+                    row.id
+                    for row in _rank_active_fact_rows(
+                        candidate_rows,
+                        query=search.query,
+                        limit=search.limit,
+                    )
+                )
+                selected_ids[index] = ids
+                hydration_ids.extend(ids)
+            facts_by_id = await _hydrate_fact_rows_by_ids(
+                self._session,
+                tuple(dict.fromkeys(hydration_ids)),
+            )
+            for index, search in indexed_searches:
+                results[index] = [
+                    fact
+                    for fact_id in selected_ids[index]
+                    if (fact := facts_by_id.get(fact_id)) is not None
+                    and _active_fact_matches_search(fact, search, now=self._now)
+                ]
+        return results
+
+    async def list_for_scope(
+        self,
+        *,
+        space_id: str,
+        memory_scope_id: str,
+        thread_id: str | None,
+        status: str | None,
+        limit: int,
+        cursor_updated_at: datetime | None = None,
+        cursor_id: str | None = None,
+        category: str | None = None,
+        tag: str | None = None,
+    ) -> list[MemoryFact]:
+        conditions = [
+            MemoryFactRow.space_id == space_id,
+            MemoryFactRow.memory_scope_id == memory_scope_id,
+        ]
+        if status:
+            conditions.append(MemoryFactRow.status == status)
+            if status == "active":
+                conditions.append(_not_expired(MemoryFactRow, self._now))
+        if category:
+            conditions.append(MemoryFactRow.category == category)
+        if thread_id is not None:
+            conditions.append(
+                or_(MemoryFactRow.thread_id == thread_id, MemoryFactRow.thread_id.is_(None))
+            )
+        if cursor_updated_at is not None and cursor_id is not None:
+            conditions.append(
+                or_(
+                    MemoryFactRow.updated_at < cursor_updated_at,
+                    (MemoryFactRow.updated_at == cursor_updated_at)
+                    & (MemoryFactRow.id < cursor_id),
+                )
+            )
+        rows = list(
+            (
+                await self._session.execute(
+                    select(MemoryFactRow)
+                    .where(*conditions)
+                    .order_by(MemoryFactRow.updated_at.desc(), MemoryFactRow.id.desc())
+                    .limit(_retrieval_candidate_limit(limit) if tag else limit)
+                )
+            ).scalars()
+        )
+        facts = []
+        for row in rows:
+            refs = await self._load_source_refs(fact_id=row.id, version=row.version)
+            fact = fact_row_to_domain(row, refs)
+            if tag and tag not in fact.tags:
+                continue
+            facts.append(fact)
+        return facts
+
+    async def delete_facts_sourced_only_by_chunks(
+        self,
+        *,
+        space_id: str,
+        memory_scope_id: str,
+        document_id: str,
+        chunk_ids: tuple[str, ...],
+        now: datetime,
+    ) -> tuple[tuple[str, int], ...]:
+        if not chunk_ids and not document_id:
+            return ()
+        chunk_id_set = set(chunk_ids)
+        candidate_rows = list(
+            (
+                await self._session.execute(
+                    select(MemoryFactRow)
+                    .join(
+                        MemorySourceRefRow,
+                        (MemorySourceRefRow.fact_id == MemoryFactRow.id)
+                        & (MemorySourceRefRow.fact_version == MemoryFactRow.version),
+                    )
+                    .where(
+                        MemoryFactRow.status == "active",
+                        MemoryFactRow.space_id == space_id,
+                        MemoryFactRow.memory_scope_id == memory_scope_id,
+                        or_(
+                            MemorySourceRefRow.chunk_id.in_(chunk_id_set),
+                            (
+                                (MemorySourceRefRow.source_type == "document")
+                                & (MemorySourceRefRow.source_id == document_id)
+                            ),
+                        ),
+                    )
+                    .distinct()
+                )
+            ).scalars()
+        )
+        deleted: list[tuple[str, int]] = []
+        for row in candidate_rows:
+            refs = await self._load_source_refs(fact_id=row.id, version=row.version)
+            if refs and all(
+                _source_ref_points_to_deleted_document(
+                    ref,
+                    document_id=document_id,
+                    chunk_ids=chunk_id_set,
+                )
+                for ref in refs
+            ):
+                forgotten = fact_row_to_domain(row, refs).forget(now=now)
+                await self.save(forgotten)
+                deleted.append((str(forgotten.id), forgotten.version))
+        return tuple(deleted)
+
+    async def _write_version(self, fact: MemoryFact) -> None:
+        existing = (
+            await self._session.execute(
+                select(MemoryFactVersionRow).where(
+                    MemoryFactVersionRow.fact_id == str(fact.id),
+                    MemoryFactVersionRow.version == fact.version,
+                )
+            )
+        ).scalar_one_or_none()
+        if existing is not None:
+            existing.text = fact.text
+            existing.status = fact.status.value
+            existing.source_refs_json = [source_ref_to_json(ref) for ref in fact.source_refs]
+            existing.created_at = fact.updated_at
+            return
+        self._session.add(
+            MemoryFactVersionRow(
+                fact_id=str(fact.id),
+                version=fact.version,
+                text=fact.text,
+                status=fact.status.value,
+                source_refs_json=[source_ref_to_json(ref) for ref in fact.source_refs],
+                reason=None,
+                created_at=fact.updated_at,
+            )
+        )
+
+    async def _replace_source_refs(self, fact: MemoryFact) -> None:
+        await self._session.execute(
+            delete(MemorySourceRefRow).where(
+                MemorySourceRefRow.fact_id == str(fact.id),
+                MemorySourceRefRow.fact_version == fact.version,
+            )
+        )
+        for ref in fact.source_refs:
+            self._session.add(
+                MemorySourceRefRow(
+                    fact_id=str(fact.id),
+                    fact_version=fact.version,
+                    source_type=ref.source_type,
+                    source_id=ref.source_id,
+                    chunk_id=ref.chunk_id,
+                    char_start=ref.char_start,
+                    char_end=ref.char_end,
+                    quote_preview=ref.quote_preview,
+                    page_number=ref.page_number,
+                    time_start_ms=ref.time_start_ms,
+                    time_end_ms=ref.time_end_ms,
+                    bbox_json=list(ref.bbox) if ref.bbox is not None else None,
+                )
+            )
+
+    async def _load_source_refs(self, *, fact_id: str, version: int) -> list[MemorySourceRefRow]:
+        return list(
+            (
+                await self._session.execute(
+                    select(MemorySourceRefRow)
+                    .where(
+                        MemorySourceRefRow.fact_id == fact_id,
+                        MemorySourceRefRow.fact_version == version,
+                    )
+                    .order_by(MemorySourceRefRow.id)
+                )
+            ).scalars()
+        )
+
+
+def _active_fact_conditions(
+    search: ActiveFactSearch,
+    *,
+    now: datetime | None,
+) -> tuple[object, ...]:
+    conditions = [
+        MemoryFactRow.space_id == search.space_id,
+        MemoryFactRow.memory_scope_id.in_(search.memory_scope_ids),
+        MemoryFactRow.status == "active",
+        MemoryFactRow.classification != "restricted",
+        _not_expired(MemoryFactRow, now),
+    ]
+    if search.category:
+        conditions.append(MemoryFactRow.category == search.category)
+    if search.thread_id is not None:
+        conditions.append(
+            or_(
+                MemoryFactRow.thread_id == search.thread_id,
+                MemoryFactRow.thread_id.is_(None),
+            )
+        )
+    return tuple(conditions)
+
+
+def _active_fact_search_group_key(search: ActiveFactSearch) -> tuple[object, ...]:
+    return (
+        search.space_id,
+        search.memory_scope_ids,
+        search.thread_id,
+        search.category,
+        search.tags_any,
+        search.tags_all,
+        search.tags_none,
+    )
+
+
+def _rank_active_fact_rows(
+    rows: list[MemoryFactRow],
+    *,
+    query: str,
+    limit: int,
+) -> list[MemoryFactRow]:
+    if limit <= 0:
+        return []
+    terms = _terms(query)
+    if not terms:
+        return rows[:limit]
+    scored_rows = [
+        (score, index, row)
+        for index, row in enumerate(rows)
+        for score in (_score(row.text, terms),)
+        if score > 0
+    ]
+    scored_rows.sort(key=lambda item: (-item[0], item[1]))
+    return [row for _, _, row in scored_rows[:limit]]
+
+
+async def _hydrate_fact_rows_by_ids(
+    session: AsyncSession,
+    fact_ids: tuple[str, ...],
+) -> dict[str, MemoryFact]:
+    rows_by_id: dict[str, MemoryFactRow] = {}
+    refs_by_fact_version: dict[tuple[str, int], list[MemorySourceRefRow]] = {}
+    for fact_id_batch in _batches(fact_ids, _MAX_FACT_HYDRATION_BINDS):
+        rows = list(
+            (
+                await session.execute(
+                    select(MemoryFactRow)
+                    .where(MemoryFactRow.id.in_(fact_id_batch))
+                    .execution_options(populate_existing=True)
+                )
+            ).scalars()
+        )
+        rows_by_id.update((row.id, row) for row in rows)
+        if not rows:
+            continue
+        refs = list(
+            (
+                await session.execute(
+                    select(MemorySourceRefRow)
+                    .where(MemorySourceRefRow.fact_id.in_(tuple(row.id for row in rows)))
+                    .order_by(
+                        MemorySourceRefRow.fact_id,
+                        MemorySourceRefRow.fact_version,
+                        MemorySourceRefRow.id,
+                    )
+                )
+            ).scalars()
+        )
+        for ref in refs:
+            refs_by_fact_version.setdefault((ref.fact_id, ref.fact_version), []).append(ref)
+    return {
+        fact_id: fact_row_to_domain(
+            row,
+            refs_by_fact_version.get((row.id, row.version), []),
+        )
+        for fact_id in fact_ids
+        if (row := rows_by_id.get(fact_id)) is not None
+    }
+
+
+def _active_fact_matches_search(
+    fact: MemoryFact,
+    search: ActiveFactSearch,
+    *,
+    now: datetime | None,
+) -> bool:
+    comparable_now = now or datetime.now(UTC)
+    expires_at = fact.expires_at
+    if expires_at is not None:
+        if expires_at.tzinfo is None and comparable_now.tzinfo is not None:
+            comparable_now = comparable_now.replace(tzinfo=None)
+        if expires_at <= comparable_now:
+            return False
+    return (
+        str(fact.space_id) == search.space_id
+        and str(fact.memory_scope_id) in search.memory_scope_ids
+        and fact.status.value == "active"
+        and fact.classification != "restricted"
+        and (not search.category or fact.category == search.category)
+        and (
+            search.thread_id is None
+            or fact.thread_id is None
+            or str(fact.thread_id) == search.thread_id
+        )
+        and _tags_match(
+            list(fact.tags),
+            tags_any=search.tags_any,
+            tags_all=search.tags_all,
+            tags_none=search.tags_none,
+        )
+    )
+
+
+def _batches(values: tuple[str, ...], size: int) -> tuple[tuple[str, ...], ...]:
+    return tuple(
+        values[offset : offset + size]
+        for offset in range(0, len(values), size)
+    )
+
+
+class PostgresFactRelationRepository(FactRelationRepositoryPort):
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def create(self, relation: MemoryFactRelation) -> MemoryFactRelation:
+        self._session.add(fact_relation_to_row(relation))
+        return relation
+
+    async def get_by_id(self, relation_id: str) -> MemoryFactRelation | None:
+        row = await self._session.get(MemoryFactRelationRow, relation_id)
+        return fact_relation_row_to_domain(row) if row else None
+
+    async def save(self, relation: MemoryFactRelation) -> MemoryFactRelation:
+        row = await self._session.get(MemoryFactRelationRow, str(relation.id))
+        if row is None:
+            raise MemoryNotFoundError("Fact relation not found")
+        row.status = relation.status.value
+        row.reason = relation.reason
+        row.observed_at = relation.observed_at
+        row.valid_from = relation.valid_from
+        row.valid_to = relation.valid_to
+        row.updated_at = relation.updated_at
+        return relation
+
+    async def find_active(
+        self,
+        *,
+        source_fact_id: str,
+        target_fact_id: str,
+        relation_type: str,
+    ) -> MemoryFactRelation | None:
+        row = (
+            await self._session.execute(
+                select(MemoryFactRelationRow).where(
+                    MemoryFactRelationRow.source_fact_id == source_fact_id,
+                    MemoryFactRelationRow.target_fact_id == target_fact_id,
+                    MemoryFactRelationRow.relation_type == relation_type,
+                    MemoryFactRelationRow.status == "active",
+                )
+            )
+        ).scalar_one_or_none()
+        return fact_relation_row_to_domain(row) if row else None
+
+    async def list_for_fact(
+        self,
+        *,
+        fact_id: str,
+        status: str | None,
+        limit: int,
+    ) -> list[MemoryFactRelation]:
+        conditions = [
+            or_(
+                MemoryFactRelationRow.source_fact_id == fact_id,
+                MemoryFactRelationRow.target_fact_id == fact_id,
+            )
+        ]
+        if status is not None:
+            conditions.append(MemoryFactRelationRow.status == status)
+        rows = (
+            await self._session.execute(
+                select(MemoryFactRelationRow)
+                .where(*conditions)
+                .order_by(MemoryFactRelationRow.updated_at.desc(), MemoryFactRelationRow.id.desc())
+                .limit(limit)
+            )
+        ).scalars()
+        return [fact_relation_row_to_domain(row) for row in rows]
+
+    async def list_for_facts(
+        self,
+        *,
+        fact_ids: tuple[str, ...],
+        status: str | None,
+        limit_per_fact: int,
+    ) -> dict[str, list[MemoryFactRelation]]:
+        unique_fact_ids = tuple(dict.fromkeys(str(fact_id) for fact_id in fact_ids if fact_id))
+        if not unique_fact_ids:
+            return {}
+        safe_limit_per_fact = max(0, int(limit_per_fact))
+        if safe_limit_per_fact <= 0:
+            return {fact_id: [] for fact_id in unique_fact_ids}
+        source_conditions = [MemoryFactRelationRow.source_fact_id.in_(unique_fact_ids)]
+        target_conditions = [MemoryFactRelationRow.target_fact_id.in_(unique_fact_ids)]
+        if status is not None:
+            source_conditions.append(MemoryFactRelationRow.status == status)
+            target_conditions.append(MemoryFactRelationRow.status == status)
+        relation_matches = union(
+            select(
+                MemoryFactRelationRow.source_fact_id.label("fact_id"),
+                MemoryFactRelationRow.id.label("relation_id"),
+                MemoryFactRelationRow.updated_at.label("updated_at"),
+            ).where(*source_conditions),
+            select(
+                MemoryFactRelationRow.target_fact_id.label("fact_id"),
+                MemoryFactRelationRow.id.label("relation_id"),
+                MemoryFactRelationRow.updated_at.label("updated_at"),
+            ).where(*target_conditions),
+        ).subquery()
+        ranked_matches = select(
+            relation_matches.c.fact_id,
+            relation_matches.c.relation_id,
+            func.row_number()
+            .over(
+                partition_by=relation_matches.c.fact_id,
+                order_by=(
+                    relation_matches.c.updated_at.desc(),
+                    relation_matches.c.relation_id.desc(),
+                ),
+            )
+            .label("relation_rank"),
+        ).subquery()
+        rows = (
+            await self._session.execute(
+                select(ranked_matches.c.fact_id, MemoryFactRelationRow)
+                .join(
+                    MemoryFactRelationRow,
+                    MemoryFactRelationRow.id == ranked_matches.c.relation_id,
+                )
+                .where(ranked_matches.c.relation_rank <= safe_limit_per_fact)
+                .order_by(ranked_matches.c.fact_id, ranked_matches.c.relation_rank)
+            )
+        ).all()
+        relations_by_fact_id: dict[str, list[MemoryFactRelation]] = {
+            fact_id: [] for fact_id in unique_fact_ids
+        }
+        for fact_id, row in rows:
+            relation = fact_relation_row_to_domain(row)
+            relations_by_fact_id[str(fact_id)].append(relation)
+        return relations_by_fact_id