diff --git a/cecli/commands/__init__.py b/cecli/commands/__init__.py index ebc7aa11823..5da19faeb05 100644 --- a/cecli/commands/__init__.py +++ b/cecli/commands/__init__.py @@ -54,7 +54,7 @@ from .read_only import ReadOnlyCommand from .read_only_stub import ReadOnlyStubCommand from .reap_agent import ReapAgentCommand -from .reasoning_effort import ReasoningEffortCommand +from .rag import RagAddCommand, RagDropCommand, RagLsCommand, RagQuestionCommand from .remove_memory import RemoveMemoryCommand from .report import ReportCommand from .reset import ResetCommand @@ -138,7 +138,10 @@ CommandRegistry.register(ReadOnlyCommand) CommandRegistry.register(ReadOnlyStubCommand) CommandRegistry.register(ReapAgentCommand) -CommandRegistry.register(ReasoningEffortCommand) +CommandRegistry.register(RagAddCommand) +CommandRegistry.register(RagDropCommand) +CommandRegistry.register(RagLsCommand) +CommandRegistry.register(RagQuestionCommand) CommandRegistry.register(RemoveMemoryCommand) CommandRegistry.register(ReportCommand) CommandRegistry.register(ResetCommand) @@ -221,6 +224,10 @@ "ReadOnlyCommand", "ReadOnlyStubCommand", "ReapAgentCommand", +"RagAddCommand", +"RagDropCommand", + "RagLsCommand", + "RagQuestionCommand", "ReasoningEffortCommand", "ReloadProgramSignal", "RemoveMemoryCommand", diff --git a/cecli/commands/rag.py b/cecli/commands/rag.py new file mode 100644 index 00000000000..f5a3626b89b --- /dev/null +++ b/cecli/commands/rag.py @@ -0,0 +1,373 @@ +"""The /rag-add and /rag-question commands. + +``/rag-add`` mirrors the UX of ``/add``: it picks files/folders (via the fuzzy +finder when no args are given, or from arguments accepting absolute/relative +paths and directory/glob expansions) and registers them as RAG objects. +Unlike the old persistent design, these RAG objects live **only for the +current conversation** — they are stored on the coder (``coder.rag_files`` ++ ``coder.rag_index``) and disappear when the cecli session ends. + +``/rag-question`` retrieves the most relevant chunks from those session RAG +objects and answers the question the way ``/help`` answers about cecli docs: +it builds a temporary coder, injects the retrieved RAG context block into the +user message, runs it, then switches back to the original coder. +""" + +import os +from pathlib import Path +from typing import List + +from cecli.commands.utils.base_command import BaseCommand +from cecli.commands.utils.helpers import ( + format_command_result, + get_file_completions, + parse_quoted_filenames, +) +from cecli.utils import run_fzf + + +def _get_rag_state(coder): + """Return ``(files, index)`` for the coder's session-scoped RAG store. + + The state lives on the coder instance so it disappears when the session + ends, mirroring how ``/add`` keeps files in ``coder.abs_fnames``. + """ + files = getattr(coder, "rag_files", None) + if files is None: + files = set() + coder.rag_files = files + index = getattr(coder, "rag_index", None) + if index is None: + index = {"entries": [], "files": []} + coder.rag_index = index + return files, index + + +class RagAddCommand(BaseCommand): + NORM_NAME = "rag-add" + DESCRIPTION = "Add files/folders as RAG objects for the current conversation only" + show_completion_notification = True + + @classmethod + async def execute(cls, io, coder, args, **kwargs): + """Register the selected files/folders as session-scoped RAG objects.""" + from cecli.rag import build_rag_index, collect_target_files + + root = coder.root + files, index = _get_rag_state(coder) + + selected = await cls._select_targets(io, coder, args) + if not selected: + return format_command_result(io, "rag-add", "No files selected for RAG") + + # Resolve paths into the project root. + resolved = [cls._resolve(root, p) for p in selected] + resolved = [p for p in resolved if p is not None] + if not resolved: + return format_command_result( + io, "rag-add", "None of the selected paths exist; nothing to add" + ) + + # Expand directories/globs into concrete files. + concrete_files = collect_target_files(resolved, root) + if not concrete_files: + return format_command_result( + io, "rag-add", "No indexable files found in selection" + ) + + for f in concrete_files: + files.add(str(f)) + + # Rebuild the in-memory keyword index over the newly added files, then + # merge it into the session index so previously added files remain. + new_index_data = build_rag_index(concrete_files) + seen_ids = {e["id"] for e in index["entries"]} + for entry in new_index_data["entries"]: + if entry["id"] not in seen_ids: + index["entries"].append(entry) + seen_ids.add(entry["id"]) + for fpath in new_index_data["files"]: + if fpath not in index["files"]: + index["files"].append(fpath) + + io.tool_output( + f"Added {len(concrete_files)} files ({len(new_index_data['entries'])} chunks) to the RAG store." + ) + io.tool_output( + "This RAG set is valid for the current conversation only. Use /rag-question to ask questions." + ) + return format_command_result( + io, "rag-add", f"RAG objects added: {len(concrete_files)} files" + ) + + @classmethod + async def _select_targets(cls, io, coder, args): + """No args -> fuzzy finder; otherwise -> parse file/folder tokens.""" + if not args.strip(): + all_files = coder.get_all_relative_files() + in_chat = coder.get_inchat_relative_files() + candidates = sorted(set(all_files) - set(in_chat)) + if not candidates: + io.tool_output("No files available to add to the RAG set.") + return [] + selected = run_fzf(candidates, multi=True, coder=coder) + return selected or [] + + # Arguments may include sub-directories and globs; parse quoted names. + return parse_quoted_filenames(args) + + @staticmethod + def _resolve(root, pattern): + pattern = os.path.expanduser(pattern) + p = Path(pattern) + if not p.is_absolute(): + p = Path(root) / p + # Keep it even if the literal path does not exist (it may be a glob). + return p + + @classmethod + def get_completions(cls, io, coder, args) -> List[str]: + """Complete both directory-based and full-file paths (like /add).""" + directory_completions = get_file_completions( + coder, + args=args, + completion_type="directory", + include_directories=True, + filter_in_chat=False, + ) + all_completions = get_file_completions( + coder, + args=args, + completion_type="all", + include_directories=True, + filter_in_chat=False, + ) + return sorted(set(directory_completions) | set(all_completions)) + + @classmethod + def get_help(cls) -> str: + help_text = super().get_help() + help_text += "\nUsage:\n" + help_text += " /rag-add # Open the fuzzy finder to pick files/folders\n" + help_text += " /rag-add [...] # Register absolute/relative paths or folders\n" + help_text += "\nRAG objects are stored for the current conversation only (like /add files).\n" + help_text += "They disappear when the cecli session ends.\n" + help_text += "Use /rag-question to ask questions against the registered RAG objects.\n" + return help_text + + +class RagQuestionCommand(BaseCommand): + NORM_NAME = "rag-question" + DESCRIPTION = "Ask a question about the RAG objects added via /rag-add" + show_completion_notification = False + + @classmethod + async def execute(cls, io, coder, args, **kwargs): + """Retrieve from the session RAG store and answer, like /help.""" + from cecli.rag import rag_context_block + + question = (args or "").strip() + if not question: + io.tool_output("Usage: /rag-question ") + return format_command_result(io, "rag-question", "Missing question", error=None) + + index = getattr(coder, "rag_index", None) + if not index or not index.get("entries"): + io.tool_output( + "No RAG objects in this conversation yet. " + "Run /rag-add first to register files/folders." + ) + return format_command_result( + io, "rag-question", "No RAG objects available", error=None + ) + + import traceback + from uuid import uuid4 as generate_unique_id + + from cecli.coders.base_coder import Coder + from cecli.commands import SwitchCoderSignal + + has_args = bool(getattr(coder, "args", None)) + + try: + # Use the editor_model from the main_model if it exists, otherwise use the main_model itself + editor_model = coder.main_model.editor_model or coder.main_model + + original_coder = coder + + create_kwargs = dict() + create_kwargs["io"] = io + create_kwargs["uuid"] = str(generate_unique_id()) + create_kwargs["from_coder"] = coder + create_kwargs["edit_format"] = "help" + create_kwargs["summarize_from_coder"] = False + create_kwargs["map_tokens"] = 0 + create_kwargs["map_mul_no_files"] = 1 + create_kwargs["main_model"] = editor_model + create_kwargs["args"] = coder.args + create_kwargs["suggest_shell_commands"] = False + create_kwargs["cache_prompts"] = False + create_kwargs["num_cache_warming_pings"] = 0 + + rag_coder = await Coder.create(**create_kwargs) + user_msg = rag_context_block(index, question) + if not user_msg.strip(): + user_msg = f"# Question: {question}\n\n# No RAG context was retrieved.\n" + user_msg += f"\n\n# Question: {question}\n" + user_msg += ( + "\nAnswer the above question using only the RAG context provided above.\n" + "If the context is insufficient, say so and what additional information you need.\n" + ) + user_msg += "\n# Announcement lines from when this session of cecli was launched:\n\n" + user_msg += "\n".join(coder.get_announcements() if has_args else []) + "\n" + + await rag_coder.run(user_msg, preproc=False) + except Exception as err: + io.tool_error(f"RAG question failed to initialize: {err}") + io.tool_error(traceback.format_exc()) + return format_command_result( + io, "rag-question", f"RAG question failed to initialize: {err}" + ) + + if coder.repo_map: + map_tokens = coder.repo_map.max_map_tokens + map_mul_no_files = coder.repo_map.map_mul_no_files + else: + map_tokens = 0 + map_mul_no_files = 1 + + raise SwitchCoderSignal( + edit_format=coder.edit_format, + summarize_from_coder=False, + from_coder=original_coder, + map_tokens=map_tokens, + map_mul_no_files=map_mul_no_files, + show_announcements=False, + ) + + @classmethod + def get_completions(cls, io, coder, args) -> List[str]: + """No completions for the free-form question text.""" + return [] + + @classmethod + def get_help(cls) -> str: + help_text = super().get_help() + help_text += "\nUsage:\n" + help_text += " /rag-question \n" + help_text += "\nRetrieves the most relevant chunks from the RAG objects added via /rag-add\n" + help_text += "in this conversation, and answers the question the way /help answers\n" + help_text += "questions about cecli docs.\n" + return help_text + + +class RagDropCommand(BaseCommand): + NORM_NAME = "rag-drop" + DESCRIPTION = "Remove RAG objects added via /rag-add from the current conversation" + show_completion_notification = False + + @classmethod + async def execute(cls, io, coder, args, **kwargs): + """Remove the specified (or all) session-scoped RAG objects.""" + files, index = _get_rag_state(coder) + + if not files and not index.get("entries"): + io.tool_output("No RAG objects in this conversation. Nothing to drop.") + return format_command_result(io, "rag-drop", "No RAG objects to drop", error=None) + + if not args.strip(): + removed = len(files) + files.clear() + index["entries"] = [] + index["files"] = [] + io.tool_output(f"Dropped all {removed} RAG files from the current conversation.") + return format_command_result(io, "rag-drop", f"Dropped all RAG objects ({removed} files)") + + import fnmatch + + filenames = parse_quoted_filenames(args) + dropped = [] + for word in filenames: + expanded_word = os.path.expanduser(word) + is_glob = any(c in expanded_word for c in "*?[") + for fpath in list(files): + if is_glob: + if fnmatch.fnmatch(fpath, expanded_word) or fnmatch.fnmatch(os.path.basename(fpath), expanded_word): + files.remove(fpath) + dropped.append(fpath) + else: + abs_path = cls._resolve(coder.root, expanded_word) + if str(abs_path) == fpath or expanded_word in fpath: + files.remove(fpath) + dropped.append(fpath) + + if not dropped: + io.tool_output(f"No RAG objects matched: {args.strip()}") + return format_command_result(io, "rag-drop", "No matching RAG objects", error=None) + + from cecli.rag import build_rag_index, collect_target_files + index_data = build_rag_index(collect_target_files(list(files), coder.root)) + index["entries"] = index_data["entries"] + index["files"] = index_data["files"] + + io.tool_output(f"Dropped {len(dropped)} RAG files: " + ", ".join(sorted(dropped))) + return format_command_result(io, "rag-drop", f"Dropped {len(dropped)} RAG objects") + + @classmethod + def get_completions(cls, io, coder, args) -> List[str]: + files = getattr(coder, "rag_files", None) + if not files: + return [] + return sorted(cls._quote_fname(os.path.basename(f)) for f in files) + + @staticmethod + def _quote_fname(fname): + if " " in fname and '"' not in fname: + return f'"{fname}"' + return fname + + @classmethod + def get_help(cls) -> str: + help_text = super().get_help() + help_text += "\nUsage:\n" + help_text += " /rag-drop # Drop all RAG objects from this conversation\n" + help_text += " /rag-drop [...] # Drop specific RAG objects\n" + help_text += "\nExamples:\n" + help_text += " /rag-drop # Remove all RAG objects\n" + help_text += " /rag-drop main.py # Remove main.py from the RAG set\n" + help_text += " /rag-drop *.md # Remove all markdown RAG objects\n" + return help_text + +class RagLsCommand(BaseCommand): + NORM_NAME = "rag-ls" + DESCRIPTION = "List all RAG objects registered in the current conversation" + show_completion_notification = False + + @classmethod + async def execute(cls, io, coder, args, **kwargs): + """List the files registered in the session RAG store.""" + files = getattr(coder, "rag_files", None) + index = getattr(coder, "rag_index", None) + + if not files and not index: + io.tool_output("\nNo RAG objects in the current conversation. Run /rag-add first.") + return format_command_result(io, "rag-ls", "No RAG objects listed") + + chunk_count = len(index.get("entries", [])) if index else 0 + io.tool_output(f"RAG files in the current conversation ({len(files)} files, {chunk_count} chunks):") + for f in sorted(files): + io.tool_output(f" {f}") + + return format_command_result(io, "rag-ls", f"Listed {len(files)} RAG files") + + @classmethod + def get_completions(cls, io, coder, args) -> List[str]: + """No completions for the list command.""" + return [] + + @classmethod + def get_help(cls) -> str: + help_text = super().get_help() + help_text += "\nUsage:\n" + help_text += " /rag-ls # List all RAG objects registered in this conversation\n" + return help_text \ No newline at end of file diff --git a/cecli/rag.py b/cecli/rag.py new file mode 100644 index 00000000000..d1550a902d6 --- /dev/null +++ b/cecli/rag.py @@ -0,0 +1,252 @@ +"""Session-scoped, in-memory RAG store over arbitrary files. + +This mirrors the UX of ``/help`` RAG (chunk-aligned splitting + a +dependency-free keyword retrieval path) but generalizes it to user-specified +files and folders that are registered for the *current conversation only*. +Unlike the old persistent chromadb design, nothing is written to disk: the +chunks live on the coder (``coder.rag_files`` + ``coder.rag_index``) and +disappear when the cecli session ends, exactly like files added with ``/add``. + +The commands in :mod:`cecli.commands.rag` register targets with ``/rag-add`` +and ask questions with ``/rag-question``. +""" + +import os +import re +from collections import Counter +from pathlib import Path + +from cecli.help import _chunk_document + +# Chunk size tuned for in-memory retrieval; a little larger than help chunks +# since user files are arbitrary. +_RAG_CHUNK_CHARS = 1500 +_MAX_RESULTS = 8 +_MAX_CONTEXT_CHARS = 16000 + +# Extensions we skip: binary/asset/generated content that has no textual value. +_SKIP_SUFFIXES = { + ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".webp", ".ico", + ".pdf", ".gz", ".zip", ".tar", ".bz2", ".7z", ".rar", + ".pyc", ".pyo", ".class", ".o", ".obj", ".bin", + ".woff", ".woff2", ".ttf", ".otf", ".eot", + ".mp3", ".mp4", ".wav", ".avi", ".mov", ".mkv", + ".lock", +} + +_TOKEN_RE = re.compile(r"[a-z0-9]+") +_KEYWORD_STOPWORDS = frozenset( + "a an and are as at be been by can could did do does for from had has have how " + "i if in into is it its me my no not of on or our should so than that the then " + "this to us was we were what when where which who why will with would you your".split() +) + + +def _is_skipped(fname: Path) -> bool: + if fname.name.startswith(".git"): + return True + parts = fname.parts + if any(p in {".git", "node_modules", ".venv", "venv", "__pycache__"} for p in parts): + return True + return fname.suffix.lower() in _SKIP_SUFFIXES + + +def collect_target_files(paths, root: str) -> list: + """Expand files/directories/globs into a de-duplicated list of target files. + + Mirrors ``/add`` semantics: a file is a file; a directory expands to all + contained files; a glob pattern expands against the project root. All + results are returned as absolute :class:`Path` objects that exist and are + not skipped. + """ + import glob as _glob + + root_path = Path(root) + seen = set() + result = [] + for raw in paths: + raw = os.path.expanduser(raw) + p = Path(raw) + if not p.is_absolute(): + p = root_path / p + if p.exists(): + candidates = [p] + else: + # Treat as glob relative to root + candidates = [Path(m) for m in _glob.glob(str(p))] + + for cand in candidates: + if cand.is_file(): + files = [cand] + elif cand.is_dir(): + files = [f for f in cand.rglob("*") if f.is_file()] + else: + continue + for f in files: + try: + f = f.resolve() + except OSError: + f = f + if _is_skipped(f): + continue + key = str(f) + if key not in seen: + seen.add(key) + result.append(f) + return sorted(set(result)) + + +def _tokenize(text): + """Split text into search tokens, dropping stopwords. + + Handles both Latin/CJK text: Latin words are tokenized via ``[a-z0-9]+``; + CJK characters are tokenized as individual characters and 2-char bigrams so + that Chinese/Japanese/Korean questions can match Chinese/Japanese/Korean + document content without a full NLP dependency. + """ + text_lower = (text or "").lower() + tokens = [] + # Latin / alphanumeric words + for word in _TOKEN_RE.findall(text_lower): + if len(word) > 1 and word not in _KEYWORD_STOPWORDS: + tokens.append(word) + # CJK unigrams + bigrams (covers Chinese, Japanese, Korean) + cjk = [c for c in text_lower if "\u4e00" <= c <= "\u9fff" or "\u3400" <= c <= "\u4dbf" or "\u3040" <= c <= "\u309f" or "\u30a0" <= c <= "\u30ff" or "\uac00" <= c <= "\ud7af"] + for ch in cjk: + if ch not in _KEYWORD_STOPWORDS: + tokens.append(ch) + for i in range(len(cjk) - 1): + bigram = cjk[i] + cjk[i + 1] + if bigram not in _KEYWORD_STOPWORDS: + tokens.append(bigram) + # De-duplicate while preserving order + return list(dict.fromkeys(tokens)) + + +def build_rag_index(files) -> dict: + """Chunk each file in-memory and build a dependency-free keyword index. + + Returns a dict holding, for each file, its list of chunk entries. Each + entry is ``{"id", "path", "document", "document_lower", "heading_tokens", + "body_tokens"}`` mirroring ``cecli.help._keyword_index`` so callers can score + relevance without any external embedding backend. + """ + import os as _os + + entries = [] + for fname in files: + try: + text = Path(fname).read_text(encoding="utf-8", errors="replace") + except (OSError, UnicodeError): + continue + + filename_tokens = set(_tokenize(_os.path.basename(str(fname)))) + for chunk_index, chunk in enumerate(_chunk_document(text, max_chars=_RAG_CHUNK_CHARS)): + if not chunk.strip(): + continue + heading_tokens = set() + for line in chunk.splitlines(): + if line.lstrip().startswith("#"): + heading_tokens.update(_tokenize(line)) + entries.append( + dict( + id=f"{fname}#{chunk_index}", + path=str(fname), + document=chunk, + document_lower=chunk.lower(), + filename_tokens=filename_tokens, + heading_tokens=heading_tokens, + body_tokens=Counter(_tokenize(chunk)), + ) + ) + return {"entries": entries, "files": sorted(set(str(f) for f in files))} + + +def _keyword_score(entry, tokens, phrase): + """Score a chunk using filename, heading, and body keyword overlap.""" + score = 0.0 + matched = 0 + for token in tokens: + token_score = 0.0 + if token in entry["filename_tokens"]: + token_score += 8.0 + if token in entry["heading_tokens"]: + token_score += 5.0 + count = entry["body_tokens"].get(token, 0) + if count: + token_score += min(count, 5) + if token_score: + matched += 1 + score += token_score + if matched and len(tokens) > 1 and phrase in entry["document_lower"]: + score += 6.0 + if matched: + score += 10.0 * matched / len(tokens) + return score + + +def retrieve_from_index(index: dict, question, n_results=_MAX_RESULTS, exclude_ids=None): + """Score chunks in a :func:`build_rag_index` result against a question. + + Returns ranked chunk dicts in the same shape the old ``RagStore.retrieve`` + produced (``id``, ``path``, ``document``, ``metadata``, ``distance``) so the + command layer can render them uniformly. + """ + entries = (index or {}).get("entries") or [] + if not entries: + return [] + + query_tokens = list(dict.fromkeys(_tokenize(question or ""))) + if not query_tokens: + return [] + + phrase = " ".join(query_tokens) + exclude = set(exclude_ids or ()) + n_results = max(1, int(n_results)) + + scored = [] + for entry in entries: + if entry["id"] in exclude: + continue + score = _keyword_score(entry, query_tokens, phrase) + if score > 0: + scored.append((score, entry)) + + scored.sort(key=lambda item: (-item[0], item[1]["id"])) + + docs = [] + for score, entry in scored: + docs.append( + dict( + id=entry["id"], + path=entry["path"], + document=entry["document"], + metadata=dict(path=entry["path"]), + distance=round(1.0 / (1.0 + score), 6), + ) + ) + if len(docs) >= n_results: + break + return docs + + +def rag_context_block(index: dict, question, n_results=_MAX_RESULTS, exclude_ids=None): + """Render retrieved chunks as an RAG context block for a prompt. + + Mirrors ``Help.ask``: returns a string block prefixed with the question + and the retrieved ```` sections, capped by ``_MAX_CONTEXT_CHARS``. + """ + documents = retrieve_from_index(index, question, n_results=n_results, exclude_ids=exclude_ids) + if not documents: + return "" + context = f"# Question: {question}\n\n# RAG context (retrieved from the files added this session):\n\n" + used = 0 + for doc in documents: + text = doc["document"] + if not text: + continue + if used and used + len(text) > _MAX_CONTEXT_CHARS: + break + used += len(text) + context += f'\n{text}\n\n\n' + return context \ No newline at end of file