treesearchlib 1.1.0__py3-none-any.whl

treesearch/config.py

@@ -0,0 +1,206 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing(xuming624@qq.com)
+@description: Unified configuration for TreeSearch.
+
+Priority (high -> low):
+    1. set_config(TreeSearchConfig(...))
+    2. Environment variables
+    3. Built-in defaults
+
+Environment variables:
+    Tokenizer: TREESEARCH_CJK_TOKENIZER
+"""
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Literal, Optional
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Index schema version
+# ---------------------------------------------------------------------------
+# Bump whenever a change in tree builder, tokenizer, FTS schema, or node_id
+# algorithm would invalidate previously-built indexes. The version is folded
+# into every file's fingerprint so old index_meta entries automatically miss
+# and the file is re-indexed on next run.
+#
+# History:
+#   "1" — original (mtime_ns:size) fingerprint with sequential int node_ids.
+#   "2" — stable hash node_ids + node-level diff + atomic per-doc transaction.
+#   "3" — docx parser emits \n\n between paragraphs for correct preview rendering,
+#         and renders tables as GFM markdown tables (header + separator + rows).
+INDEX_SCHEMA_VERSION = "3"
+
+# ---------------------------------------------------------------------------
+# Environment variable names
+# ---------------------------------------------------------------------------
+_ENV_CJK_TOKENIZER = "TREESEARCH_CJK_TOKENIZER"
+_ENV_FINGERPRINT_MODE = "TREESEARCH_FINGERPRINT_MODE"
+_ENV_PRUNE = "TREESEARCH_PRUNE"
+_ENV_SHADOW_MD = "TREESEARCH_ENABLE_SHADOW_MD"
+_ENV_ALLOWED_SOURCE_TYPES = "TREESEARCH_ALLOWED_SOURCE_TYPES"
+
+
+def _env_int(cfg: "TreeSearchConfig", attr: str, min_val: int, env_name: str) -> None:
+    """Read an integer env var and set it on *cfg* if present and >= *min_val*."""
+    raw = os.getenv(env_name)
+    if raw is not None:
+        try:
+            v = int(raw)
+            if v >= min_val:
+                setattr(cfg, attr, v)
+        except ValueError:
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Configuration dataclass
+# ---------------------------------------------------------------------------
+
+@dataclass
+class TreeSearchConfig:
+    """Single configuration class for TreeSearch.
+
+    Priority: set_config() > env vars > defaults.
+    """
+    # Search
+    max_nodes_per_doc: int = 5
+    top_k_docs: int = 3
+
+    # Index
+    if_add_node_summary: bool = True
+    if_add_doc_description: bool = False
+    if_add_node_text: bool = True
+    if_add_node_id: bool = True
+    if_thinning: bool = False
+    min_thinning_chars: int = 15000  # min chars to keep a sub-tree during thinning
+    summary_chars_threshold: int = 600  # nodes shorter than this use full text as summary
+    max_concurrency: int = field(default_factory=lambda: min(os.cpu_count() or 4, 256))
+    # PST 解析的并发上限。PST 重量级：每文件启 sidecar 子进程 + 加载 GB 级文件 +
+    # 全量附件提取，多个并发易耗尽内存/CPU/IO 触发 sidecar 崩溃（exit 0xC000013A）。
+    # 1 = PST 串行索引（最稳）；调大可加速但会增加资源压力。
+    max_pst_concurrency: int = 1
+    max_dir_files: int = 10_000  # safety cap for directory walk
+
+    # Text length limits
+    max_node_chars: int = 8000  # max characters per node text when indexing into FTS5
+    max_result_chars: int = 32000  # max total characters of returned search result texts
+
+    # FTS
+    fts_db_path: str = ""  # empty = same DB as tree storage (default: index.db)
+    fts_title_weight: float = 5.0
+    fts_summary_weight: float = 2.0
+    fts_body_weight: float = 10.0
+    fts_code_weight: float = 1.0
+    fts_front_matter_weight: float = 2.0
+
+    # Tree Search
+    search_mode: str = "auto"  # "auto" | "flat" | "tree" | "auto" degrades to flat for code-only docs
+    anchor_top_k: int = 5  # max anchor nodes per document
+    max_anchor_per_doc: int = 3  # anchors to expand per document
+    max_expansions: int = 40  # max total node expansions in tree walk
+    max_hops: int = 3  # max depth offset from anchor
+    max_siblings: int = 2  # max sibling nodes to expand per step
+    min_frontier_score: float = 0.1  # stop if best frontier score below this
+    early_stop_score: float = 0.95  # stop early if a path reaches this score
+    path_top_k: int = 3  # top paths to return
+
+    # Tokenizer
+    cjk_tokenizer: str = "auto"  # "auto" | "jieba" | "bigram" | "char"
+
+    # Incremental indexing
+    fingerprint_mode: Literal["stat", "content"] = "stat"
+    # "stat":    fast `(mtime_ns:size)` fingerprint. Re-indexes after a `touch`.
+    # "content": samples first/middle/last 64KB of large files (full md5 for
+    #            files <1MB) — robust against `touch`/CI replay; ~1ms cost on small
+    #            files, dozens of ms on multi-GB files. Opt-in.
+    content_fingerprint_size_threshold: int = 1_000_000  # bytes; full md5 below this
+    content_fingerprint_sample_bytes: int = 64 * 1024     # head/mid/tail sample size
+
+    # Auto FTS5 maintenance: run `optimize` every N reindexed documents
+    # within a single build_index call. 0 disables.
+    auto_optimize_threshold: int = 1000
+
+    # Default prune policy when build_index sees a directory in `paths`.
+    # Set to False to never auto-delete orphans even on directory walks.
+    prune_orphans_on_directory: bool = True
+
+    # Failed files auto-skip: after N consecutive parse failures, skip the file
+    max_index_fail_count: int = 3
+
+    # Excel parser limits
+    xlsx_max_rows_per_sheet: int = 10000  # max rows to index per sheet
+    xlsx_max_consecutive_empty_rows: int = 100  # stop parsing after this many consecutive empty rows
+
+    # Shadow Markdown: generate a hidden .md copy for binary files (PDF, DOCX, etc.)
+    # so ripgrep fallback can search them. Disable to speed up indexing.
+    enable_shadow_md: bool = True
+
+    # Allowed source types: if non-empty, only files matching these source types
+    # will be indexed. Empty list means all types are allowed (backward compatible).
+    # Valid values: markdown, code, text, json, jsonl, csv, html, xml, pdf, doc,
+    # docx, pptx, excel
+    allowed_source_types: list[str] = field(default_factory=list)
+
+    @classmethod
+    def from_env(cls) -> "TreeSearchConfig":
+        """Create config from environment variables, falling back to defaults."""
+        config = cls()
+
+        env_cjk = os.getenv(_ENV_CJK_TOKENIZER)
+        if env_cjk:
+            config.cjk_tokenizer = env_cjk
+
+        env_fp = os.getenv(_ENV_FINGERPRINT_MODE)
+        if env_fp in ("stat", "content"):
+            config.fingerprint_mode = env_fp
+
+        env_prune = os.getenv(_ENV_PRUNE)
+        if env_prune is not None:
+            config.prune_orphans_on_directory = env_prune.lower() in ("1", "true", "yes")
+
+        env_shadow = os.getenv(_ENV_SHADOW_MD)
+        if env_shadow is not None:
+            config.enable_shadow_md = env_shadow.lower() in ("1", "true", "yes")
+
+        _env_int(config, "xlsx_max_rows_per_sheet", 1, "TREESEARCH_XLSX_MAX_ROWS_PER_SHEET")
+        _env_int(config, "xlsx_max_consecutive_empty_rows", 1, "TREESEARCH_XLSX_MAX_CONSECUTIVE_EMPTY_ROWS")
+
+        env_source_types = os.getenv(_ENV_ALLOWED_SOURCE_TYPES)
+        if env_source_types:
+            config.allowed_source_types = [
+                t.strip() for t in env_source_types.split(",") if t.strip()
+            ]
+
+        return config
+
+
+# ---------------------------------------------------------------------------
+# Global singleton
+# ---------------------------------------------------------------------------
+_default_config: Optional[TreeSearchConfig] = None
+
+
+def get_config(reload: bool = False) -> TreeSearchConfig:
+    """Get global configuration (lazy singleton).
+
+    First call reads env vars + defaults. Subsequent calls return cached instance.
+    """
+    global _default_config
+    if reload or _default_config is None:
+        _default_config = TreeSearchConfig.from_env()
+    return _default_config
+
+
+def set_config(config: TreeSearchConfig) -> None:
+    """Set global configuration (highest priority)."""
+    global _default_config
+    _default_config = config
+
+
+def reset_config() -> None:
+    """Reset global config. Next get_config() re-initializes from env."""
+    global _default_config
+    _default_config = None