"""Dataset-scoped Skill lookup helpers for the agentic retriever.""" from __future__ import annotations from typing import List, Optional, Sequence, Union from uuid import UUID from cognee.modules.engine.models import Skill from cognee.shared.logging_utils import get_logger logger = get_logger("cognee.tools.resolve_skills") def _skill_in_dataset_scope(skill: Skill, dataset_id: UUID) -> bool: return skill.is_active and str(dataset_id) in (skill.dataset_scope or []) async def resolve_skills( skills: Optional[Sequence[Union[str, Skill]]] = None, *, dataset_id: Optional[UUID] = None, ) -> List[Skill]: """Resolve explicit skills inside one dataset. v1 intentionally rejects unscoped skill lookup and multi-dataset lookup. """ if dataset_id is None: raise ValueError("Skill lookup requires one explicit dataset.") resolved: List[Skill] = [] seen_ids = set() for item in skills or []: skill = item if isinstance(item, Skill) else None if isinstance(item, str): skill = await find_skill_by_name(item, dataset_id=dataset_id) elif skill is None: logger.warning("Skill entries must be Skill or str; got %s", type(item).__name__) if skill is None: logger.warning("Skill %r not found in dataset %s; skipping", item, dataset_id) continue if not _skill_in_dataset_scope(skill, dataset_id): logger.warning("Skill %r is outside dataset scope; skipping", skill.name) continue if skill.id not in seen_ids: resolved.append(skill) seen_ids.add(skill.id) return resolved async def find_skill_by_id(skill_id: str, *, dataset_id: UUID) -> Optional[Skill]: raw_nodes = await _load_skill_nodes() for raw in raw_nodes: skill = _coerce_skill(raw) if skill is None: continue if str(skill.id) == str(skill_id) and _skill_in_dataset_scope(skill, dataset_id): return skill return None async def find_skill_by_name(name: str, *, dataset_id: UUID) -> Optional[Skill]: raw_nodes = await _load_skill_nodes(name=name) for raw in raw_nodes: skill = _coerce_skill(raw) if skill is None or skill.name != name: continue if _skill_in_dataset_scope(skill, dataset_id): return skill return None async def _load_skill_nodes(name: Optional[str] = None): try: from cognee.infrastructure.databases.graph import get_graph_engine except Exception: return [] try: graph_engine = await get_graph_engine() except Exception: return [] get_by_type = getattr(graph_engine, "get_nodes_by_type", None) if get_by_type is not None: try: return await get_by_type(node_type=Skill) except Exception as exc: logger.warning("Skill lookup by type failed: %s", exc) return [] get_nodeset = getattr(graph_engine, "get_nodeset_subgraph", None) if get_nodeset is not None and name is not None: try: nodes, _ = await get_nodeset(node_type=Skill, node_name=[name]) if nodes: return nodes except Exception as exc: logger.warning("Skill lookup by nodeset failed: %s", exc) get_graph_data = getattr(graph_engine, "get_graph_data", None) if get_graph_data is None: return [] try: nodes, _ = await get_graph_data() return nodes except Exception as exc: logger.warning("Skill lookup by full graph scan failed: %s", exc) return [] def _coerce_skill(raw) -> Optional[Skill]: if isinstance(raw, Skill): return raw node_id = None if isinstance(raw, (list, tuple)) and len(raw) > 1: node_id = raw[0] raw = raw[1] data = raw.model_dump() if hasattr(raw, "model_dump") else raw if not isinstance(data, dict): return None data = {k: v for k, v in data.items() if k != "metadata"} if node_id is not None and "id" not in data: data["id"] = node_id try: return Skill.model_validate(data) except Exception: return None