ragjson 0.1.0__py3-none-any.whl

ragjson/commands.py

@@ -0,0 +1,478 @@
+"""CLI command implementations.
+
+Each function corresponds to a `rag` subcommand defined in spec Section 6.
+"""
+
+import os
+import sys
+import tarfile
+from concurrent.futures import as_completed
+from pathlib import Path
+
+from ragjson import RAGJSON_VERSION
+from ragjson.atomic_ops import (
+    DirectoryLock,
+    clean_tmp_files,
+    ensure_rag_dirs,
+)
+from ragjson.config import RagConfig, load_config, require_embedding_config
+from ragjson.doc_parser import is_supported_file, parse_document
+from ragjson.embedding import EmbeddingClient
+from ragjson.ragjson_file import (
+    STATUS_DELETED,
+    STATUS_EMBEDDED,
+    STATUS_EMBEDDING,
+    STATUS_INITIAL,
+    STATUS_SLICED,
+    STATUS_SLICING,
+    SliceData,
+    create_initial_ragjson,
+    create_ragjson_filename,
+    is_source_changed,
+    populate_source_info,
+    ragjson_path_for,
+    read_ragjson,
+    update_status,
+    write_ragjson,
+)
+from ragjson.splitter import split_text
+from ragjson.utils import bold, green, red, yellow
+
+
+# ============================================================
+# rag init <path>
+# ============================================================
+
+def cmd_init(path: str) -> None:
+    """Initialize .rag/ and .rag/tmp/ directories at the given path.
+
+    Per spec Section 6.1: rag init <path>
+    Creates the target directory if it does not exist.
+    """
+    target = Path(path).resolve()
+    if target.exists() and not target.is_dir():
+        print(f"Error: Path is not a directory: {target}", file=sys.stderr)
+        sys.exit(1)
+
+    # Create target directory if needed
+    target.mkdir(parents=True, exist_ok=True)
+
+    rag_dir, tmp_dir = ensure_rag_dirs(target)
+    print(f"Initialized RAG directory: {target}")
+    print(f"  Created: {rag_dir}")
+    print(f"  Created: {tmp_dir}")
+
+
+# ============================================================
+# rag status [path] [-r]
+# ============================================================
+
+def cmd_status(path: str, recursive: bool = False) -> None:
+    """Show knowledge base status with MD5 comparison.
+
+    Per spec Section 6.2:
+    - Green: EMBEDDED and MD5 match
+    - Red: MD5 mismatch or no ragjson (needs update)
+    - Yellow: intermediate states
+    """
+    target = Path(path).resolve()
+    if not target.is_dir():
+        print(f"Error: Path is not a directory: {target}", file=sys.stderr)
+        sys.exit(1)
+
+    dirs_to_scan = _collect_dirs(target, recursive)
+    has_any = False
+
+    for scan_dir in dirs_to_scan:
+        rag_dir = scan_dir / ".rag"
+        files = _list_supported_files(scan_dir)
+        if not files:
+            continue
+
+        has_any = True
+        print(f"\n{bold(str(scan_dir))}")
+
+        for f in files:
+            ragjson_path = ragjson_path_for(scan_dir, f.name)
+            data = read_ragjson(ragjson_path)
+
+            if data is None:
+                # No ragjson file — needs processing
+                print(f"  {red('NEW'):30s} {f.name}")
+            elif data.processing.status == STATUS_EMBEDDED:
+                if is_source_changed(data, f):
+                    # MD5 changed — needs reprocessing
+                    print(f"  {red('CHANGED'):30s} {f.name}")
+                else:
+                    # Up to date
+                    print(f"  {green('EMBEDDED'):30s} {f.name}")
+            elif data.processing.status == STATUS_DELETED:
+                print(f"  {yellow('DELETED'):30s} {f.name}")
+            else:
+                # Intermediate state
+                status = data.processing.status
+                print(f"  {yellow(status):30s} {f.name}")
+
+        # Check for orphaned ragjson files (source deleted)
+        if rag_dir.exists():
+            for rj in rag_dir.glob("*.ragjson"):
+                source_name = rj.name.removesuffix(".ragjson")
+                source_path = scan_dir / source_name
+                if not source_path.exists():
+                    data = read_ragjson(rj)
+                    if data and data.processing.status != STATUS_DELETED:
+                        print(f"  {yellow('ORPHANED'):30s} {source_name}")
+
+    if not has_any:
+        print("No supported files found.")
+
+
+# ============================================================
+# rag update [path] [-r] [-j N] [--splitter TYPE] [--parser TYPE]
+# ============================================================
+
+def cmd_update(
+    path: str,
+    recursive: bool = False,
+    jobs: int = 0,
+    splitter_type: str = "recursive",
+    parser_type: str = "markitdown",
+) -> None:
+    """Incremental update: process files with changed MD5.
+
+    Per spec Section 6.1: rag update [path] [-r]
+    Supports -j for parallel directory processing.
+    """
+    target = Path(path).resolve()
+    if not target.is_dir():
+        print(f"Error: Path is not a directory: {target}", file=sys.stderr)
+        sys.exit(1)
+
+    config = load_config()
+    require_embedding_config(config)
+
+    dirs_to_scan = _collect_dirs(target, recursive)
+
+    # Collect directories that have files needing updates
+    update_tasks = []
+    for scan_dir in dirs_to_scan:
+        files_needing_update = _files_needing_update(scan_dir)
+        if files_needing_update:
+            update_tasks.append((scan_dir, files_needing_update))
+
+    if not update_tasks:
+        print("All files are up to date.")
+        return
+
+    total_files = sum(len(files) for _, files in update_tasks)
+    print(f"Found {total_files} file(s) to update in {len(update_tasks)} directory(ies)")
+
+    if jobs == 0:
+        jobs = os.cpu_count() or 1
+
+    if len(update_tasks) == 1:
+        # Single directory — no need for parallelism
+        scan_dir, files = update_tasks[0]
+        _process_directory(scan_dir, files, config, splitter_type, parser_type)
+    else:
+        # Multiple directories — process in parallel
+        # Note: We use ThreadPoolExecutor instead of ProcessPoolExecutor
+        # because the heavy work (API calls) is I/O-bound, and threads
+        # share the config without pickling issues.
+        from concurrent.futures import ThreadPoolExecutor
+
+        with ThreadPoolExecutor(max_workers=min(jobs, len(update_tasks))) as pool:
+            futures = {}
+            for scan_dir, files in update_tasks:
+                future = pool.submit(
+                    _process_directory,
+                    scan_dir, files, config, splitter_type, parser_type,
+                )
+                futures[future] = scan_dir
+
+            for future in as_completed(futures):
+                scan_dir = futures[future]
+                try:
+                    future.result()
+                except Exception as e:
+                    print(f"Error processing {scan_dir}: {e}", file=sys.stderr)
+
+
+def _process_directory(
+    scan_dir: Path,
+    files: list[Path],
+    config: RagConfig,
+    splitter_type: str,
+    parser_type: str,
+) -> None:
+    """Process all files in a single directory with locking.
+
+    Implements the full update pipeline per spec Section 5.
+    """
+    rag_dir = scan_dir / ".rag"
+    ensure_rag_dirs(scan_dir)
+
+    # Acquire directory lock (spec Section 3.3)
+    with DirectoryLock(rag_dir):
+        # Clean up any residual .tmp files (spec Section 3.2)
+        cleaned = clean_tmp_files(rag_dir)
+        if cleaned:
+            print(f"  Cleaned {len(cleaned)} residual temp file(s) in {rag_dir}")
+
+        # Initialize embedding client
+        embed_client = EmbeddingClient(config)
+        model_name, dimension = embed_client.get_model_info()
+
+        for source_path in files:
+            filename = source_path.name
+            ragjson_name = create_ragjson_filename(filename)
+            ragjson_path = ragjson_path_for(scan_dir, filename)
+
+            print(f"  Processing: {filename}")
+
+            # Read existing or create new ragjson
+            data = read_ragjson(ragjson_path)
+            if data is None:
+                data = create_initial_ragjson(source_path)
+                write_ragjson(data, scan_dir, ragjson_name)
+            else:
+                # Reset to INITIAL if source changed
+                if data.processing.status not in (STATUS_INITIAL,):
+                    update_status(data, STATUS_INITIAL)
+                    populate_source_info(data, source_path)
+                    write_ragjson(data, scan_dir, ragjson_name)
+
+            try:
+                # --- SLICING phase ---
+                update_status(data, STATUS_SLICING)
+                write_ragjson(data, scan_dir, ragjson_name)
+
+                # Parse document to markdown
+                print(f"    Parsing with {parser_type}...")
+                md_text = parse_document(source_path, parser_type, config)
+
+                if not md_text.strip():
+                    print(f"    Warning: Empty content from {filename}, skipping")
+                    update_status(data, STATUS_INITIAL)
+                    write_ragjson(data, scan_dir, ragjson_name)
+                    continue
+
+                # Split text into chunks
+                chunks = split_text(
+                    md_text,
+                    splitter_type=splitter_type,
+                    chunk_size=config.splitter.chunk_size,
+                    chunk_overlap=config.splitter.chunk_overlap,
+                )
+
+                if not chunks:
+                    print(f"    Warning: No chunks produced for {filename}, skipping")
+                    update_status(data, STATUS_INITIAL)
+                    write_ragjson(data, scan_dir, ragjson_name)
+                    continue
+
+                # Store slices (without vectors yet)
+                data.slices = [
+                    SliceData(index=i, content=chunk, vector_base64="")
+                    for i, chunk in enumerate(chunks)
+                ]
+                update_status(data, STATUS_SLICED)
+                write_ragjson(data, scan_dir, ragjson_name)
+
+                # --- EMBEDDING phase ---
+                update_status(data, STATUS_EMBEDDING)
+                write_ragjson(data, scan_dir, ragjson_name)
+
+                # Generate embeddings
+                print(f"    Embedding {len(chunks)} chunk(s)...")
+                vector_b64_list = embed_client.embed_and_encode(chunks)
+
+                # Attach vectors to slices
+                for i, vec_b64 in enumerate(vector_b64_list):
+                    data.slices[i].vector_base64 = vec_b64
+
+                # Update embedding info
+                data.embedding.model_name = model_name
+                data.embedding.dimension = dimension
+
+                # Final write: EMBEDDED
+                update_status(data, STATUS_EMBEDDED)
+                write_ragjson(data, scan_dir, ragjson_name)
+                print(f"    {green('Done')}: {len(chunks)} slices, {dimension}-dim vectors")
+
+            except Exception as e:
+                print(f"    {red('Error')}: {filename}: {e}", file=sys.stderr)
+                # Reset to INITIAL so it can be retried
+                try:
+                    update_status(data, STATUS_INITIAL)
+                    write_ragjson(data, scan_dir, ragjson_name)
+                except Exception:
+                    pass
+
+
+# ============================================================
+# rag clean [path] [-r]
+# ============================================================
+
+def cmd_clean(path: str, recursive: bool = False) -> None:
+    """Clean residual .tmp files from .rag/tmp/ directories.
+
+    Per spec Section 6.1: rag clean [path] [-r]
+    """
+    target = Path(path).resolve()
+    if not target.is_dir():
+        print(f"Error: Path is not a directory: {target}", file=sys.stderr)
+        sys.exit(1)
+
+    dirs_to_scan = _collect_dirs(target, recursive)
+    total_cleaned = 0
+
+    for scan_dir in dirs_to_scan:
+        rag_dir = scan_dir / ".rag"
+        if not rag_dir.exists():
+            continue
+
+        cleaned = clean_tmp_files(rag_dir)
+        if cleaned:
+            print(f"  Cleaned {len(cleaned)} temp file(s) in {rag_dir / 'tmp'}")
+            total_cleaned += len(cleaned)
+
+    if total_cleaned == 0:
+        print("No residual temp files found.")
+    else:
+        print(f"Total: cleaned {total_cleaned} temp file(s).")
+
+
+# ============================================================
+# rag cleanup [path] [-r]
+# ============================================================
+
+def cmd_cleanup(path: str, recursive: bool = False) -> None:
+    """Physically delete .ragjson files with DELETED status.
+
+    Per spec Section 6.1: rag cleanup [path] [-r]
+    """
+    target = Path(path).resolve()
+    if not target.is_dir():
+        print(f"Error: Path is not a directory: {target}", file=sys.stderr)
+        sys.exit(1)
+
+    dirs_to_scan = _collect_dirs(target, recursive)
+    total_removed = 0
+
+    for scan_dir in dirs_to_scan:
+        rag_dir = scan_dir / ".rag"
+        if not rag_dir.exists():
+            continue
+
+        for rj_path in rag_dir.glob("*.ragjson"):
+            data = read_ragjson(rj_path)
+            if data and data.processing.status == STATUS_DELETED:
+                rj_path.unlink()
+                print(f"  Removed: {rj_path}")
+                total_removed += 1
+
+        # Also check for orphaned ragjson (source file missing, not DELETED)
+        for rj_path in rag_dir.glob("*.ragjson"):
+            source_name = rj_path.name.removesuffix(".ragjson")
+            source_path = scan_dir / source_name
+            if not source_path.exists():
+                data = read_ragjson(rj_path)
+                if data and data.processing.status != STATUS_DELETED:
+                    # Mark as DELETED first, then remove
+                    update_status(data, STATUS_DELETED)
+                    write_ragjson(data, scan_dir, rj_path.name)
+                    rj_path.unlink()
+                    print(f"  Removed orphaned: {rj_path}")
+                    total_removed += 1
+
+    if total_removed == 0:
+        print("No DELETED or orphaned ragjson files found.")
+    else:
+        print(f"Total: removed {total_removed} ragjson file(s).")
+
+
+# ============================================================
+# rag export <path> <output>
+# ============================================================
+
+def cmd_export(path: str, output: str) -> None:
+    """Export all .ragjson files into a tar.gz archive.
+
+    Per spec Section 6.1: rag export <path> <output>
+    """
+    target = Path(path).resolve()
+    if not target.is_dir():
+        print(f"Error: Path is not a directory: {target}", file=sys.stderr)
+        sys.exit(1)
+
+    output_path = Path(output).resolve()
+
+    # Collect all .ragjson files recursively
+    ragjson_files = list(target.rglob(".rag/*.ragjson"))
+
+    if not ragjson_files:
+        print("No .ragjson files found to export.")
+        return
+
+    # Create tar.gz archive preserving relative directory structure
+    with tarfile.open(output_path, "w:gz") as tar:
+        for rj_file in ragjson_files:
+            # Store path relative to the target directory
+            arcname = rj_file.relative_to(target)
+            tar.add(rj_file, arcname=arcname)
+
+    print(f"Exported {len(ragjson_files)} ragjson file(s) to: {output_path}")
+
+
+# ============================================================
+# Helper functions
+# ============================================================
+
+def _collect_dirs(base: Path, recursive: bool) -> list[Path]:
+    """Collect directories to scan, optionally recursive."""
+    dirs = [base]
+    if recursive:
+        for child in sorted(base.rglob("*")):
+            if child.is_dir() and child.name != ".rag":
+                dirs.append(child)
+    return dirs
+
+
+def _list_supported_files(scan_dir: Path) -> list[Path]:
+    """List supported files directly in a directory (not recursive)."""
+    files = []
+    for child in sorted(scan_dir.iterdir()):
+        if child.is_file() and is_supported_file(child):
+            files.append(child)
+    return files
+
+
+def _files_needing_update(scan_dir: Path) -> list[Path]:
+    """Find files in a directory that need processing.
+
+    A file needs update if:
+    - No corresponding .ragjson exists
+    - The .ragjson status is INITIAL (or stuck intermediate state)
+    - The source MD5 has changed
+    """
+    rag_dir = scan_dir / ".rag"
+    result = []
+
+    for f in _list_supported_files(scan_dir):
+        ragjson_path = ragjson_path_for(scan_dir, f.name)
+        data = read_ragjson(ragjson_path)
+
+        if data is None:
+            result.append(f)
+        elif data.processing.status == STATUS_EMBEDDED:
+            if is_source_changed(data, f):
+                result.append(f)
+        elif data.processing.status in (STATUS_SLICING, STATUS_EMBEDDING):
+            # Stuck intermediate state — needs recovery
+            result.append(f)
+        elif data.processing.status in (STATUS_INITIAL, STATUS_SLICED):
+            result.append(f)
+        # DELETED status: skip (will be handled by cleanup)
+
+    return result

ragjson/config.py

@@ -0,0 +1,113 @@
+"""Configuration loader for ~/.rag.config (JSON format)."""
+
+import json
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class EmbeddingModelConfig:
+    """Embedding model service configuration."""
+    base_url: str = ""
+    api_key: str = ""
+    model_name: str = "text-embedding-3-small"
+    dimension: int = 1536
+
+
+@dataclass
+class MinerUConfig:
+    """MinerU document parser service configuration."""
+    backend: str = "http-client"
+    server_url: str = ""
+    image_dir: str = "/tmp/mineru_output"
+    image_url_prefix: str = ""
+
+
+@dataclass
+class SplitterConfig:
+    """Text splitter configuration."""
+    chunk_size: int = 1000
+    chunk_overlap: int = 200
+
+
+@dataclass
+class RagConfig:
+    """Top-level configuration combining all sections."""
+    embedding_model: EmbeddingModelConfig = field(default_factory=EmbeddingModelConfig)
+    mineru: MinerUConfig = field(default_factory=MinerUConfig)
+    splitter: SplitterConfig = field(default_factory=SplitterConfig)
+
+
+def get_config_path() -> Path:
+    """Return the path to the configuration file."""
+    return Path.home() / ".rag.config"
+
+
+def load_config() -> RagConfig:
+    """Load and parse the configuration from ~/.rag.config.
+
+    Returns:
+        RagConfig instance with parsed values.
+
+    Raises:
+        SystemExit: If config file is missing or malformed.
+    """
+    config_path = get_config_path()
+
+    if not config_path.exists():
+        print(
+            f"Error: Configuration file not found: {config_path}\n"
+            f"Please create {config_path} with the required settings.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    try:
+        raw = json.loads(config_path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as e:
+        print(f"Error: Invalid JSON in {config_path}: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    config = RagConfig()
+
+    # Parse embedding_model section
+    if "embedding_model" in raw:
+        emb = raw["embedding_model"]
+        config.embedding_model = EmbeddingModelConfig(
+            base_url=emb.get("base_url", ""),
+            api_key=emb.get("api_key", ""),
+            model_name=emb.get("model_name", "text-embedding-3-small"),
+            dimension=emb.get("dimension", 1536),
+        )
+
+    # Parse MinerU section
+    if "mineru" in raw:
+        mineru = raw["mineru"]
+        config.mineru = MinerUConfig(
+            backend=mineru.get("backend", "http-client"),
+            server_url=mineru.get("server_url", ""),
+            image_dir=mineru.get("image_dir", "/tmp/mineru_output"),
+            image_url_prefix=mineru.get("image_url_prefix", ""),
+        )
+
+    # Parse splitter section
+    if "splitter" in raw:
+        splitter = raw["splitter"]
+        config.splitter = SplitterConfig(
+            chunk_size=splitter.get("chunk_size", 1000),
+            chunk_overlap=splitter.get("chunk_overlap", 200),
+        )
+
+    return config
+
+
+def require_embedding_config(config: RagConfig) -> None:
+    """Validate that embedding model configuration is present."""
+    if not config.embedding_model.base_url or not config.embedding_model.api_key:
+        print(
+            "Error: Embedding model configuration missing.\n"
+            f"Please set 'embedding_model.base_url' and 'embedding_model.api_key' in {get_config_path()}",
+            file=sys.stderr,
+        )
+        sys.exit(1)