topoteretes--cognee/cognee/modules/agent_memory/runtime.py

from __future__ import annotations

import contextvars
import inspect
import json
from dataclasses import dataclass
from typing import Any, Optional
from uuid import UUID

from cognee.exceptions import CogneeValidationError
from cognee.modules.observability import new_span
from cognee.modules.users.methods import get_default_user
from cognee.modules.users.models import User
from cognee.modules.users.permissions.methods import get_all_user_permission_datasets
from cognee.shared.logging_utils import get_logger
from cognee.modules.agent_memory.sanitization import (
    MAX_SERIALIZED_VALUE_LENGTH,
    sanitize_value,
    truncate_text,
)

logger = get_logger("agent_memory")

MAX_MEMORY_CONTEXT_LENGTH = 4000


@dataclass(slots=True)
class AgentMemoryConfig:
    """Validated decorator configuration used for one wrapped agent invocation."""

    with_memory: bool
    with_session_memory: bool
    save_session_traces: bool
    memory_query_fixed: Optional[str]
    memory_query_from_method: Optional[str]
    memory_system_prompt: Optional[str]
    memory_top_k: int
    memory_only_context: bool
    session_memory_last_n: int
    session_id: Optional[str]
    user: Optional[User]
    dataset_name: Optional[str]
    session_trace_summary: bool
    persist_session_trace_after: Optional[int]
    persist_session_trace_raw_content: bool
    persist_session_trace_node_set_name: Optional[str]


@dataclass(slots=True)
class AgentScope:
    """Authorized dataset scope resolved for Cognee-backed memory retrieval."""

    user: User
    dataset_name: str
    dataset_id: UUID


@dataclass(slots=True)
class AgentMemoryContext:
    """Per-call execution state shared across retrieval, wrapped call, and trace persistence."""

    origin_function: str
    config: AgentMemoryConfig
    method_params: dict[str, Any]
    user: Optional[User] = None
    scope: Optional[AgentScope] = None
    memory_query: str = ""
    memory_context: str = ""
    method_return_value: Any = None
    status: str = "running"
    error_message: str = ""


_agent_memory_context_var: contextvars.ContextVar[Optional[AgentMemoryContext]] = (
    contextvars.ContextVar("agent_memory_context", default=None)
)


def get_current_agent_memory_context() -> Optional[AgentMemoryContext]:
    """Return the active agent-memory execution context for the current async task."""
    return _agent_memory_context_var.get()


def set_current_agent_memory_context(
    context: AgentMemoryContext,
) -> contextvars.Token[Optional[AgentMemoryContext]]:
    """Store the active agent-memory context and return a reset token."""
    return _agent_memory_context_var.set(context)


def reset_current_agent_memory_context(
    token: contextvars.Token[Optional[AgentMemoryContext]],
) -> None:
    """Restore the previously active agent-memory context."""
    _agent_memory_context_var.reset(token)


def validate_agent_memory_config(
    *,
    with_memory: bool,
    with_session_memory: bool,
    save_session_traces: bool,
    memory_query_fixed: Optional[str],
    memory_query_from_method: Optional[str],
    memory_system_prompt: Optional[str],
    memory_top_k: int,
    memory_only_context: bool,
    session_memory_last_n: int,
    session_id: Optional[str],
    user: Optional[User],
    dataset_name: Optional[str],
    session_trace_summary: bool,
    persist_session_trace_after: Optional[int],
    persist_session_trace_raw_content: bool,
    persist_session_trace_node_set_name: Optional[str],
) -> AgentMemoryConfig:
    """Validate and normalize the public decorator configuration."""
    from cognee.infrastructure.databases.cache.config import get_cache_config

    if not isinstance(with_memory, bool):
        raise CogneeValidationError("with_memory must be a boolean.", log=False)
    if not isinstance(with_session_memory, bool):
        raise CogneeValidationError("with_session_memory must be a boolean.", log=False)
    if not isinstance(save_session_traces, bool):
        raise CogneeValidationError("save_session_traces must be a boolean.", log=False)
    if not isinstance(memory_only_context, bool):
        raise CogneeValidationError("memory_only_context must be a boolean.", log=False)
    if not isinstance(session_trace_summary, bool):
        raise CogneeValidationError("session_trace_summary must be a boolean.", log=False)
    if not isinstance(persist_session_trace_raw_content, bool):
        raise CogneeValidationError(
            "persist_session_trace_raw_content must be a boolean.",
            log=False,
        )
    if persist_session_trace_node_set_name is not None and not isinstance(
        persist_session_trace_node_set_name, str
    ):
        raise CogneeValidationError(
            "persist_session_trace_node_set_name must be a string when provided.",
            log=False,
        )
    if memory_query_fixed is not None and not isinstance(memory_query_fixed, str):
        raise CogneeValidationError("memory_query_fixed must be a string when provided.", log=False)
    if memory_query_from_method is not None and not isinstance(memory_query_from_method, str):
        raise CogneeValidationError(
            "memory_query_from_method must be a string when provided.",
            log=False,
        )
    if memory_system_prompt is not None and not isinstance(memory_system_prompt, str):
        raise CogneeValidationError(
            "memory_system_prompt must be a string when provided.",
            log=False,
        )
    if memory_query_fixed is not None and not memory_query_fixed.strip():
        raise CogneeValidationError(
            "memory_query_fixed must not be blank when provided.",
            log=False,
        )
    if memory_query_from_method is not None and not memory_query_from_method.strip():
        raise CogneeValidationError(
            "memory_query_from_method must not be blank when provided.",
            log=False,
        )
    if memory_system_prompt is not None and not memory_system_prompt.strip():
        raise CogneeValidationError(
            "memory_system_prompt must not be blank when provided.",
            log=False,
        )
    if (
        persist_session_trace_node_set_name is not None
        and not persist_session_trace_node_set_name.strip()
    ):
        raise CogneeValidationError(
            "persist_session_trace_node_set_name must not be blank when provided.",
            log=False,
        )
    if memory_query_fixed is not None and memory_query_from_method is not None:
        raise CogneeValidationError(
            "Only one of memory_query_fixed or memory_query_from_method can be provided to cognee.agent_memory.",
            log=False,
        )
    if not isinstance(session_memory_last_n, int) or session_memory_last_n < 1:
        raise CogneeValidationError(
            "session_memory_last_n must be a positive integer.",
            log=False,
        )
    if persist_session_trace_after is not None and (
        not isinstance(persist_session_trace_after, int) or persist_session_trace_after < 1
    ):
        raise CogneeValidationError(
            "persist_session_trace_after must be a positive integer when provided.",
            log=False,
        )
    if persist_session_trace_after is not None and not save_session_traces:
        raise CogneeValidationError(
            "persist_session_trace_after requires save_session_traces=True.",
            log=False,
        )
    cache_config = get_cache_config()
    if not cache_config.caching and (
        with_session_memory or save_session_traces or persist_session_trace_after is not None
    ):
        raise CogneeValidationError(
            (
                "Caching must be enabled to use with_session_memory, save_session_traces, "
                "or persist_session_trace_after with cognee.agent_memory."
            ),
            log=False,
        )
    if persist_session_trace_after is not None and (
        not isinstance(persist_session_trace_after, int) or persist_session_trace_after < 1
    ):
        raise CogneeValidationError(
            "persist_session_trace_after must be a positive integer when provided.",
            log=False,
        )
    if persist_session_trace_after is not None and not save_session_traces:
        raise CogneeValidationError(
            "persist_session_trace_after requires save_session_traces=True.",
            log=False,
        )
    cache_config = get_cache_config()
    if not cache_config.caching and (
        with_session_memory or save_session_traces or persist_session_trace_after is not None
    ):
        raise CogneeValidationError(
            (
                "Caching must be enabled to use with_session_memory, save_session_traces, "
                "or persist_session_trace_after with cognee.agent_memory."
            ),
            log=False,
        )
    if session_id is not None and (not isinstance(session_id, str) or not session_id.strip()):
        raise CogneeValidationError(
            "session_id must be a non-empty string when provided.",
            log=False,
        )
    if user is not None and not hasattr(user, "id"):
        raise CogneeValidationError("user must have an id attribute.", log=False)
    if dataset_name is not None and (not isinstance(dataset_name, str) or not dataset_name.strip()):
        raise CogneeValidationError(
            "dataset_name must be a non-empty string when provided.",
            log=False,
        )

    return AgentMemoryConfig(
        with_memory=with_memory,
        with_session_memory=with_session_memory,
        save_session_traces=save_session_traces,
        memory_query_fixed=(
            memory_query_fixed.strip() if isinstance(memory_query_fixed, str) else None
        ),
        memory_query_from_method=(
            memory_query_from_method.strip() if isinstance(memory_query_from_method, str) else None
        ),
        memory_system_prompt=(
            memory_system_prompt.strip() if isinstance(memory_system_prompt, str) else None
        ),
        memory_top_k=memory_top_k,
        memory_only_context=memory_only_context,
        session_memory_last_n=session_memory_last_n,
        session_id=session_id.strip() if isinstance(session_id, str) else None,
        user=user,
        dataset_name=dataset_name.strip() if isinstance(dataset_name, str) else None,
        session_trace_summary=session_trace_summary,
        persist_session_trace_after=persist_session_trace_after,
        persist_session_trace_raw_content=persist_session_trace_raw_content,
        persist_session_trace_node_set_name=(
            persist_session_trace_node_set_name.strip()
            if isinstance(persist_session_trace_node_set_name, str)
            else None
        ),
    )


async def resolve_agent_user(config: AgentMemoryConfig) -> User:
    """Resolve the effective user for agent-memory search/session operations."""
    return config.user or await get_default_user()


async def resolve_agent_dataset_scope(config: AgentMemoryConfig, resolved_user: User) -> AgentScope:
    """Resolve the dataset scope for Cognee search using a user with read and write access."""
    requested_dataset_name = config.dataset_name or "main_dataset"

    readable_datasets = await get_all_user_permission_datasets(resolved_user, "read")
    writable_datasets = await get_all_user_permission_datasets(resolved_user, "write")

    readable_by_id = {dataset.id: dataset for dataset in readable_datasets}
    writable_ids = {dataset.id: dataset for dataset in writable_datasets}
    matching_datasets = [
        dataset
        for dataset in readable_by_id.values()
        if dataset.id in writable_ids and dataset.name == requested_dataset_name
    ]

    if len(matching_datasets) > 1:
        raise CogneeValidationError(
            (
                f"Multiple datasets named {requested_dataset_name!r} grant both read and write "
                f"permissions to user {resolved_user.id}. Please use a unique dataset name."
            ),
            log=False,
        )
    if not matching_datasets:
        raise CogneeValidationError(
            (
                f"User {resolved_user.id} must have both read and write permissions for dataset "
                f"{requested_dataset_name!r} to use cognee.agent_memory."
            ),
            log=False,
        )

    authorized_dataset = matching_datasets[0]

    return AgentScope(
        user=resolved_user,
        dataset_name=authorized_dataset.name,
        dataset_id=authorized_dataset.id,
    )


def build_method_params(func, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
    """Bind wrapped call arguments to parameter names and sanitize them for storage."""
    bound_args = inspect.signature(func).bind_partial(*args, **kwargs)
    bound_args.apply_defaults()
    return {key: sanitize_value(value) for key, value in bound_args.arguments.items()}


def normalize_optional_text(value: Any, limit: int = MAX_SERIALIZED_VALUE_LENGTH) -> Optional[str]:
    """Convert a value into a bounded non-empty string, or return None when unusable."""
    if value is None:
        return None

    if not isinstance(value, str):
        value = str(sanitize_value(value))

    stripped = value.strip()
    if not stripped:
        return None

    return stripped[:limit]


def get_query_text_from_method_param(
    memory_query_from_method: Optional[str],
    method_params: dict[str, Any],
) -> Optional[str]:
    """Extract a bounded retrieval query from a configured wrapped-method parameter."""
    if not memory_query_from_method or memory_query_from_method not in method_params:
        return None

    return normalize_optional_text(method_params[memory_query_from_method])


def derive_query_text(
    memory_query_fixed: Optional[str],
    memory_query_from_method: Optional[str],
    method_params: dict[str, Any],
) -> Optional[str]:
    """Resolve the retrieval query from dynamic, fixed, or fallback method inputs."""
    query_from_method = get_query_text_from_method_param(memory_query_from_method, method_params)
    if query_from_method:
        return query_from_method

    if memory_query_fixed:
        return memory_query_fixed

    for key, value in method_params.items():
        if key in {"user", "dataset_name", "session_id"}:
            continue
        if not isinstance(value, str):
            continue
        normalized_value = normalize_optional_text(value)
        if normalized_value:
            return normalized_value

    return None


async def retrieve_memory_context(context: AgentMemoryContext) -> str:
    """Fetch memory text for the current agent execution across enabled memory sources."""
    memory_parts: list[str] = []

    session_memory = await retrieve_session_memory_context(context)
    if session_memory:
        memory_parts.append(f"Recent Session Memory:\n{session_memory}")

    cognee_memory = await retrieve_cognee_memory_context(context)
    if cognee_memory:
        memory_parts.append(f"Relevant Cognee Memory:\n{cognee_memory}")

    if not memory_parts:
        return ""

    return truncate_text("\n\n".join(memory_parts), MAX_MEMORY_CONTEXT_LENGTH)


async def retrieve_cognee_memory_context(context: AgentMemoryContext) -> str:
    """Fetch dataset-backed Cognee search memory when enabled."""
    if not context.config.with_memory or context.scope is None:
        context.memory_query = ""
        return ""

    query_text = derive_query_text(
        context.config.memory_query_fixed,
        context.config.memory_query_from_method,
        context.method_params,
    )
    context.memory_query = query_text or ""
    if not query_text:
        logger.info("Skipping agent memory retrieval because no usable query could be derived.")
        return ""

    with new_span("cognee.agent_memory.retrieve") as span:
        span.set_attribute("cognee.agent_memory.query_length", len(query_text))
        try:
            from cognee.api.v1.search import SearchType, search

            results = await search(
                query_text=query_text,
                query_type=SearchType.GRAPH_SUMMARY_COMPLETION,
                user=context.scope.user,
                dataset_ids=[context.scope.dataset_id],
                system_prompt=context.config.memory_system_prompt,
                top_k=context.config.memory_top_k,
                only_context=context.config.memory_only_context,
            )
        except Exception as error:
            logger.warning(
                "Agent memory retrieval failed for %s: %s",
                context.origin_function,
                error,
                exc_info=False,
            )
            span.set_attribute("cognee.agent_memory.retrieval_failed", True)
            return ""

        memory_context = truncate_text(normalize_search_results(results), MAX_MEMORY_CONTEXT_LENGTH)
        span.set_attribute("cognee.agent_memory.context_length", len(memory_context))
        return memory_context


async def retrieve_session_memory_context(context: AgentMemoryContext) -> str:
    """Fetch recent trace feedback from the session-backed trace store when enabled."""
    if not context.config.with_session_memory or context.user is None:
        return ""

    from cognee.infrastructure.session.get_session_manager import get_session_manager

    session_manager = get_session_manager()
    try:
        feedback_values = await session_manager.get_agent_trace_feedback(
            user_id=str(context.user.id),
            session_id=context.config.session_id,
            last_n=context.config.session_memory_last_n,
        )
    except Exception as error:
        logger.warning(
            "Session agent memory retrieval failed for %s: %s",
            context.origin_function,
            error,
            exc_info=False,
        )
        return ""

    normalized_feedback = [
        normalized
        for value in feedback_values
        if (normalized := normalize_optional_text(value)) is not None
    ]
    if not normalized_feedback:
        return ""

    return "\n".join(normalized_feedback)


async def persist_trace(context: AgentMemoryContext) -> None:
    """Persist one agent trace step into session-backed storage."""
    if not context.config.save_session_traces or context.user is None:
        return

    from cognee.infrastructure.session.get_session_manager import get_session_manager

    session_manager = get_session_manager()
    user_id = str(context.user.id)
    try:
        await session_manager.add_agent_trace_step(
            user_id=user_id,
            session_id=context.config.session_id,
            origin_function=context.origin_function,
            status=context.status,
            generate_feedback_with_llm=context.config.session_trace_summary,
            memory_query=context.memory_query,
            memory_context=context.memory_context,
            method_params=context.method_params,
            method_return_value=sanitize_value(context.method_return_value),
            error_message=truncate_text(context.error_message, MAX_SERIALIZED_VALUE_LENGTH),
        )
    except Exception as error:
        logger.warning(
            "Agent trace persistence failed for %s: %s",
            context.origin_function,
            error,
            exc_info=False,
        )
        return

    if context.config.persist_session_trace_after is None:
        return

    resolved_session_id = context.config.session_id or session_manager.default_session_id

    try:
        trace_count = await session_manager.get_agent_trace_count(
            user_id=user_id,
            session_id=resolved_session_id,
        )
        if trace_count == 0 or trace_count % context.config.persist_session_trace_after != 0:
            return

        from cognee.memify_pipelines.persist_agent_trace_feedbacks_in_knowledge_graph import (
            persist_agent_trace_feedbacks_in_knowledge_graph_pipeline,
        )

        persist_kwargs = {
            "user": context.user,
            "session_ids": [resolved_session_id],
            "dataset": context.config.dataset_name or "main_dataset",
            "raw_trace_content": context.config.persist_session_trace_raw_content,
            "last_n_steps": context.config.persist_session_trace_after,
            "run_in_background": False,
        }
        if context.config.persist_session_trace_node_set_name is not None:
            persist_kwargs["node_set_name"] = context.config.persist_session_trace_node_set_name

        await persist_agent_trace_feedbacks_in_knowledge_graph_pipeline(
            **persist_kwargs,
        )
    except Exception as error:
        logger.warning(
            "Agent trace memify persistence failed for %s: %s",
            context.origin_function,
            error,
            exc_info=False,
        )


def normalize_search_results(results: Any) -> str:
    """Flatten heterogeneous search outputs into a single text blob."""
    if results is None:
        return ""
    if isinstance(results, str):
        return results
    if isinstance(results, list):
        normalized_items = [normalize_search_results(item) for item in results]
        return "\n".join(item for item in normalized_items if item).strip()
    if isinstance(results, dict):
        if "search_result" in results:
            return normalize_search_results(results["search_result"])
        return json.dumps(sanitize_value(results), ensure_ascii=False)
    if hasattr(results, "search_result"):
        return normalize_search_results(results.search_result)
    return str(results)