code-context-engine 0.4.0__py3-none-any.whl

context_engine/project_commands.py

@@ -0,0 +1,296 @@
+"""Project-specific commands, rules, and preferences loaded at session start.
+
+Supports two levels:
+- **Workspace** (optional): parent directory's .cce/commands.yaml — global
+  defaults that apply to all projects under it.
+- **Project**: the project's own .cce/commands.yaml — extends or overrides
+  the workspace config.
+
+Example .cce/commands.yaml:
+    rules:
+      - NEVER generate down() in migrations — forward-only
+      - Use UUID for primary keys
+    preferences:
+      database: PostgreSQL
+      auth: Sanctum
+      style: "Clean architecture"
+    before_push:
+      - composer test
+      - phpstan analyse
+    before_commit:
+      - php-cs-fixer fix --dry-run
+    on_start:
+      - echo "Deploy freeze until Friday"
+    custom:
+      deploy: kubectl apply -f k8s/
+"""
+import logging
+from pathlib import Path
+
+import yaml
+
+log = logging.getLogger(__name__)
+
+COMMANDS_DIR = ".cce"
+COMMANDS_FILE = "commands.yaml"
+
+VALID_HOOKS = {"before_push", "before_commit", "on_start", "custom"}
+# Sections that are lists (merged by appending, deduped)
+_LIST_SECTIONS = {"rules", "before_push", "before_commit", "on_start"}
+# Sections that are dicts (merged by update)
+_DICT_SECTIONS = {"preferences", "custom"}
+
+
+def _commands_path(project_dir: str) -> Path:
+    return Path(project_dir) / COMMANDS_DIR / COMMANDS_FILE
+
+
+def _load_yaml(path: Path) -> dict:
+    """Load a YAML file. Returns {} on any error."""
+    if not path.exists():
+        return {}
+    try:
+        data = yaml.safe_load(path.read_text()) or {}
+    except (yaml.YAMLError, OSError) as exc:
+        log.warning("Failed to parse %s: %s", path, exc)
+        return {}
+    if not isinstance(data, dict):
+        log.warning("%s is not a valid YAML mapping", path)
+        return {}
+    return data
+
+
+def _find_workspace_dir(project_dir: str) -> Path | None:
+    """Find the nearest parent with .cce/commands.yaml (not project_dir itself)."""
+    current = Path(project_dir).resolve().parent
+    home = Path.home()
+    # Walk up but stop at home directory (don't scan /Users or /)
+    while current != current.parent and current != home.parent:
+        candidate = current / COMMANDS_DIR / COMMANDS_FILE
+        if candidate.exists():
+            return current
+        current = current.parent
+    return None
+
+
+def _merge_configs(workspace: dict, project: dict) -> dict:
+    """Merge workspace config into project config. Project wins on conflicts."""
+    merged = {}
+    all_keys = set(workspace.keys()) | set(project.keys())
+    for key in all_keys:
+        ws_val = workspace.get(key)
+        pj_val = project.get(key)
+        if key in _LIST_SECTIONS:
+            # Merge lists, project items come after workspace, deduplicate
+            ws_list = ws_val if isinstance(ws_val, list) else []
+            pj_list = pj_val if isinstance(pj_val, list) else []
+            merged_list = []
+            seen_strs: set[str] = set()
+            for item in ws_list + pj_list:
+                item_key = str(item)
+                if item_key not in seen_strs:
+                    seen_strs.add(item_key)
+                    merged_list.append(item)
+            if merged_list:
+                merged[key] = merged_list
+        elif key in _DICT_SECTIONS:
+            # Merge dicts, project overrides workspace
+            ws_dict = ws_val if isinstance(ws_val, dict) else {}
+            pj_dict = pj_val if isinstance(pj_val, dict) else {}
+            combined = {**ws_dict, **pj_dict}
+            if combined:
+                merged[key] = combined
+        else:
+            # Unknown section: project wins, fallback to workspace
+            merged[key] = pj_val if pj_val is not None else ws_val
+    return merged
+
+
+def load_commands(project_dir: str) -> dict:
+    """Load merged config: workspace (optional) + project."""
+    project_config = _load_yaml(_commands_path(project_dir))
+    workspace_dir = _find_workspace_dir(project_dir)
+    if workspace_dir is None:
+        return project_config
+    workspace_config = _load_yaml(workspace_dir / COMMANDS_DIR / COMMANDS_FILE)
+    return _merge_configs(workspace_config, project_config)
+
+
+def load_project_only(project_dir: str) -> dict:
+    """Load only the project-level config (no workspace merge)."""
+    return _load_yaml(_commands_path(project_dir))
+
+
+def save_commands(project_dir: str, commands: dict) -> None:
+    """Save project commands to .cce/commands.yaml."""
+    path = _commands_path(project_dir)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(yaml.dump(commands, default_flow_style=False, sort_keys=False))
+
+
+def add_command(project_dir: str, hook: str, command: str) -> None:
+    """Add a command to a hook. Creates the file if it doesn't exist."""
+    if hook not in VALID_HOOKS:
+        raise ValueError(f"Invalid hook '{hook}'. Valid hooks: {', '.join(sorted(VALID_HOOKS))}")
+    if hook == "custom":
+        raise ValueError("Use add_custom_command() for custom commands")
+    commands = load_project_only(project_dir)
+    hook_list = commands.setdefault(hook, [])
+    if not isinstance(hook_list, list):
+        raise ValueError(f"Hook '{hook}' is not a list in commands.yaml")
+    if command in hook_list:
+        return
+    hook_list.append(command)
+    save_commands(project_dir, commands)
+
+
+def add_rule(project_dir: str, rule: str) -> None:
+    """Add a rule. Creates the file if it doesn't exist."""
+    commands = load_project_only(project_dir)
+    rules = commands.setdefault("rules", [])
+    if not isinstance(rules, list):
+        raise ValueError("'rules' section must be a list in commands.yaml")
+    if rule in rules:
+        return
+    rules.append(rule)
+    save_commands(project_dir, commands)
+
+
+def set_preference(project_dir: str, key: str, value: str) -> None:
+    """Set a preference key-value pair."""
+    commands = load_project_only(project_dir)
+    prefs = commands.setdefault("preferences", {})
+    if not isinstance(prefs, dict):
+        raise ValueError("'preferences' section must be a mapping in commands.yaml")
+    prefs[key] = value
+    save_commands(project_dir, commands)
+
+
+def add_custom_command(project_dir: str, name: str, command: str) -> None:
+    """Add a named custom command."""
+    commands = load_project_only(project_dir)
+    custom = commands.setdefault("custom", {})
+    if not isinstance(custom, dict):
+        raise ValueError("'custom' section must be a mapping in commands.yaml")
+    custom[name] = command
+    save_commands(project_dir, commands)
+
+
+def remove_command(project_dir: str, hook: str, command: str) -> bool:
+    """Remove a command from a hook. Returns True if removed."""
+    commands = load_project_only(project_dir)
+    if hook not in commands:
+        return False
+    if hook == "custom":
+        custom = commands.get("custom", {})
+        if command in custom:
+            del custom[command]
+            if not custom:
+                del commands["custom"]
+            save_commands(project_dir, commands)
+            return True
+        return False
+    hook_list = commands.get(hook, [])
+    if not isinstance(hook_list, list):
+        return False
+    if command in hook_list:
+        hook_list.remove(command)
+        if not hook_list:
+            del commands[hook]
+        save_commands(project_dir, commands)
+        return True
+    return False
+
+
+def remove_rule(project_dir: str, rule: str) -> bool:
+    """Remove a rule. Returns True if removed."""
+    commands = load_project_only(project_dir)
+    rules = commands.get("rules", [])
+    if not isinstance(rules, list) or rule not in rules:
+        return False
+    rules.remove(rule)
+    if not rules:
+        del commands["rules"]
+    save_commands(project_dir, commands)
+    return True
+
+
+def remove_preference(project_dir: str, key: str) -> bool:
+    """Remove a preference. Returns True if removed."""
+    commands = load_project_only(project_dir)
+    prefs = commands.get("preferences", {})
+    if not isinstance(prefs, dict) or key not in prefs:
+        return False
+    del prefs[key]
+    if not prefs:
+        del commands["preferences"]
+    save_commands(project_dir, commands)
+    return True
+
+
+_GITIGNORE_ENTRIES = [
+    # CCE local cache and per-machine files
+    (".cce/", "CCE local cache (per-machine, not for version control)"),
+    (".claude/settings.local.json", "Claude Code local settings written by cce init"),
+]
+
+
+def ensure_gitignore(project_dir: str) -> None:
+    """Add CCE-related entries to .gitignore if not already present."""
+    gitignore = Path(project_dir) / ".gitignore"
+    content = gitignore.read_text() if gitignore.exists() else ""
+
+    additions = []
+    for entry, comment in _GITIGNORE_ENTRIES:
+        if entry not in content:
+            additions.append(f"# {comment}\n{entry}")
+
+    if not additions:
+        return
+
+    block = "\n\n# CCE (code-context-engine)\n" + "\n".join(additions) + "\n"
+    gitignore.write_text(content.rstrip() + block)
+
+
+def format_for_prompt(commands: dict, label: str = "Project") -> str:
+    """Format commands as markdown for the init prompt."""
+    if not commands:
+        return ""
+    lines = []
+
+    # Rules
+    rules = commands.get("rules", [])
+    if rules and isinstance(rules, list):
+        lines.append(f"### {label} Rules")
+        for r in rules:
+            lines.append(f"- {r}")
+
+    # Preferences
+    prefs = commands.get("preferences", {})
+    if prefs and isinstance(prefs, dict):
+        lines.append(f"### {label} Preferences")
+        for k, v in prefs.items():
+            lines.append(f"- **{k}:** {v}")
+
+    # Commands
+    hook_labels = {
+        "before_push": "Before push",
+        "before_commit": "Before commit",
+        "on_start": "On session start",
+    }
+    cmd_lines = []
+    for hook, hook_label in hook_labels.items():
+        cmds = commands.get(hook, [])
+        if cmds and isinstance(cmds, list):
+            cmd_str = ", ".join(f"`{c}`" for c in cmds)
+            cmd_lines.append(f"- **{hook_label}:** {cmd_str}")
+    custom = commands.get("custom", {})
+    if custom and isinstance(custom, dict):
+        cmd_lines.append("- **Custom commands:**")
+        for name, cmd in custom.items():
+            cmd_lines.append(f"  - `{name}`: `{cmd}`")
+    if cmd_lines:
+        lines.append(f"### {label} Commands")
+        lines.extend(cmd_lines)
+
+    return "\n".join(lines) if lines else ""

context_engine/retrieval/confidence.py

@@ -0,0 +1,47 @@
+"""Confidence scoring for retrieved chunks.
+
+The score is a weighted sum of three factors, each normalised to [0, 1]:
+
+- vector similarity: 1 - (cosine distance from query embedding).
+- keyword / file-hint match: a lightweight query-parser bonus when the chunk's
+  file path or content hits the parsed query intent. Replaces what used to be
+  labelled "graph hops" before the graph store was removed.
+- recency: exponential decay based on the chunk's `modified_ts` metadata.
+
+The weights live here as module constants so they're easy to find and tune.
+"""
+import time
+from context_engine.models import Chunk
+
+_VECTOR_WEIGHT = 0.5
+_KEYWORD_WEIGHT = 0.3
+_RECENCY_WEIGHT = 0.2
+_MAX_KEYWORD_DISTANCE = 5
+_RECENCY_HALF_LIFE = 7 * 24 * 3600  # 1 week
+
+
+class ConfidenceScorer:
+    def score(
+        self,
+        chunk: Chunk,
+        vector_distance: float,
+        keyword_distance: int,
+    ) -> float:
+        vector_score = max(0.0, 1.0 - vector_distance)
+        keyword_score = max(0.0, 1.0 - (keyword_distance / _MAX_KEYWORD_DISTANCE))
+        recency_score = self._recency_score(chunk)
+        combined = (
+            _VECTOR_WEIGHT * vector_score
+            + _KEYWORD_WEIGHT * keyword_score
+            + _RECENCY_WEIGHT * recency_score
+        )
+        return min(1.0, max(0.0, combined))
+
+    def _recency_score(self, chunk: Chunk) -> float:
+        modified_ts = chunk.metadata.get("modified_ts")
+        if modified_ts is None:
+            return 0.5
+        age_seconds = time.time() - modified_ts
+        if age_seconds <= 0:
+            return 1.0
+        return 0.5 ** (age_seconds / _RECENCY_HALF_LIFE)

context_engine/retrieval/query_parser.py

@@ -0,0 +1,105 @@
+"""Query understanding — intent classification and keyword extraction."""
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class QueryIntent(Enum):
+    CODE_LOOKUP = "code_lookup"
+    DECISION_RECALL = "decision_recall"
+    ARCHITECTURE = "architecture"
+    GENERAL = "general"
+
+
+_DECISION_PATTERNS = [
+    r"what did we decide",
+    r"decision about",
+    r"why did we",
+    r"last session",
+    r"previous discussion",
+    r"agreed on",
+]
+_ARCHITECTURE_PATTERNS = [
+    r"how is .+ structured",
+    r"architecture",
+    r"module.+structure",
+    r"component.+design",
+    r"how does .+ work",
+    r"overview of",
+    r"explain the .+ system",
+]
+_CODE_PATTERNS = [
+    r"find .+ function",
+    r"show me .+ class",
+    r"where is .+ defined",
+    r"implementation of",
+    r"\.py|\.js|\.ts",
+    r"function|class|method|def |import ",
+]
+_FILE_PATH_RE = re.compile(r"[a-zA-Z0-9_./-]+\.[a-zA-Z]{1,10}")
+# Natural-language stop words we always strip.
+_STOP_WORDS = {
+    "the", "a", "an", "is", "are", "was", "were", "do", "does", "did",
+    "what", "how", "why", "where", "when", "who", "which",
+    "in", "on", "at", "to", "for", "of", "with", "about",
+    "me", "my", "we", "our", "it", "its", "i", "you",
+    "tell", "give",
+}
+# Code-flavoured words that look like stop words in prose ("show me get
+# functions") but are critical naming prefixes in code. Strip them when the
+# intent is conversational, keep them when the intent is code lookup so
+# `getUser` / `set_config` / `find_by_id` matches survive keyword extraction.
+_CODE_PREFIX_WORDS = {"show", "find", "get", "set", "fetch", "save", "validate", "create", "update", "delete"}
+
+
+@dataclass
+class ParsedQuery:
+    original: str
+    intent: QueryIntent
+    keywords: list[str] = field(default_factory=list)
+    file_hints: list[str] = field(default_factory=list)
+
+
+class QueryParser:
+    def parse(self, query: str) -> ParsedQuery:
+        lower = query.lower()
+        intent = self._classify_intent(lower)
+        keywords = self._extract_keywords(query, intent=intent)
+        file_hints = _FILE_PATH_RE.findall(query)
+        return ParsedQuery(
+            original=query, intent=intent, keywords=keywords, file_hints=file_hints
+        )
+
+    def _classify_intent(self, query: str) -> QueryIntent:
+        for p in _DECISION_PATTERNS:
+            if re.search(p, query):
+                return QueryIntent.DECISION_RECALL
+        for p in _ARCHITECTURE_PATTERNS:
+            if re.search(p, query):
+                return QueryIntent.ARCHITECTURE
+        for p in _CODE_PATTERNS:
+            if re.search(p, query):
+                return QueryIntent.CODE_LOOKUP
+        return QueryIntent.GENERAL
+
+    def _extract_keywords(
+        self, query: str, intent: QueryIntent = QueryIntent.GENERAL
+    ) -> list[str]:
+        identifiers = re.findall(r"[A-Z][a-zA-Z0-9]+", query)
+        words = re.findall(r"[a-zA-Z_][a-zA-Z0-9_]*", query)
+        # For code-lookup intent, keep prefix words like `get`/`find`/`save`
+        # so the user's literal verb survives into FTS keyword scoring.
+        stop_words = (
+            _STOP_WORDS if intent == QueryIntent.CODE_LOOKUP
+            else _STOP_WORDS | _CODE_PREFIX_WORDS
+        )
+        meaningful = [
+            w for w in words if w.lower() not in stop_words and len(w) > 2
+        ]
+        seen = set()
+        result = []
+        for kw in identifiers + meaningful:
+            if kw not in seen:
+                seen.add(kw)
+                result.append(kw)
+        return result

context_engine/retrieval/retriever.py

@@ -0,0 +1,199 @@
+"""Hybrid retrieval — vector search + FTS BM25 + RRF merging + confidence scoring."""
+import logging
+
+from context_engine.models import Chunk
+from context_engine.storage.backend import StorageBackend
+from context_engine.indexer.embedder import Embedder
+from context_engine.retrieval.confidence import ConfidenceScorer
+from context_engine.retrieval.query_parser import QueryIntent, QueryParser
+
+log = logging.getLogger(__name__)
+
+_DEPRIORITISED_PATHS = {"tests/", "test_", "docs/", "spec", "plan"}
+_RRF_K = 60
+# Confidence weight in the final blend. The remainder goes to RRF, normalised to
+# [0,1] by the best score in the candidate set so an exact-match FTS rank-1 hit
+# scores the same as a vector rank-1 hit instead of being clamped to ~1.0.
+_CONFIDENCE_WEIGHT = 0.5
+# When the parsed query looks like a code lookup, give FTS more pull because
+# exact-identifier hits are usually what the user wants.
+_FTS_BOOST_CODE_LOOKUP = 1.5
+
+
+class HybridRetriever:
+    def __init__(self, backend: StorageBackend, embedder: Embedder) -> None:
+        self._backend = backend
+        self._embedder = embedder
+        self._scorer = ConfidenceScorer()
+        self._parser = QueryParser()
+        self._fts_warned = False
+
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int = 10,
+        confidence_threshold: float = 0.0,
+        max_tokens: int | None = None,
+    ) -> list[Chunk]:
+        parsed = self._parser.parse(query)
+        query_embedding = self._embedder.embed_query(query)
+
+        # embed_query returns tuple for LRU cache hashability; vector_store
+        # now handles the conversion internally via _to_list().
+
+        vector_results = await self._backend.vector_search(
+            query_embedding=query_embedding,
+            top_k=max(top_k * 3, 1),
+        )
+
+        # FTS search with graceful fallback
+        fts_ids: dict[str, int] = {}
+        try:
+            fts_results = await self._backend.fts_search(query, top_k=top_k * 3)
+            fts_ids = {id_: rank for rank, (id_, _) in enumerate(fts_results)}
+        except Exception:
+            if not self._fts_warned:
+                log.warning("FTS search unavailable; falling back to vector-only")
+                self._fts_warned = True
+
+        # Build vector rankings and chunk map
+        vector_ranks: dict[str, int] = {}
+        chunk_map: dict[str, Chunk] = {}
+        seen_keys: set[str] = set()
+
+        for rank, chunk in enumerate(vector_results):
+            dedup_key = f"{chunk.file_path}:{chunk.start_line}-{chunk.end_line}"
+            if dedup_key in seen_keys:
+                continue
+            seen_keys.add(dedup_key)
+            vector_ranks[chunk.id] = rank
+            chunk_map[chunk.id] = chunk
+
+        # Hydrate FTS-only results
+        fts_only_ids = [id_ for id_ in fts_ids if id_ not in chunk_map]
+        if fts_only_ids:
+            try:
+                hydrated = await self._backend.get_chunks_by_ids(fts_only_ids)
+                for chunk in hydrated:
+                    chunk_map[chunk.id] = chunk
+            except Exception as exc:
+                log.warning("Failed to hydrate FTS-only chunks: %s", exc)
+
+        # Compute RRF scores. Boost FTS contribution when the parsed intent
+        # is CODE_LOOKUP — exact identifier matches are almost always what the
+        # user wants and would otherwise be drowned by semantic-similarity hits.
+        fts_weight = (
+            _FTS_BOOST_CODE_LOOKUP if parsed.intent == QueryIntent.CODE_LOOKUP else 1.0
+        )
+        all_ids = set(vector_ranks.keys()) | set(fts_ids.keys())
+        rrf_scores: dict[str, float] = {}
+        for id_ in all_ids:
+            score = 0.0
+            if id_ in vector_ranks:
+                score += 1.0 / (_RRF_K + vector_ranks[id_])
+            if id_ in fts_ids:
+                score += fts_weight * (1.0 / (_RRF_K + fts_ids[id_]))
+            rrf_scores[id_] = score
+
+        # Normalise RRF to [0, 1] by the best score in this candidate set.
+        # The previous `min(rrf * _RRF_K, 1.0)` saturated nearly every result to
+        # ~1.0, so confidence_score dominated the blend and FTS rank carried
+        # almost no signal past the top few. Rank-normalising restores gradient.
+        max_rrf = max(rrf_scores.values()) if rrf_scores else 0.0
+
+        # Score with confidence scorer
+        scored: list[tuple[Chunk, float]] = []
+        for id_, rrf_score in rrf_scores.items():
+            chunk = chunk_map.get(id_)
+            if chunk is None:
+                continue
+
+            distance = chunk.metadata.get("_distance", 0.0)
+            normalised_distance = min(max(distance / 2.0, 0.0), 1.0)
+            keyword_distance = self._estimate_keyword_distance(chunk, parsed)
+            conf_score = self._scorer.score(
+                chunk,
+                vector_distance=normalised_distance,
+                keyword_distance=keyword_distance,
+            )
+
+            normalised_rrf = (rrf_score / max_rrf) if max_rrf > 0 else 0.0
+            final_score = (
+                _CONFIDENCE_WEIGHT * conf_score
+                + (1.0 - _CONFIDENCE_WEIGHT) * normalised_rrf
+            )
+            final_score = self._apply_path_penalty(chunk.file_path, final_score)
+            chunk.confidence_score = final_score
+
+            if final_score >= confidence_threshold:
+                scored.append((chunk, final_score))
+
+        scored.sort(key=lambda x: x[1], reverse=True)
+        ranked = [chunk for chunk, _ in scored[:top_k]]
+
+        # Graph expansion: fetch 1-2 bonus chunks from files reachable via
+        # CALLS/IMPORTS edges from the top results.
+        if ranked and hasattr(self._backend, "get_related_file_paths"):
+            try:
+                top_files = list({c.file_path for c in ranked[:3]})
+                related_files = await self._backend.get_related_file_paths(top_files)
+                qe_list = (
+                    list(query_embedding)
+                    if not isinstance(query_embedding, list)
+                    else query_embedding
+                )
+                for rel_fp in related_files[:2]:  # max 2 bonus files
+                    bonus = await self._backend.vector_search(
+                        query_embedding=qe_list,
+                        top_k=2,
+                        filters={"file_path": rel_fp},
+                    )
+                    for b in bonus:
+                        dedup_key = (
+                            f"{b.file_path}:{b.start_line}-{b.end_line}"
+                        )
+                        if dedup_key not in seen_keys:
+                            seen_keys.add(dedup_key)
+                            dist = b.metadata.get("_distance", 1.0)
+                            b.confidence_score = max(0.0, 1.0 - dist) * 0.85
+                            if b.confidence_score >= confidence_threshold:
+                                ranked.append(b)
+            except Exception as exc:
+                log.debug("Graph expansion skipped: %s", exc)
+
+        if max_tokens is None:
+            return ranked
+
+        packed: list[Chunk] = []
+        budget = max_tokens
+        for chunk in ranked:
+            tokens = chunk.token_count
+            if tokens <= budget:
+                packed.append(chunk)
+                budget -= tokens
+            elif chunk.compressed_content:
+                compressed_tokens = max(1, int(len(chunk.compressed_content) / 3.3))
+                if compressed_tokens <= budget:
+                    packed.append(chunk)
+                    budget -= compressed_tokens
+        return packed
+
+    @staticmethod
+    def _apply_path_penalty(file_path: str, score: float) -> float:
+        if file_path.startswith("git:"):
+            return score
+        fp_lower = file_path.lower()
+        for marker in _DEPRIORITISED_PATHS:
+            if marker in fp_lower:
+                return score * 0.8
+        return score
+
+    def _estimate_keyword_distance(self, chunk, parsed) -> int:
+        if parsed.file_hints:
+            for hint in parsed.file_hints:
+                if hint in chunk.file_path:
+                    return 0
+        for keyword in parsed.keywords:
+            if keyword.lower() in chunk.content.lower():
+                return 0
+        return 2