topoteretes--cognee/cognee/modules/tools/skill_runs.py

"""Record SkillRun feedback through the remember() typed-entry path."""

from __future__ import annotations

from types import SimpleNamespace
from typing import Optional
from uuid import NAMESPACE_URL, UUID, uuid5

from cognee.memory.entries import SkillRunEntry
from cognee.context_global_variables import set_database_global_context_variables
from cognee.modules.engine.models import (
    CandidateSkill,
    NodeSet,
    SkillRun,
    ToolCall,
    UNSCORED_SKILL_RUN_SCORE,
)
from cognee.modules.engine.operations.setup import setup
from cognee.modules.engine.utils.generate_node_id import generate_node_id
from cognee.modules.pipelines.layers.resolve_authorized_user_datasets import (
    resolve_authorized_user_datasets,
)
from cognee.modules.pipelines.models import PipelineContext
from cognee.modules.tools.resolve_skills import resolve_skills
from cognee.shared.logging_utils import get_logger
from cognee.tasks.storage.add_data_points import add_data_points


logger = get_logger("cognee.tools.skill_runs")


def _skill_run_data_id(dataset_id: UUID, run_id: str) -> UUID:
    """Stable pseudo data id used to attach SkillRun writes to dataset ACL tables."""
    return uuid5(NAMESPACE_URL, f"cognee:skill-runs:{dataset_id}:{run_id}")


def _make_storage_context(user, dataset, run_id: str) -> PipelineContext:
    return PipelineContext(
        user=user,
        dataset=dataset,
        data_item=SimpleNamespace(id=_skill_run_data_id(dataset.id, run_id)),
        pipeline_name="skill_runs_pipeline",
    )


def _coerce_tool_trace(raw_trace: list[dict]) -> list[ToolCall]:
    calls: list[ToolCall] = []
    for index, item in enumerate(raw_trace):
        if not isinstance(item, dict):
            logger.warning("Skipping non-dict tool_trace item at index %d", index)
            continue
        calls.append(ToolCall(**item))
    return calls


def _candidate_skill_for_id(skill_id: str, selected_skill) -> CandidateSkill:
    is_selected = str(skill_id) == str(selected_skill.id)
    if not is_selected:
        return CandidateSkill(skill_id=str(skill_id))

    return CandidateSkill(
        skill_id=str(selected_skill.id),
        skill_name=selected_skill.name,
        skill_description=selected_skill.description,
        skill_text=selected_skill.skill_text or selected_skill.search_text,
    )


async def remember_skill_run_entry(
    entry: SkillRunEntry,
    *,
    dataset_name: str,
    session_id: Optional[str],
    user=None,
) -> tuple[SkillRun, object]:
    """Persist a dataset-scoped SkillRun from the remember() typed-entry path."""
    await setup()

    if user is None:
        from cognee.modules.users.methods import get_default_user

        user = await get_default_user()

    user, authorized_datasets = await resolve_authorized_user_datasets(dataset_name, user)
    dataset = authorized_datasets[0]

    owner_id = getattr(dataset, "owner_id", None) or getattr(user, "id", None)
    if owner_id is None:
        raise ValueError("SkillRun persistence requires a dataset owner or user.")

    async with set_database_global_context_variables(dataset.id, owner_id):
        resolved_skills = await resolve_skills([entry.selected_skill_id], dataset_id=dataset.id)
        if not resolved_skills:
            raise ValueError(
                f"Skill '{entry.selected_skill_id}' was not found or is not visible "
                f"in dataset '{dataset.name}'"
            )

        selected_skill = resolved_skills[0]
        candidate_ids = entry.candidate_skill_ids or [str(selected_skill.id)]
        success_score = (
            UNSCORED_SKILL_RUN_SCORE if entry.success_score is None else entry.success_score
        )

        run = SkillRun(
            run_id=entry.run_id,
            selected_skill_id=str(selected_skill.id),
            selected_skill_name=selected_skill.name,
            selected_skill=selected_skill,
            dataset_scope=[str(dataset.id)],
            task_text=entry.task_text,
            result_summary=entry.result_summary,
            success_score=success_score,
            session_id=session_id or "agentic",
            candidate_skills=[
                _candidate_skill_for_id(skill_id, selected_skill) for skill_id in candidate_ids
            ],
            task_pattern_id=entry.task_pattern_id,
            router_version=entry.router_version,
            tool_trace=_coerce_tool_trace(entry.tool_trace),
            error_type=entry.error_type,
            error_message=entry.error_message,
            started_at_ms=entry.started_at_ms,
            latency_ms=entry.latency_ms,
            feedback=entry.feedback,
        )
        run.belongs_to_set = [
            NodeSet(id=generate_node_id(f"NodeSet:{entry.node_set}"), name=entry.node_set)
        ]

        await add_data_points([run], ctx=_make_storage_context(user, dataset, entry.run_id))

    return run, dataset