topoteretes--cognee/cognee/api/v1/session/session.py

from typing import List, Optional

from cognee.context_global_variables import session_user
from cognee.exceptions import CogneeValidationError
from cognee.infrastructure.databases.cache.models import SessionQAEntry
from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
from cognee.infrastructure.session.get_session_manager import get_session_manager
from cognee.modules.users.exceptions.exceptions import UserNotFoundError
from cognee.modules.users.methods import get_default_user
from cognee.modules.users.models import User
from cognee.shared.logging_utils import get_logger

logger = get_logger("session_api_sdk")


async def _resolve_user(user: Optional[User]) -> User:
    if user is not None:
        if getattr(user, "id", None) is None:
            raise CogneeValidationError(
                message="Session user must have an id.",
                name="SessionPreconditionError",
            )
        return user
    ctx_user = session_user.get()
    if ctx_user is not None and getattr(ctx_user, "id", None) is not None:
        return ctx_user
    try:
        return await get_default_user()
    except (DatabaseNotCreatedError, UserNotFoundError) as error:
        raise CogneeValidationError(
            message=(
                "Session prerequisites not met: no default user found. "
                "Initialize Cognee before using session APIs by running "
                "`await cognee.add(...)` followed by `await cognee.cognify()`."
            ),
            name="SessionPreconditionError",
        ) from error


async def get_session(
    session_id: str = "default_session",
    last_n: Optional[int] = None,
    user: Optional[User] = None,
) -> List[SessionQAEntry]:
    resolved_user = await _resolve_user(user)
    user_id = str(resolved_user.id)

    try:
        sm = get_session_manager()
        raw = await sm.get_session(
            user_id=user_id,
            session_id=session_id,
            last_n=last_n,
            formatted=False,
        )
    except Exception as e:
        logger.warning("get_session: error from SessionManager: %s", e)
        return []

    if not raw:
        return []

    result: List[SessionQAEntry] = []
    for entry in raw:
        if isinstance(entry, dict):
            try:
                result.append(SessionQAEntry.model_validate(entry))
            except Exception as e:
                logger.warning("get_session: skip invalid entry: %s", e)
        elif isinstance(entry, SessionQAEntry):
            result.append(entry)
        else:
            logger.warning("get_session: skip non-dict non-SessionQAEntry entry: %s", type(entry))
    return result


async def add_feedback(
    session_id: str,
    qa_id: str,
    feedback_text: Optional[str] = None,
    feedback_score: Optional[int] = None,
    user: Optional[User] = None,
) -> bool:
    resolved_user = await _resolve_user(user)
    user_id = str(resolved_user.id)

    try:
        sm = get_session_manager()
        return await sm.add_feedback(
            user_id=user_id,
            session_id=session_id,
            qa_id=qa_id,
            feedback_text=feedback_text,
            feedback_score=feedback_score,
        )
    except Exception as e:
        logger.warning("add_feedback: error from SessionManager: %s", e)
        return False


async def add_frequency_weights(
    session_id: str,
    qa_id: str,
    node_ids: Optional[list[str]] = None,
    edge_ids: Optional[list[str]] = None,
    user: Optional[User] = None,
) -> bool:
    """Add or update frequency weight data for a QA entry.

    This function stores the graph elements (node_ids and edge_ids) that were used
    in generating the answer for a QA entry. This data is later processed by
    apply_frequency_weights to increment the frequency weights of those elements.

    The frequency_weights_applied flag is reset to False so the entry will be
    reprocessed by the apply_frequency_weights pipeline.

    Args:
        session_id: Session identifier.
        qa_id: QA entry identifier.
        node_ids: List of node IDs used in generating the answer.
        edge_ids: List of edge IDs used in generating the answer.
        user: User that owns the session. If None, uses session/context user or default user.

    Returns:
        True if updated, False if QA not found or cache unavailable.
    """
    from cognee.tasks.memify.frequency_weights_constants import (
        MEMIFY_METADATA_FREQUENCY_WEIGHTS_APPLIED_KEY,
    )

    resolved_user = await _resolve_user(user)
    user_id = str(resolved_user.id)

    used_graph_element_ids: dict[str, list[str]] = {}
    if node_ids:
        used_graph_element_ids["node_ids"] = node_ids
    if edge_ids:
        used_graph_element_ids["edge_ids"] = edge_ids

    try:
        sm = get_session_manager()
        return await sm.update_qa(
            user_id=user_id,
            session_id=session_id,
            qa_id=qa_id,
            used_graph_element_ids=used_graph_element_ids if used_graph_element_ids else None,
            memify_metadata={MEMIFY_METADATA_FREQUENCY_WEIGHTS_APPLIED_KEY: False},
        )
    except Exception as e:
        logger.warning("add_frequency_weights: error from SessionManager: %s", e)
        return False


async def delete_feedback(
    session_id: str,
    qa_id: str,
    user: Optional[User] = None,
) -> bool:
    """
    Clear feedback for a QA entry (sets feedback_text and feedback_score to None).

    When user is None, uses session context or default user.

    Args:
        session_id: Session identifier.
        qa_id: QA entry identifier to clear feedback for.
        user: User that owns the session. If None, uses session/context user or default user.

    Returns:
        True if feedback was cleared, False if QA not found or cache unavailable.
    """
    resolved_user = await _resolve_user(user)
    user_id = str(resolved_user.id)

    try:
        sm = get_session_manager()
        return await sm.delete_feedback(
            user_id=user_id,
            session_id=session_id,
            qa_id=qa_id,
        )
    except Exception as e:
        logger.warning("delete_feedback: error from SessionManager: %s", e)
        return False