diff --git a/CLAUDE.md b/CLAUDE.md index 6ced8cc51c..d446412cfe 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -255,6 +255,7 @@ Available search types (from `cognee/modules/search/types/SearchType.py`), passe - **TEMPORAL** - Time-aware graph search - **FEELING_LUCKY** - Automatic search type selection - **CODING_RULES** - Code-specific search rules +- **SKILLS** - Semantic discovery of skill playbooks (metadata-only, no LLM; requires exactly one dataset) `recall()` picks one of these automatically when `query_type` is omitted. The CLI is narrower: `cognee-cli recall --query-type` accepts only the choices in `cognee/cli/config.py:SEARCH_TYPE_CHOICES` and defaults to `HYBRID_COMPLETION`; the rest are SDK-only. @@ -704,6 +705,15 @@ Opt-in LLM check that runs as the last `cognify()` task (default **off**). After - **Applies to `remember()` too** — and to session memory bridged back by `improve()` — since those build their graphs through `cognify()`. The exception is `remember(content_type="code")`, which runs the separate code-graph pipeline. - **Scope / limitations**: only the 1-hop neighbourhood of the touched entities is compared; structural edges (`contains`, `is_part_of`, `made_from`, `exists_in`, `contradicts`) and edges with an unnamed endpoint are skipped; the temporal cognify path is not covered. +### Skills (Procedural Memory) +Dataset-scoped `SKILL.md` playbooks agents can discover, load on demand, execute, and improve from run history. + +- **Ingest**: `remember(content_type="skills", dataset_name=...)` (folder, file, or inline via `skills_text`/`skill_name`) — requires an explicit dataset; re-ingest upserts (deterministic ids). HTTP: `POST /skills`. +- **Discover**: `SearchType.SKILLS` — one vector search over the `Skill_search_text` collection, no LLM, **metadata-only results** (never the procedure body; progressive disclosure). Requires exactly one dataset; skills outside that dataset's scope, inactive skills, and empty-scope legacy skills are filtered out. Missing collection returns `[]`, not an error. +- **Skill gate**: `recall()` runs a deterministic regex gate (`cognee/api/v1/recall/skill_gate.py`); procedural-sounding queries trigger a concurrent SKILLS lookup whose hits are appended tagged `source="skills"`. Additive and fail-safe; only fires when exactly one dataset is targeted. Disable with `SKILL_GATE_ENABLED=false`. +- **Execute**: `SearchType.AGENTIC_COMPLETION` with `skills=[...]` — LLM sees name+description, loads bodies via the `load_skill` tool (12k char cap). +- **Improve**: `SkillRun` records (via `remember()` skill-run entries) feed LLM-drafted `SkillImprovementProposal`s; preview then apply by proposal id (`/proposals` router). + ### Code Files (cognify CODE route) Supported code files (`.py`, `.go`, `.ts`, `.java`, `.rs`, … — the extension list lives on `code_loader`) are recognized at add time through the loader system: the code loader claims the file, stores it under its real extension, and `ingest_data` tags the record with `system_metadata = {"source": "code"}`. Cognify then routes such items down the CODE route, which runs the deterministic enola code graph pipeline per file — typed `CodeSymbol`/`CodeModule`/… nodes with `calls`/`imports`/`has_method` edges, **no LLM calls**. diff --git a/README.md b/README.md index 1ba79f2e8a..0baeba3d2d 100644 --- a/README.md +++ b/README.md @@ -198,6 +198,31 @@ A third flag, `DATASET_QUEUE_ENABLED=false`, removes the per-process concurrency on datasets; it saves a little latency but risks file-lock leaks and resource exhaustion when multiple datasets run in parallel — leave it on for servers. +### Procedural memory (skills) + +Beyond facts, cognee stores **skills** — dataset-scoped `SKILL.md` playbooks that +agents can discover semantically, load on demand, execute, and improve from run +history: + +```python +# Ingest a folder of SKILL.md files into a dataset +await cognee.remember("./skills", content_type="skills", dataset_name="ops") + +# Discover by meaning — returns name/description/metadata, never the full procedure +results = await cognee.search("how do I deploy to staging", + query_type=SearchType.SKILLS, datasets=["ops"]) + +# recall() also runs a deterministic (no-LLM) skill gate: procedural questions +# automatically get matching skills appended, tagged source="skills" +results = await cognee.recall("how do I deploy to staging", datasets=["ops"]) +``` + +Skills are authored as `SKILL.md` files (one directory per skill, frontmatter for +description/tools/maintainer), scoped to exactly one dataset at ingest time, and +executed through `SearchType.AGENTIC_COMPLETION`, where the agent loads full +procedure bodies on demand via the `load_skill` tool. Disable the recall gate with +`SKILL_GATE_ENABLED=false`. + ## Run with Docker Prefer containers? Cognee publishes prebuilt images to Docker Hub on every push to `main`: diff --git a/cognee/api/v1/recall/recall.py b/cognee/api/v1/recall/recall.py index fb234f7129..2f84c1b8dc 100644 --- a/cognee/api/v1/recall/recall.py +++ b/cognee/api/v1/recall/recall.py @@ -1,3 +1,4 @@ +import asyncio import re from typing import Annotated, Literal from uuid import UUID @@ -34,6 +35,7 @@ ResponseMarkerEntry, ResponseQAEntry, ResponseSessionContextEntry, + ResponseSkillEntry, ResponseToolEntry, ) from cognee.modules.recall.types.SearchResultItem import SearchResultItem @@ -900,6 +902,60 @@ async def _run_code() -> list[RecallResponse]: ) return tagged + async def _run_skill_gate(gate_top_k: int) -> list[RecallResponse]: + """Metadata-only SKILLS lookup for the deterministic skill gate. + + Skipped silently unless exactly one dataset is targeted (skill + lookup is single-dataset by invariant). Any failure contributes + nothing instead of failing the recall. + """ + from cognee.modules.search.methods.search import authorized_search + + try: + gate_user = user + if gate_user is None: + gate_user = await get_default_user() + + if dataset_ids and len(dataset_ids) == 1: + gate_dataset_ids = list(dataset_ids) + elif datasets and len(datasets) == 1: + authorized = await get_authorized_existing_datasets( + datasets, "read", gate_user + ) + if len(authorized) != 1: + return [] + gate_dataset_ids = [dataset.id for dataset in authorized] + else: + return [] + + payloads = await authorized_search( + query_text=query_text, + query_type=SearchType.SKILLS, + user=gate_user, + dataset_ids=gate_dataset_ids, + top_k=gate_top_k, + ) + except Exception as error: + logger.warning("Skill gate lookup failed (non-fatal): %s", error) + return [] + + entries: list[RecallResponse] = [] + for payload in payloads or []: + for item in getattr(payload, "completion", None) or []: + if not isinstance(item, dict): + continue + name = item.get("name") or "" + description = item.get("description") or "" + entries.append( + ResponseSkillEntry( + source="skills", + text=f"{name}: {description}" if description else name, + skill={k: v for k, v in item.items() if k != "score"}, + score=item.get("score"), + ) + ) + return entries + runners = { "session": _run_session, "trace": _run_trace, @@ -909,22 +965,49 @@ async def _run_code() -> list[RecallResponse]: "code": _run_code, } + # Deterministic skill gate: a procedural-looking query triggers a + # concurrent metadata-only SKILLS lookup (one vector search, no + # LLM call). Additive only — the main lanes never wait on it, and + # explicit SKILLS / AGENTIC_COMPLETION calls bypass it. + skills_task = None + if ( + "graph" in sources + and not only_context + and query_type not in (SearchType.SKILLS, SearchType.AGENTIC_COMPLETION) + ): + from cognee.api.v1.recall.skill_gate import ( + DEFAULT_SKILL_GATE_TOP_K, + should_search_skills, + skill_gate_enabled, + ) + + if skill_gate_enabled() and should_search_skills(query_text).fired: + skills_task = asyncio.create_task(_run_skill_gate(DEFAULT_SKILL_GATE_TOP_K)) + session_result_count = 0 - for src in sources: - runner = runners.get(src) - if runner is None: - continue - # Auto mode special case: session hit short-circuits graph. - if auto_fallthrough and src == "graph" and merged: - break - # on_empty: the other sources gave cognee enough context — don't - # go back to the external database. - if src == "tools" and tools_trigger == "on_empty" and merged: - continue - part = await runner() - if src == "session": - session_result_count = len(part) - merged.extend(part) + try: + for src in sources: + runner = runners.get(src) + if runner is None: + continue + # Auto mode special case: session hit short-circuits graph. + if auto_fallthrough and src == "graph" and merged: + break + # on_empty: the other sources gave cognee enough context — don't + # go back to the external database. + if src == "tools" and tools_trigger == "on_empty" and merged: + continue + part = await runner() + if src == "session": + session_result_count = len(part) + merged.extend(part) + except BaseException: + if skills_task is not None: + skills_task.cancel() + raise + + if skills_task is not None: + merged.extend(await skills_task) if session_result_count: span.set_attribute(COGNEE_SESSION_ENTRY_COUNT, session_result_count) diff --git a/cognee/api/v1/recall/skill_gate.py b/cognee/api/v1/recall/skill_gate.py new file mode 100644 index 0000000000..11fecf3a54 --- /dev/null +++ b/cognee/api/v1/recall/skill_gate.py @@ -0,0 +1,86 @@ +""" +Deterministic skill gate for recall(). + +Decides — with regexes only, no LLM and no I/O — whether a query looks +procedural enough to warrant a skill lookup. When the gate fires (and exactly +one dataset is targeted), recall() runs a metadata-only SKILLS search +concurrently with the main search and appends the hits tagged +``source="skills"``. The gate is additive: the main answer is never replaced +or blocked by it. + +Disable with ``SKILL_GATE_ENABLED=false``. +""" + +import os +import re +from dataclasses import dataclass, field + +from cognee.api.v1.recall.query_router import _is_negated +from cognee.shared.logging_utils import get_logger + +logger = get_logger("skill_gate") + +# How many skills the gate lane asks for. Deliberately small: gate hits are a +# side-channel next to the main answer, not the answer itself. +DEFAULT_SKILL_GATE_TOP_K = 3 + +# Each rule: (pattern, weight). Weights accumulate; the gate fires at the +# threshold. Weak signals (bare ops verbs) score 2.0 so one alone never fires; +# procedural phrasings score 3.0+ and fire on their own. +_GATE_RULES: list[tuple[re.Pattern, float]] = [ + (re.compile(r"\b(how (do|can|should|would) (i|we|you)|how to)\b", re.IGNORECASE), 3.0), + (re.compile(r"\b(steps? (to|for)|step.by.step)\b", re.IGNORECASE), 3.0), + (re.compile(r"\b(procedure|playbook|runbook|checklist|workflow)\b", re.IGNORECASE), 4.0), + (re.compile(r"\b(walk (me|us) through|guide (to|for|on))\b", re.IGNORECASE), 3.0), + (re.compile(r"\bwhat('?s| is) the (process|procedure)\b", re.IGNORECASE), 4.0), + (re.compile(r"\bskills?\b", re.IGNORECASE), 4.0), + ( + re.compile( + r"\b(set(ting)? up|setup|install(ing)?|configur(e|ing)|deploy(ing)?" + r"|migrat(e|ing)|provision(ing)?|onboard(ing)?|rotate|troubleshoot(ing)?)\b", + re.IGNORECASE, + ), + 2.0, + ), +] + +_GATE_THRESHOLD = 3.0 + + +@dataclass +class GateResult: + """Gate decision with the score and matched fragments for observability.""" + + fired: bool + score: float = 0.0 + matched: list[str] = field(default_factory=list) + + +def skill_gate_enabled() -> bool: + """``SKILL_GATE_ENABLED`` env flag; on unless explicitly disabled.""" + return os.getenv("SKILL_GATE_ENABLED", "true").strip().lower() not in ("false", "0", "no") + + +def should_search_skills(query: str) -> GateResult: + """Decide whether ``query`` warrants a skill lookup. Pure function, no I/O. + + Every rule whose pattern matches (and is not negated, same suppression as + the recall query router) adds its weight; the gate fires when the total + reaches the threshold. + """ + q = (query or "").strip() + if not q: + return GateResult(fired=False) + + score = 0.0 + matched: list[str] = [] + for pattern, weight in _GATE_RULES: + match = pattern.search(q) + if match and not _is_negated(q, match): + score += weight + matched.append(match.group(0)) + + fired = score >= _GATE_THRESHOLD + if fired: + logger.info("skill_gate fired: score=%.1f matched=%s query=%r", score, matched, q) + return GateResult(fired=fired, score=score, matched=matched) diff --git a/cognee/api/v1/search/search.py b/cognee/api/v1/search/search.py index 3291d5932f..4f494ac6a3 100644 --- a/cognee/api/v1/search/search.py +++ b/cognee/api/v1/search/search.py @@ -337,14 +337,19 @@ async def search( ): operation_context.set_dataset(target_dataset_ids[0]) - if query_type is SearchType.AGENTIC_COMPLETION: + if query_type in (SearchType.AGENTIC_COMPLETION, SearchType.SKILLS): active_dataset_refs = dataset_ids if dataset_ids else datasets if isinstance(active_dataset_refs, UUID): active_dataset_refs = [active_dataset_refs] if not active_dataset_refs or len(active_dataset_refs) != 1: + if query_type is SearchType.AGENTIC_COMPLETION: + raise CogneeValidationError( + message="Agentic skill search requires exactly one explicit dataset.", + name="InvalidAgenticDatasetScope", + ) raise CogneeValidationError( - message="Agentic skill search requires exactly one explicit dataset.", - name="InvalidAgenticDatasetScope", + message="SKILLS search requires exactly one explicit dataset.", + name="InvalidSkillsDatasetScope", ) if any(v is not None for v in agentic_overrides.values()): diff --git a/cognee/cli/config.py b/cognee/cli/config.py index ca91cdbe9e..59743550c0 100644 --- a/cognee/cli/config.py +++ b/cognee/cli/config.py @@ -44,6 +44,7 @@ "CODE", "CYPHER", "GRAPH_REPORT", + "SKILLS", ] DEFAULT_SEARCH_TYPE = "HYBRID_COMPLETION" diff --git a/cognee/modules/recall/methods/normalize_search_payload.py b/cognee/modules/recall/methods/normalize_search_payload.py index c261e247a7..48b0b24904 100644 --- a/cognee/modules/recall/methods/normalize_search_payload.py +++ b/cognee/modules/recall/methods/normalize_search_payload.py @@ -36,6 +36,7 @@ SearchType.CHUNKS_LEXICAL: SearchResultKind.CHUNK, SearchType.SUMMARIES: SearchResultKind.SUMMARY, SearchType.AGENTIC_COMPLETION: SearchResultKind.GRAPH_COMPLETION, + SearchType.SKILLS: SearchResultKind.SKILL, } diff --git a/cognee/modules/recall/types/RecallResponse.py b/cognee/modules/recall/types/RecallResponse.py index 52f014971e..473d4f5204 100644 --- a/cognee/modules/recall/types/RecallResponse.py +++ b/cognee/modules/recall/types/RecallResponse.py @@ -53,6 +53,22 @@ class ResponseToolEntry(BaseModel): structured: Optional[dict] = None +class ResponseSkillEntry(BaseModel): + """One skill surfaced by the deterministic skill gate. + + Metadata-only: ``skill`` carries the projected Skill fields and never the + procedure body — progressive disclosure keeps bodies behind the + ``load_skill`` tool or ``GET /skills/{skill_id}``. ``text`` is a + renderable "name: description" line; ``score`` is the raw vector distance + (lower is better) when available. + """ + + source: Literal["skills"] + text: str + skill: dict + score: Optional[float] = None + + class ResponseMarkerEntry(BaseModel): """System-generated marker (not data), e.g. "memory still warming up". @@ -74,6 +90,7 @@ class ResponseMarkerEntry(BaseModel): | ResponseGraphEntry | ResponseCodeEntry | ResponseToolEntry + | ResponseSkillEntry | ResponseMarkerEntry, Field(discriminator="source"), ] diff --git a/cognee/modules/recall/types/SearchResultItem.py b/cognee/modules/recall/types/SearchResultItem.py index 7301008f81..adc080631a 100644 --- a/cognee/modules/recall/types/SearchResultItem.py +++ b/cognee/modules/recall/types/SearchResultItem.py @@ -37,6 +37,7 @@ class SearchResultKind(str, Enum): SUMMARY = "summary" CODING_RULE = "coding_rule" CODE = "code" + SKILL = "skill" NATURAL_LANGUAGE = "natural_language" TEMPORAL = "temporal" STRUCTURED = "structured" # when a response_model was supplied diff --git a/cognee/modules/retrieval/skills_retriever.py b/cognee/modules/retrieval/skills_retriever.py new file mode 100644 index 0000000000..3bae260cca --- /dev/null +++ b/cognee/modules/retrieval/skills_retriever.py @@ -0,0 +1,175 @@ +from typing import Any, List, Optional, Union +from uuid import UUID + +from cognee.shared.logging_utils import get_logger +from cognee.infrastructure.databases.unified import get_unified_engine +from cognee.modules.retrieval.base_retriever import BaseRetriever +from cognee.modules.retrieval.exceptions.exceptions import QueryValidationError +from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError + +logger = get_logger("SkillsRetriever") + +SKILL_COLLECTION = "Skill_search_text" + +# Strict dataset/active filtering shrinks the candidate set, so fetch more +# than top_k from the vector engine and trim after filtering. +_FETCH_MULTIPLIER = 4 +_MIN_FETCH_LIMIT = 20 + + +def _project_skill_payload(payload: dict) -> dict: + """Project a Skill vector payload onto the metadata-only skill shape. + + Mirrors ``list_skills._skill_to_dict`` but works on raw payload dicts and + never includes ``procedure`` / ``skill_text`` / ``search_text``: search + results must preserve the progressive-disclosure design — bodies load via + the ``load_skill`` tool or ``GET /skills/{skill_id}``. ``.get()`` with + defaults keeps skills ingested before newer fields existed from crashing. + """ + return { + "id": str(payload.get("id") or ""), + "name": payload.get("name") or "", + "description": payload.get("description") or "", + "maintainer": payload.get("maintainer") or "", + "maintainer_url": payload.get("maintainer_url") or "", + "version": payload.get("skill_version") or "", + "tags": list(payload.get("tags") or []), + "license": payload.get("license") or "", + "declared_tools": list(payload.get("declared_tools") or []), + "dataset_scope": [str(entry) for entry in (payload.get("dataset_scope") or [])], + "is_active": bool(payload.get("is_active", True)), + "source_repo_url": payload.get("source_repo_url") or "", + "source_dir": payload.get("source_dir") or "", + } + + +class SkillsRetriever(BaseRetriever): + """ + Retriever for semantic discovery of dataset-scoped Skill playbooks. + + Searches the ``Skill_search_text`` vector collection (populated by skill + ingestion via the standard Embeddable indexing path) and returns + metadata-only skill payloads — never the procedure body. + + Requires exactly one explicit dataset: only skills whose ``dataset_scope`` + contains that dataset (and that are ``is_active``) are returned. Skills + with an empty ``dataset_scope`` are excluded. + + A missing collection returns an empty result instead of raising + ``NoDataError`` (unlike SummariesRetriever): "no skills ingested yet" is a + normal state, and the recall skill gate must degrade to a no-op. + + Public methods: + - __init__ + - get_retrieved_objects + - get_context_from_objects + - get_completion_from_context + """ + + # Deterministic, non-generative search type: skip the conversational + # session analysis (which may call an LLM before retrieval). + supports_session_turn_preparation = False + + def __init__( + self, + top_k: Optional[int] = 5, + dataset_id: Optional[Union[str, UUID]] = None, + session_id: Optional[str] = None, + ): + """Initialize retriever with search parameters. ``dataset_id`` is required.""" + if dataset_id is None: + raise QueryValidationError( + message="SKILLS search requires exactly one explicit dataset." + ) + self.top_k = top_k if top_k is not None else 5 + self.dataset_id = str(dataset_id) + self.session_id = session_id + + async def get_retrieved_objects(self, query: str) -> Any: + """ + Retrieves skill hits for the query, filtered to this dataset. + + Over-fetches from the vector engine, then keeps only active skills + whose ``dataset_scope`` contains the retriever's dataset, deduplicated + by id and trimmed to ``top_k``. + """ + logger.info( + f"Starting skill retrieval for query: '{query[:100]}{'...' if len(query) > 100 else ''}'" + ) + + unified = await get_unified_engine() + vector_engine = unified.vector + + fetch_limit = max(self.top_k * _FETCH_MULTIPLIER, _MIN_FETCH_LIMIT) + + try: + results = await vector_engine.search( + SKILL_COLLECTION, query, limit=fetch_limit, include_payload=True + ) + except CollectionNotFoundError: + logger.info( + "%s collection not found — no skills ingested yet; returning no results", + SKILL_COLLECTION, + ) + return [] + + filtered = [] + seen_ids: set = set() + for result in results: + payload = getattr(result, "payload", None) or {} + if not payload.get("is_active", True): + continue + scope = [str(entry) for entry in (payload.get("dataset_scope") or [])] + if self.dataset_id not in scope: + continue + payload_id = str(payload.get("id") or getattr(result, "id", "") or "") + if payload_id: + if payload_id in seen_ids: + continue + seen_ids.add(payload_id) + filtered.append(result) + if len(filtered) >= self.top_k: + break + + logger.info(f"Found {len(filtered)} in-scope skill(s) from vector search") + return filtered + + async def get_context_from_objects(self, query: str, retrieved_objects: Any) -> str: + """ + Formats retrieved skills as a name + description listing. + + Same shape the agentic retriever puts in its system prompt, so the + context is directly usable to offer skills to an LLM without leaking + procedure bodies. + """ + if not retrieved_objects: + return "" + lines = [] + for result in retrieved_objects: + payload = getattr(result, "payload", None) or {} + name = payload.get("name") or "" + description = payload.get("description") or "" + lines.append(f"- `{name}`: {description}") + return "\n".join(lines) + + async def get_completion_from_context( + self, query: str, retrieved_objects: Any, context: Any + ) -> Union[List[str], List[dict]]: + """ + Returns metadata-only skill payloads; no LLM completion is generated. + + Each dict carries the projected skill fields plus the vector ``score`` + (raw backend distance — lower is better). + """ + if not retrieved_objects: + return [] + completions = [] + for result in retrieved_objects: + payload = getattr(result, "payload", None) or {} + projected = _project_skill_payload(payload) + score = getattr(result, "score", None) + if isinstance(score, (int, float)) and not isinstance(score, bool): + projected["score"] = float(score) + completions.append(projected) + logger.info(f"Returning {len(completions)} skill payload(s)") + return completions diff --git a/cognee/modules/search/methods/get_search_type_retriever_instance.py b/cognee/modules/search/methods/get_search_type_retriever_instance.py index b6594a93d1..0aafbbe29a 100644 --- a/cognee/modules/search/methods/get_search_type_retriever_instance.py +++ b/cognee/modules/search/methods/get_search_type_retriever_instance.py @@ -35,6 +35,7 @@ from cognee.modules.retrieval.agentic_retriever import AgenticRetriever from cognee.modules.retrieval.code_retriever import CodeRetriever from cognee.modules.retrieval.graph_report_retriever import GraphReportRetriever +from cognee.modules.retrieval.skills_retriever import SkillsRetriever from cognee.context_global_variables import session_user @@ -91,6 +92,7 @@ async def get_search_type_retriever_instance( neighborhood_depth = kwargs.get("neighborhood_depth") neighborhood_seed_top_k = kwargs.get("neighborhood_seed_top_k") include_references = kwargs.get("include_references", False) + dataset = kwargs.get("dataset") # Registry mapping search types to their corresponding retriever classes and input parameters search_core_registry: dict[SearchType, Tuple[BaseRetriever, dict]] = { @@ -99,6 +101,16 @@ async def get_search_type_retriever_instance( {"config": retriever_specific_config}, ), SearchType.SUMMARIES: (SummariesRetriever, {"top_k": top_k, "session_id": session_id}), + SearchType.SKILLS: ( + SkillsRetriever, + { + "top_k": top_k, + # SKILLS is single-dataset by invariant; SkillsRetriever raises + # QueryValidationError when no dataset reaches the factory. + "dataset_id": dataset.id if dataset is not None else None, + "session_id": session_id, + }, + ), SearchType.CHUNKS: ( ChunksRetriever, { @@ -350,7 +362,6 @@ async def get_search_type_retriever_instance( ) if query_type is SearchType.AGENTIC_COMPLETION: - dataset = kwargs.get("dataset") dataset_id = dataset.id if dataset is not None else None user = kwargs.get("user") try: diff --git a/cognee/modules/search/types/SearchType.py b/cognee/modules/search/types/SearchType.py index 4f33af0e42..25d67ad08f 100644 --- a/cognee/modules/search/types/SearchType.py +++ b/cognee/modules/search/types/SearchType.py @@ -21,3 +21,4 @@ class SearchType(str, Enum): AGENTIC_COMPLETION = "AGENTIC_COMPLETION" CODE = "CODE" GRAPH_REPORT = "GRAPH_REPORT" + SKILLS = "SKILLS" diff --git a/cognee/tests/unit/api/v1/recall/test_skill_gate.py b/cognee/tests/unit/api/v1/recall/test_skill_gate.py new file mode 100644 index 0000000000..14d6b69505 --- /dev/null +++ b/cognee/tests/unit/api/v1/recall/test_skill_gate.py @@ -0,0 +1,276 @@ +import importlib +import types +from uuid import uuid4 + +import pytest + +from cognee.api.v1.recall.skill_gate import ( + should_search_skills, + skill_gate_enabled, +) +from cognee.modules.search.types import SearchType + + +# ── gate classification: pure regex, no LLM ────────────────────────────────── + + +@pytest.mark.parametrize( + "query", + [ + "how do I deploy to staging", + "how to rotate the API keys", + "steps to onboard a new tenant", + "what is the process for a release", + "walk me through the database migration", + "is there a runbook for incident response", + "setup guide for the staging cluster", + "which skills are available", + ], +) +def test_gate_fires_on_procedural_queries(query): + assert should_search_skills(query).fired + + +@pytest.mark.parametrize( + "query", + [ + "what is our churn rate", + "who owns the billing service", + "deploy notes from yesterday", + "summary of last week's incidents", + "", + ], +) +def test_gate_stays_closed_on_non_procedural_queries(query): + assert not should_search_skills(query).fired + + +def test_gate_negation_suppresses_match(): + result = should_search_skills("do not walk me through it") + assert not result.fired + + +def test_gate_result_carries_score_and_matches(): + result = should_search_skills("how do I deploy to staging") + assert result.score >= 3.0 + assert any("how do" in fragment.lower() for fragment in result.matched) + + +def test_gate_enabled_flag(monkeypatch): + monkeypatch.delenv("SKILL_GATE_ENABLED", raising=False) + assert skill_gate_enabled() is True + monkeypatch.setenv("SKILL_GATE_ENABLED", "false") + assert skill_gate_enabled() is False + monkeypatch.setenv("SKILL_GATE_ENABLED", "0") + assert skill_gate_enabled() is False + monkeypatch.setenv("SKILL_GATE_ENABLED", "true") + assert skill_gate_enabled() is True + + +# ── recall() wiring: the gate appends source="skills" entries ───────────────── + + +def _make_user(): + return types.SimpleNamespace(id=uuid4(), tenant_id=None) + + +def _skill_item(**overrides): + item = { + "id": str(uuid4()), + "name": "deploy-checklist", + "description": "Steps to deploy to staging", + "score": 0.3, + } + item.update(overrides) + return item + + +@pytest.fixture +def api_recall_mod(): + return importlib.import_module("cognee.api.v1.recall.recall") + + +def _patch_recall_plumbing(monkeypatch, api_recall_mod, search_calls, skill_items): + """Patch recall's external dependencies; record authorized_search calls.""" + + async def dummy_set_session_user_context_variable(_user): + return None + + async def dummy_authorized_search(**kwargs): + from cognee.modules.search.models.SearchResultPayload import SearchResultPayload + + search_calls.append(kwargs) + if kwargs.get("query_type") is SearchType.SKILLS: + # A real payload: the graph lane normalizes it when SKILLS is the + # explicit query_type; the gate lane only reads .completion. + return [ + SearchResultPayload( + result_object=None, + context=None, + completion=list(skill_items), + search_type=SearchType.SKILLS, + only_context=False, + ) + ] + return [] + + def dummy_get_remote_client(): + return None + + async def dummy_log_search_history(*args, **kwargs): + return None + + monkeypatch.setattr( + api_recall_mod, + "set_session_user_context_variable", + dummy_set_session_user_context_variable, + ) + serve_state = importlib.import_module("cognee.api.v1.serve.state") + search_methods = importlib.import_module("cognee.modules.search.methods.search") + search_operations = importlib.import_module("cognee.modules.search.operations") + monkeypatch.setattr(serve_state, "get_remote_client", dummy_get_remote_client) + monkeypatch.setattr(search_methods, "authorized_search", dummy_authorized_search) + monkeypatch.setattr(search_operations, "log_search_history", dummy_log_search_history) + + +@pytest.mark.asyncio +async def test_gate_appends_skill_entries_for_procedural_query(monkeypatch, api_recall_mod): + user = _make_user() + dataset_id = uuid4() + search_calls = [] + _patch_recall_plumbing(monkeypatch, api_recall_mod, search_calls, [_skill_item()]) + + out = await api_recall_mod.recall( + query_text="how do I deploy to staging", + query_type=SearchType.GRAPH_COMPLETION, + dataset_ids=[dataset_id], + auto_route=False, + user=user, + ) + + called_types = [call.get("query_type") for call in search_calls] + assert SearchType.SKILLS in called_types + skills_call = next(call for call in search_calls if call.get("query_type") is SearchType.SKILLS) + assert skills_call["dataset_ids"] == [dataset_id] + + skill_entries = [entry for entry in out if getattr(entry, "source", None) == "skills"] + assert len(skill_entries) == 1 + entry = skill_entries[0] + assert entry.text == "deploy-checklist: Steps to deploy to staging" + assert entry.score == 0.3 + assert entry.skill["name"] == "deploy-checklist" + assert "score" not in entry.skill + + +@pytest.mark.asyncio +async def test_gate_skipped_for_non_procedural_query(monkeypatch, api_recall_mod): + user = _make_user() + search_calls = [] + _patch_recall_plumbing(monkeypatch, api_recall_mod, search_calls, [_skill_item()]) + + out = await api_recall_mod.recall( + query_text="who owns the billing service", + query_type=SearchType.GRAPH_COMPLETION, + dataset_ids=[uuid4()], + auto_route=False, + user=user, + ) + + assert SearchType.SKILLS not in [call.get("query_type") for call in search_calls] + assert out == [] + + +@pytest.mark.asyncio +async def test_gate_skipped_when_disabled(monkeypatch, api_recall_mod): + monkeypatch.setenv("SKILL_GATE_ENABLED", "false") + user = _make_user() + search_calls = [] + _patch_recall_plumbing(monkeypatch, api_recall_mod, search_calls, [_skill_item()]) + + await api_recall_mod.recall( + query_text="how do I deploy to staging", + query_type=SearchType.GRAPH_COMPLETION, + dataset_ids=[uuid4()], + auto_route=False, + user=user, + ) + + assert SearchType.SKILLS not in [call.get("query_type") for call in search_calls] + + +@pytest.mark.asyncio +async def test_gate_skipped_without_exactly_one_dataset(monkeypatch, api_recall_mod): + """The skill invariant: lookup requires exactly one dataset — else skip silently.""" + user = _make_user() + search_calls = [] + _patch_recall_plumbing(monkeypatch, api_recall_mod, search_calls, [_skill_item()]) + + await api_recall_mod.recall( + query_text="how do I deploy to staging", + query_type=SearchType.GRAPH_COMPLETION, + dataset_ids=[uuid4(), uuid4()], + auto_route=False, + user=user, + ) + + assert SearchType.SKILLS not in [call.get("query_type") for call in search_calls] + + +@pytest.mark.asyncio +async def test_gate_bypassed_for_explicit_skills_query_type(monkeypatch, api_recall_mod): + user = _make_user() + search_calls = [] + _patch_recall_plumbing(monkeypatch, api_recall_mod, search_calls, [_skill_item()]) + + await api_recall_mod.recall( + query_text="how do I deploy to staging", + query_type=SearchType.SKILLS, + dataset_ids=[uuid4()], + auto_route=False, + user=user, + ) + + # Exactly one SKILLS call — the graph lane's own — not a second gate call. + assert [call.get("query_type") for call in search_calls].count(SearchType.SKILLS) == 1 + + +@pytest.mark.asyncio +async def test_gate_failure_never_fails_recall(monkeypatch, api_recall_mod): + """A gate-lane exception is swallowed; the main lanes still answer.""" + user = _make_user() + search_calls = [] + + async def dummy_set_session_user_context_variable(_user): + return None + + async def dummy_authorized_search(**kwargs): + search_calls.append(kwargs) + if kwargs.get("query_type") is SearchType.SKILLS: + raise RuntimeError("skill lane exploded") + return [] + + async def dummy_log_search_history(*args, **kwargs): + return None + + monkeypatch.setattr( + api_recall_mod, + "set_session_user_context_variable", + dummy_set_session_user_context_variable, + ) + serve_state = importlib.import_module("cognee.api.v1.serve.state") + search_methods = importlib.import_module("cognee.modules.search.methods.search") + search_operations = importlib.import_module("cognee.modules.search.operations") + monkeypatch.setattr(serve_state, "get_remote_client", lambda: None) + monkeypatch.setattr(search_methods, "authorized_search", dummy_authorized_search) + monkeypatch.setattr(search_operations, "log_search_history", dummy_log_search_history) + + out = await api_recall_mod.recall( + query_text="how do I deploy to staging", + query_type=SearchType.GRAPH_COMPLETION, + dataset_ids=[uuid4()], + auto_route=False, + user=user, + ) + + assert out == [] + assert SearchType.SKILLS in [call.get("query_type") for call in search_calls] diff --git a/cognee/tests/unit/modules/retrieval/skills_retriever_test.py b/cognee/tests/unit/modules/retrieval/skills_retriever_test.py new file mode 100644 index 0000000000..71d745e4e5 --- /dev/null +++ b/cognee/tests/unit/modules/retrieval/skills_retriever_test.py @@ -0,0 +1,221 @@ +import pytest +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 + +from cognee.modules.retrieval.skills_retriever import SkillsRetriever +from cognee.modules.retrieval.exceptions.exceptions import QueryValidationError +from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError + +DATASET_ID = str(uuid4()) +OTHER_DATASET_ID = str(uuid4()) + + +def _make_unified_mock(vector_engine): + """Create a mock unified engine that exposes the given vector engine.""" + unified = AsyncMock() + unified.vector = vector_engine + unified.graph = AsyncMock() + return unified + + +@pytest.fixture +def mock_vector_engine(): + """Create a mock vector engine.""" + engine = AsyncMock() + engine.search = AsyncMock() + return engine + + +def _skill_payload(**overrides): + """A full Skill vector payload, procedure body included.""" + payload = { + "id": str(uuid4()), + "name": "deploy-checklist", + "description": "Steps to deploy to staging", + "procedure": "SECRET PROCEDURE BODY", + "skill_text": "deploy-checklist\n\nSECRET PROCEDURE BODY", + "search_text": "deploy-checklist\n\nSECRET PROCEDURE BODY", + "declared_tools": ["bash"], + "maintainer": "ops-team", + "maintainer_url": "https://example.com", + "skill_version": "1.2", + "tags": ["ops"], + "license": "MIT", + "source_repo_url": "", + "source_dir": "skills/deploy-checklist", + "dataset_scope": [DATASET_ID], + "is_active": True, + } + payload.update(overrides) + return payload + + +def _result(payload, score=0.1): + return SimpleNamespace(id=payload.get("id"), payload=payload, score=score) + + +def _patch_engine(monkeypatch, vector_engine): + import cognee.modules.retrieval.skills_retriever as mod + + async def _get_unified_engine(): + return _make_unified_mock(vector_engine) + + monkeypatch.setattr(mod, "get_unified_engine", _get_unified_engine) + + +def test_init_requires_dataset_id(): + """SKILLS search is single-dataset by invariant; no dataset is an error.""" + with pytest.raises(QueryValidationError, match="exactly one explicit dataset"): + SkillsRetriever() + + +def test_init_defaults(): + retriever = SkillsRetriever(dataset_id=DATASET_ID) + assert retriever.top_k == 5 + assert retriever.dataset_id == DATASET_ID + + +def test_init_custom_top_k_and_uuid_dataset(): + dataset_uuid = uuid4() + retriever = SkillsRetriever(top_k=10, dataset_id=dataset_uuid) + assert retriever.top_k == 10 + assert retriever.dataset_id == str(dataset_uuid) + + +def test_init_top_k_none_falls_back_to_default(): + retriever = SkillsRetriever(top_k=None, dataset_id=DATASET_ID) + assert retriever.top_k == 5 + + +@pytest.mark.asyncio +async def test_get_objects_filters_scope_and_active(monkeypatch, mock_vector_engine): + """Only active skills scoped to this dataset survive; empty scope is excluded.""" + in_scope = _skill_payload(name="in-scope") + inactive = _skill_payload(name="inactive", is_active=False) + other_dataset = _skill_payload(name="other-dataset", dataset_scope=[OTHER_DATASET_ID]) + empty_scope = _skill_payload(name="empty-scope", dataset_scope=[]) + + mock_vector_engine.search.return_value = [ + _result(in_scope), + _result(inactive), + _result(other_dataset), + _result(empty_scope), + ] + _patch_engine(monkeypatch, mock_vector_engine) + + retriever = SkillsRetriever(dataset_id=DATASET_ID) + objects = await retriever.get_retrieved_objects("deploy") + + assert [obj.payload["name"] for obj in objects] == ["in-scope"] + + +@pytest.mark.asyncio +async def test_get_objects_overfetches_and_trims_to_top_k(monkeypatch, mock_vector_engine): + """Fetch limit exceeds top_k (filtering shrinks results); output is trimmed.""" + payloads = [_skill_payload(name=f"skill-{i}") for i in range(5)] + mock_vector_engine.search.return_value = [_result(p) for p in payloads] + _patch_engine(monkeypatch, mock_vector_engine) + + retriever = SkillsRetriever(top_k=2, dataset_id=DATASET_ID) + objects = await retriever.get_retrieved_objects("deploy") + + assert len(objects) == 2 + mock_vector_engine.search.assert_awaited_once_with( + "Skill_search_text", "deploy", limit=20, include_payload=True + ) + + +@pytest.mark.asyncio +async def test_get_objects_dedupes_by_id(monkeypatch, mock_vector_engine): + payload = _skill_payload(name="dup") + mock_vector_engine.search.return_value = [_result(payload), _result(payload)] + _patch_engine(monkeypatch, mock_vector_engine) + + retriever = SkillsRetriever(dataset_id=DATASET_ID) + objects = await retriever.get_retrieved_objects("deploy") + + assert len(objects) == 1 + + +@pytest.mark.asyncio +async def test_get_objects_collection_not_found_returns_empty(monkeypatch, mock_vector_engine): + """No skills ingested yet is a normal state — no NoDataError (unlike SUMMARIES).""" + mock_vector_engine.search.side_effect = CollectionNotFoundError("Collection not found") + _patch_engine(monkeypatch, mock_vector_engine) + + retriever = SkillsRetriever(dataset_id=DATASET_ID) + objects = await retriever.get_retrieved_objects("deploy") + + assert objects == [] + + +@pytest.mark.asyncio +async def test_get_objects_empty_results(monkeypatch, mock_vector_engine): + mock_vector_engine.search.return_value = [] + _patch_engine(monkeypatch, mock_vector_engine) + + retriever = SkillsRetriever(dataset_id=DATASET_ID) + objects = await retriever.get_retrieved_objects("deploy") + + assert objects == [] + + +@pytest.mark.asyncio +async def test_context_lists_names_and_descriptions_only(): + retriever = SkillsRetriever(dataset_id=DATASET_ID) + objects = [_result(_skill_payload(name="deploy-checklist"))] + + context = await retriever.get_context_from_objects("deploy", objects) + + assert context == "- `deploy-checklist`: Steps to deploy to staging" + assert "SECRET PROCEDURE BODY" not in context + + +@pytest.mark.asyncio +async def test_context_empty_without_objects(): + retriever = SkillsRetriever(dataset_id=DATASET_ID) + assert await retriever.get_context_from_objects("deploy", []) == "" + + +@pytest.mark.asyncio +async def test_completion_projection_is_metadata_only(): + """Results never expose the procedure body — progressive disclosure.""" + retriever = SkillsRetriever(dataset_id=DATASET_ID) + payload = _skill_payload() + objects = [_result(payload, score=0.42)] + + completion = await retriever.get_completion_from_context("deploy", objects, "") + + assert len(completion) == 1 + projected = completion[0] + assert projected["name"] == "deploy-checklist" + assert projected["description"] == "Steps to deploy to staging" + assert projected["version"] == "1.2" + assert projected["declared_tools"] == ["bash"] + assert projected["dataset_scope"] == [DATASET_ID] + assert projected["score"] == 0.42 + for stripped_field in ("procedure", "skill_text", "search_text"): + assert stripped_field not in projected + + +@pytest.mark.asyncio +async def test_completion_tolerates_sparse_legacy_payload(): + """Skills ingested before newer fields existed must not crash the projection.""" + retriever = SkillsRetriever(dataset_id=DATASET_ID) + sparse = {"id": str(uuid4()), "name": "old-skill", "dataset_scope": [DATASET_ID]} + objects = [SimpleNamespace(id=sparse["id"], payload=sparse, score=None)] + + completion = await retriever.get_completion_from_context("deploy", objects, "") + + assert completion[0]["name"] == "old-skill" + assert completion[0]["description"] == "" + assert completion[0]["tags"] == [] + assert completion[0]["is_active"] is True + assert "score" not in completion[0] + + +@pytest.mark.asyncio +async def test_completion_empty_without_objects(): + retriever = SkillsRetriever(dataset_id=DATASET_ID) + assert await retriever.get_completion_from_context("deploy", [], "") == [] diff --git a/cognee/tests/unit/modules/search/test_get_search_type_retriever_instance.py b/cognee/tests/unit/modules/search/test_get_search_type_retriever_instance.py index 6d829c0d85..3965fbce08 100644 --- a/cognee/tests/unit/modules/search/test_get_search_type_retriever_instance.py +++ b/cognee/tests/unit/modules/search/test_get_search_type_retriever_instance.py @@ -359,6 +359,34 @@ async def search(collection_name, *args, **kwargs): assert count_retrieved_objects(payload.result_object) == 2 +@pytest.mark.asyncio +async def test_skills_retriever_registered_with_dataset(): + import types + from uuid import uuid4 + + import cognee.modules.search.methods.get_search_type_retriever_instance as mod + from cognee.modules.retrieval.skills_retriever import SkillsRetriever + + dataset = types.SimpleNamespace(id=uuid4()) + + retriever_instance = await mod.get_search_type_retriever_instance( + SearchType.SKILLS, query_text="how do I deploy", top_k=4, dataset=dataset + ) + + assert isinstance(retriever_instance, SkillsRetriever) + assert retriever_instance.top_k == 4 + assert retriever_instance.dataset_id == str(dataset.id) + + +@pytest.mark.asyncio +async def test_skills_requires_dataset(): + import cognee.modules.search.methods.get_search_type_retriever_instance as mod + from cognee.modules.retrieval.exceptions.exceptions import QueryValidationError + + with pytest.raises(QueryValidationError, match="exactly one explicit dataset"): + await mod.get_search_type_retriever_instance(SearchType.SKILLS, query_text="q") + + @pytest.mark.asyncio async def test_chunks_lexical_returns_bm25_retriever(): import cognee.modules.search.methods.get_search_type_retriever_instance as mod