topoteretes--cognee/cognee/modules/retrieval/utils/global_context.py

from typing import Any, Optional

from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
from cognee.modules.retrieval.utils.brute_force_triplet_search import get_memory_fragment

# Vector collection for the persisted GlobalContextSummary.text index.
GLOBAL_CONTEXT_SUMMARY_COLLECTION = "GlobalContextSummary_text"


async def load_root_text() -> Optional[str]:
    """
    Return the dataset's root GlobalContextSummary text, or None when no root
    exists in scope. Relies on backend access control to scope the query to
    one dataset; within scope a dataset has at most one root.
    """
    fragment = await get_memory_fragment(
        properties_to_project=["id", "text", "type", "is_root"],
        memory_fragment_filter=[{"type": ["GlobalContextSummary"]}],
    )
    for node in fragment.nodes.values():
        if _is_root(node.attributes):
            text = node.attributes.get("text")
            return text or None
    return None


async def search_top_global_context_summaries(
    query: str,
    top_k: int,
    vector_engine: Any,
    query_vector: Optional[list[float]] = None,
) -> list[str]:
    """Return up to top_k non-root GlobalContextSummary texts ranked by similarity."""
    if top_k <= 0:
        return []
    try:
        results = await vector_engine.search(
            GLOBAL_CONTEXT_SUMMARY_COLLECTION,
            None if query_vector is not None else query,
            query_vector=query_vector,
            limit=top_k + 1,  # +1 in case the root is in the top results
            include_payload=True,
        )
    except CollectionNotFoundError:
        return []

    summaries: list[str] = []
    for result in results or []:
        payload = getattr(result, "payload", None) or {}
        if _is_root(payload):
            continue
        text = payload.get("text")
        if text:
            summaries.append(text)
        if len(summaries) >= top_k:
            break
    return summaries


def format_global_context_prelude(
    root_text: Optional[str],
    top_summaries: list[str],
) -> str:
    """Build the prepend string. Empty when neither input has content."""
    blocks: list[str] = []
    if root_text:
        blocks.append(f"World summary:\n{root_text}")
    if top_summaries:
        joined = "\n\n".join(top_summaries)
        blocks.append(f"Relevant areas:\n{joined}")
    return "\n\n".join(blocks)


def _is_root(attributes: dict) -> bool:
    value = attributes.get("is_root")
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        return value.lower() == "true"
    return bool(value)