from __future__ import annotations import asyncio import importlib import os import tempfile import unittest from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, patch from uuid import uuid4 _SKILL = """\ --- description: Review code changes for correctness and regressions. allowed-tools: memory_search --- # Instructions Read the diff, identify concrete bugs, and cite file paths. """ def _run(coro): return asyncio.run(coro) def _make_skill_dir(slug: str = "code-review") -> Path: root = Path(tempfile.mkdtemp(dir=Path.cwd())) skill_dir = root / slug skill_dir.mkdir() (skill_dir / "SKILL.md").write_text(_SKILL) return root class TestSkillContract(unittest.TestCase): def setUp(self) -> None: # `set_database_global_context_variables` is a no-op when backend access # control is off; with default Kuzu+LanceDB it auto-enables and then tries # to look up the SimpleNamespace user in the relational DB. Disable it for # the duration of each test in this class. patcher = patch( "cognee.context_global_variables.backend_access_control_enabled", return_value=False, ) patcher.start() self.addCleanup(patcher.stop) def test_add_skills_persists_one_dataset_scoped_skill(self): from cognee.modules.engine.models import NodeSet from cognee.modules.tools.ingest_skills import add_skills root = _make_skill_dir() dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) try: with patch( "cognee.modules.tools.ingest_skills.add_data_points", new_callable=AsyncMock, ) as mock_add: skills = _run(add_skills(root, user=user, dataset=dataset)) assert len(skills) == 1 skill = skills[0] assert skill.name == "code-review" assert skill.dataset_scope == [str(dataset.id)] # source_file/source_dir are native normalized paths (see # skill_parser._normalize_path), so use os.sep rather than a # hard-coded "/" — on Windows the separator is "\". assert skill.source_file.endswith(os.path.join("code-review", "SKILL.md")) assert skill.source_dir.endswith("code-review") assert isinstance(skill.belongs_to_set[0], NodeSet) assert skill.belongs_to_set[0].name == "skills" assert skill.content_hash assert "Review code changes" in skill.search_text assert skill.skill_text == skill.search_text assert mock_add.await_count == 1 assert mock_add.await_args.kwargs["ctx"].dataset is dataset finally: import shutil shutil.rmtree(root) def test_same_skill_slug_in_two_datasets_gets_distinct_ids(self): from cognee.modules.tools.ingest_skills import add_skills root = _make_skill_dir() dataset_a = SimpleNamespace(id=uuid4(), name="a") dataset_b = SimpleNamespace(id=uuid4(), name="b") try: with patch( "cognee.modules.tools.ingest_skills.add_data_points", new_callable=AsyncMock, ): skill_a = _run(add_skills(root, dataset=dataset_a))[0] skill_b = _run(add_skills(root, dataset=dataset_b))[0] assert skill_a.name == skill_b.name assert skill_a.id != skill_b.id assert skill_a.dataset_scope == [str(dataset_a.id)] assert skill_b.dataset_scope == [str(dataset_b.id)] finally: import shutil shutil.rmtree(root) def test_agentic_completion_uses_agentic_retriever_only_when_explicit(self): from cognee.modules.retrieval.agentic_retriever import AgenticRetriever from cognee.modules.retrieval.graph_completion_cot_retriever import ( GraphCompletionCotRetriever, ) from cognee.modules.search.exceptions import UnsupportedSearchTypeError from cognee.modules.search.methods.get_search_type_retriever_instance import ( get_search_type_retriever_instance, ) from cognee.modules.search.types import SearchType user = SimpleNamespace(id=uuid4()) dataset = SimpleNamespace(id=uuid4(), name="project") retriever = _run( get_search_type_retriever_instance( SearchType.AGENTIC_COMPLETION, "review this", user=user, dataset=dataset, retriever_specific_config={"skills": ["code-review"]}, ) ) assert isinstance(retriever, AgenticRetriever) cot = _run( get_search_type_retriever_instance( SearchType.GRAPH_COMPLETION_COT, "review this", retriever_specific_config={}, ) ) assert isinstance(cot, GraphCompletionCotRetriever) with self.assertRaises(UnsupportedSearchTypeError): _run( get_search_type_retriever_instance( SearchType.GRAPH_COMPLETION, "review this", retriever_specific_config={"skills": ["code-review"]}, ) ) def test_agentic_retriever_skips_memory_retrieval_when_graph_has_no_edges(self): from cognee.modules.retrieval.agentic_retriever import AgenticRetriever from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4()) retriever = AgenticRetriever(skills=["code-review"], user=user, dataset=dataset) async def run_retrieval(): with ( patch.object(AgenticRetriever, "_graph_has_edges", new_callable=AsyncMock) as edges, patch( "cognee.modules.retrieval.agentic_retriever.resolve_skills", new_callable=AsyncMock, return_value=[], ), patch( "cognee.modules.retrieval.agentic_retriever.list_tools_for_dataset", new_callable=AsyncMock, return_value=[], ), patch.object( GraphCompletionRetriever, "get_retrieved_objects", new_callable=AsyncMock, ) as parent_retrieval, ): edges.return_value = False result = await retriever.get_retrieved_objects(query="review this") return result, parent_retrieval result, parent_retrieval = _run(run_retrieval()) assert result["triplets"] == [] parent_retrieval.assert_not_awaited() def test_load_skill_loads_only_active_context_skills(self): from cognee.modules.engine.models import Skill from cognee.modules.tools.builtin.load_skill import handler from cognee.modules.tools.context import active_skills_var, opened_skills_var from cognee.modules.tools.errors import ToolInvocationError skill = Skill(name="code-review", description="Review code.", procedure="step 1") async def run_handler(): active_token = active_skills_var.set({skill.name: skill}) opened: set[str] = set() opened_token = opened_skills_var.set(opened) try: body = await handler({"name": "code-review"}) with self.assertRaises(ToolInvocationError): await handler({"name": "other"}) finally: active_skills_var.reset(active_token) opened_skills_var.reset(opened_token) return body, opened body, opened = _run(run_handler()) assert "step 1" in body assert opened == {"code-review"} def test_agentic_skill_run_write_is_dataset_scoped(self): from cognee.modules.engine.models import Skill from cognee.modules.engine.models.SkillRun import UNSCORED_SKILL_RUN_SCORE from cognee.modules.retrieval.agentic_retriever import AgenticRetriever dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) skill = Skill( name="code-review", description="Review code.", skill_text="code-review\n\nReview code.", search_text="code-review\n\nReview code.", dataset_scope=[str(dataset.id)], ) add_data_points_module = importlib.import_module("cognee.tasks.storage.add_data_points") with patch.object( add_data_points_module, "add_data_points", new_callable=AsyncMock ) as mock_add: _run( AgenticRetriever._record_skill_runs( [skill], "review this", "done", user=user, dataset=dataset, session_id="session-1", ) ) run = mock_add.await_args.args[0][0] assert run.selected_skill_id == str(skill.id) assert run.selected_skill_name == "code-review" assert run.candidate_skills[0].skill_name == "code-review" assert run.candidate_skills[0].skill_description == "Review code." assert run.candidate_skills[0].skill_text == skill.skill_text assert run.candidate_skills[0].metadata["index_fields"] == ["skill_description"] assert run.dataset_scope == [str(dataset.id)] assert run.session_id == "session-1" assert run.success_score == UNSCORED_SKILL_RUN_SCORE assert mock_add.await_args.kwargs["ctx"].dataset is dataset def test_remember_skill_run_entry_persists_with_dataset_context(self): from cognee.memory import SkillRunEntry from cognee.modules.engine.models import Skill from cognee.modules.tools.skill_runs import remember_skill_run_entry dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) skill = Skill( name="code-review", description="Review code.", procedure="step 1", skill_text="code-review\n\nReview code.\n\nstep 1", search_text="code-review\n\nReview code.\n\nstep 1", dataset_scope=[str(dataset.id)], ) entry = SkillRunEntry( run_id="run-1", selected_skill_id="code-review", task_text="Review the diff", success_score=0.2, feedback=-0.5, ) async def run_entry(): with ( patch("cognee.modules.tools.skill_runs.setup", new_callable=AsyncMock), patch( "cognee.modules.tools.skill_runs.resolve_authorized_user_datasets", new_callable=AsyncMock, return_value=(user, [dataset]), ), patch( "cognee.modules.tools.skill_runs.resolve_skills", new_callable=AsyncMock, return_value=[skill], ), patch( "cognee.modules.tools.skill_runs.add_data_points", new_callable=AsyncMock, ) as mock_add, ): run, resolved_dataset = await remember_skill_run_entry( entry, dataset_name="project", session_id="session-1", user=user, ) return run, resolved_dataset, mock_add run, resolved_dataset, mock_add = _run(run_entry()) assert resolved_dataset is dataset assert run.selected_skill_id == str(skill.id) assert run.selected_skill_name == "code-review" assert run.selected_skill is skill assert run.candidate_skills[0].skill_id == str(skill.id) assert run.candidate_skills[0].skill_name == "code-review" assert run.candidate_skills[0].skill_description == "Review code." assert "step 1" in run.candidate_skills[0].skill_text assert run.candidate_skills[0].metadata["index_fields"] == ["skill_description"] assert run.dataset_scope == [str(dataset.id)] assert run.belongs_to_set[0].name == "skills" assert mock_add.await_args.kwargs["ctx"].dataset is dataset def test_skill_improvement_proposal_does_not_mutate_skill(self): from cognee.modules.engine.models import Skill, SkillRun from cognee.modules.memify.skill_improvement import ( SkillImprovementDraft, improve_skill, ) dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) skill = Skill(name="code-review", description="Review code.", procedure="old") run = SkillRun( run_id="run-1", session_id="session", task_text="review", selected_skill_id=str(skill.id), selected_skill_name=skill.name, dataset_scope=[str(dataset.id)], success_score=0.1, ) async def create_proposal(): with ( patch( "cognee.modules.memify.skill_improvement.find_skill_by_name", new_callable=AsyncMock, return_value=skill, ), patch( "cognee.modules.memify.skill_improvement._find_recent_failure_runs", new_callable=AsyncMock, return_value=[run], ), patch( "cognee.modules.memify.skill_improvement._generate_proposal", new_callable=AsyncMock, return_value=SkillImprovementDraft( proposed_procedure="new", rationale="The old procedure missed regression checks.", confidence=0.8, ), ), patch( "cognee.modules.memify.skill_improvement.add_data_points", new_callable=AsyncMock, ) as mock_add, ): proposal = await improve_skill("code-review", dataset=dataset, user=user) return proposal, mock_add proposal, mock_add = _run(create_proposal()) assert skill.procedure == "old" assert proposal.old_procedure == "old" assert proposal.proposed_procedure == "# code-review\n\nnew" assert proposal.runs_used == ["run-1"] assert proposal.skill is skill assert proposal.runs == [run] assert proposal.dataset_scope == [str(dataset.id)] assert proposal.belongs_to_set[0].name == "skills" assert mock_add.await_args.args[0][0] is proposal def test_skill_improvement_apply_requires_existing_proposal_id(self): from cognee.modules.memify.skill_improvement import improve_skill dataset = SimpleNamespace(id=uuid4(), name="project") with self.assertRaisesRegex(ValueError, "proposal_id"): _run(improve_skill("code-review", dataset=dataset, apply=True)) def test_skill_improvement_apply_updates_only_target_skill(self): from cognee.modules.engine.models import Skill, SkillImprovementProposal from cognee.modules.memify.skill_improvement import improve_skill dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) skill = Skill(name="code-review", description="Review code.", procedure="old") proposal = SkillImprovementProposal( proposal_id="proposal-1", skill_id=str(skill.id), skill_name=skill.name, dataset_scope=[str(dataset.id)], old_procedure="old", proposed_procedure="new", ) async def apply_proposal(): with ( patch( "cognee.modules.memify.skill_improvement._find_proposal", new_callable=AsyncMock, return_value=proposal, ), patch( "cognee.modules.memify.skill_improvement.find_skill_by_id", new_callable=AsyncMock, return_value=skill, ), patch( "cognee.modules.memify.skill_improvement.add_data_points", new_callable=AsyncMock, ) as mock_add, ): applied = await improve_skill( "code-review", dataset=dataset, user=user, proposal_id="proposal-1", apply=True, ) return applied, mock_add applied, mock_add = _run(apply_proposal()) assert skill.procedure == "# code-review\n\nnew" assert skill.skill_text == "code-review\n\nReview code.\n\n# code-review\n\nnew" assert skill.search_text == skill.skill_text assert applied.status == "applied" assert mock_add.await_args.args[0] == [skill, proposal] def test_skill_improvement_apply_preserves_existing_skill_body_heading(self): from cognee.modules.engine.models import Skill, SkillImprovementProposal from cognee.modules.memify.skill_improvement import improve_skill dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) skill = Skill(name="code-review", description="Review code.", procedure="old") proposal = SkillImprovementProposal( proposal_id="proposal-1", skill_id=str(skill.id), skill_name=skill.name, dataset_scope=[str(dataset.id)], old_procedure="old", proposed_procedure="# code-review\n\n- Read the diff.", ) async def apply_proposal(): with ( patch( "cognee.modules.memify.skill_improvement._find_proposal", new_callable=AsyncMock, return_value=proposal, ), patch( "cognee.modules.memify.skill_improvement.find_skill_by_id", new_callable=AsyncMock, return_value=skill, ), patch( "cognee.modules.memify.skill_improvement.add_data_points", new_callable=AsyncMock, ), ): return await improve_skill( "code-review", dataset=dataset, user=user, proposal_id="proposal-1", apply=True, ) _run(apply_proposal()) assert skill.procedure == "# code-review\n\n- Read the diff." def test_get_proposal_returns_proposal_preview(self): from cognee.modules.engine.models import SkillImprovementProposal from cognee.modules.memify.skill_improvement import get_proposal dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) proposal = SkillImprovementProposal( proposal_id="proposal-1", skill_id=str(uuid4()), skill_name="code-review", dataset_scope=[str(dataset.id)], old_procedure="old", proposed_procedure="# code-review\n\nnew", rationale="Missed regression checks.", confidence=0.8, status="proposed", ) async def fetch(): with patch( "cognee.modules.memify.skill_improvement._find_proposal", new_callable=AsyncMock, return_value=proposal, ) as mock_find: result = await get_proposal("proposal-1", dataset=dataset, user=user) return result, mock_find result, mock_find = _run(fetch()) assert result is proposal assert result.old_procedure == "old" assert result.proposed_procedure == "# code-review\n\nnew" assert result.rationale == "Missed regression checks." assert result.confidence == 0.8 assert mock_find.await_args.kwargs["proposal_id"] == "proposal-1" assert mock_find.await_args.kwargs["dataset_id"] == dataset.id def test_get_proposal_returns_none_for_unknown_id(self): from cognee.modules.memify.skill_improvement import get_proposal dataset = SimpleNamespace(id=uuid4(), name="project") user = SimpleNamespace(id=uuid4(), tenant_id=uuid4()) async def fetch(): with patch( "cognee.modules.memify.skill_improvement._find_proposal", new_callable=AsyncMock, return_value=None, ): return await get_proposal("missing", dataset=dataset, user=user) assert _run(fetch()) is None def test_get_proposal_requires_dataset_with_id(self): from cognee.modules.memify.skill_improvement import get_proposal with self.assertRaisesRegex(ValueError, "explicit dataset"): _run(get_proposal("proposal-1", dataset=SimpleNamespace(id=None))) def test_materialize_inline_skill_writes_slug_dir(self): from cognee.api.v1.remember.remember import _materialize_inline_skill tmp_dir, source = _materialize_inline_skill(_SKILL, "code-review") try: md = source / "code-review" / "SKILL.md" assert md.is_file() assert "Review code changes" in md.read_text() finally: tmp_dir.cleanup() def test_materialize_inline_skill_defaults_and_sanitizes_slug(self): from cognee.api.v1.remember.remember import _materialize_inline_skill tmp_dir, source = _materialize_inline_skill("# x", None) try: assert (source / "skill" / "SKILL.md").is_file() finally: tmp_dir.cleanup() # A traversal-y name collapses to a single safe path segment. tmp_dir2, source2 = _materialize_inline_skill("# x", "../../etc/evil") try: assert (source2 / "evil" / "SKILL.md").is_file() finally: tmp_dir2.cleanup()