source-kb 0.2.2__py3-none-any.whl

core/rag/bm25_index.py

@@ -0,0 +1,268 @@
+"""BM25 keyword-based fallback index.
+
+Used when embedding API is unavailable or embed_backend is "none".
+Provides the same query interface as the vector index.
+
+Storage: {knowledge_dir}/.bm25-index.json
+
+Usage:
+    from core.rag.bm25_index import BM25Index
+
+    index = BM25Index(Path("knowledge/my-kb/.bm25-index.json"))
+    index.build(Path("knowledge/my-kb"))
+    index.save()
+
+    results = index.search("points issuance", top_k=5)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import math
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# BM25 parameters
+_K1 = 1.5
+_B = 0.75
+
+
+@dataclass
+class SearchResult:
+    """Single search result."""
+
+    doc_path: str
+    chunk_text: str
+    score: float
+    metadata: dict | None = None
+
+
+class BM25Index:
+    """BM25 keyword index with Chinese tokenization support."""
+
+    def __init__(self, index_path: Path):
+        self._path = index_path
+        self._documents: list[dict] = []  # [{path, chunk_idx, text, terms, tf}]
+        self._idf: dict[str, float] = {}
+        self._avg_dl: float = 0.0
+        self._total_docs: int = 0
+
+    def build(self, knowledge_dir: Path) -> None:
+        """Build index from all .md files in knowledge_dir.
+
+        Splits documents into chunks (~500-1000 chars) and tokenizes.
+        """
+        self._documents = []
+        md_files = sorted(knowledge_dir.rglob("*.md"))
+
+        for md_file in md_files:
+            if md_file.name.startswith("."):
+                continue
+            try:
+                content = md_file.read_text(encoding="utf-8", errors="replace")
+            except OSError:
+                continue
+
+            rel_path = str(md_file.relative_to(knowledge_dir))
+            chunks = _split_into_chunks(content)
+
+            for i, chunk in enumerate(chunks):
+                terms = _tokenize(chunk)
+                tf = _compute_tf(terms)
+                self._documents.append({
+                    "path": rel_path,
+                    "chunk_idx": i,
+                    "text": chunk[:500],  # Store truncated for display
+                    "terms": terms,
+                    "tf": tf,
+                    "dl": len(terms),
+                })
+
+        # Compute IDF
+        self._total_docs = len(self._documents)
+        if self._total_docs == 0:
+            return
+
+        self._avg_dl = sum(d["dl"] for d in self._documents) / self._total_docs
+        df: dict[str, int] = {}
+        for doc in self._documents:
+            seen: set[str] = set()
+            for term in doc["terms"]:
+                if term not in seen:
+                    df[term] = df.get(term, 0) + 1
+                    seen.add(term)
+
+        self._idf = {
+            term: math.log((self._total_docs - freq + 0.5) / (freq + 0.5) + 1)
+            for term, freq in df.items()
+        }
+
+        logger.info("[bm25] Built index: %d docs, %d chunks, %d unique terms",
+                    len(md_files), self._total_docs, len(self._idf))
+
+    def search(self, query: str, top_k: int = 10) -> list[SearchResult]:
+        """Search index using BM25 scoring."""
+        if not self._documents:
+            return []
+
+        query_terms = _tokenize(query)
+        if not query_terms:
+            return []
+
+        scores: list[tuple[int, float]] = []
+
+        for idx, doc in enumerate(self._documents):
+            score = 0.0
+            dl = doc["dl"]
+            tf = doc["tf"]
+
+            for term in query_terms:
+                if term not in tf:
+                    continue
+                term_tf = tf[term]
+                idf = self._idf.get(term, 0.0)
+                # BM25 formula
+                numerator = term_tf * (_K1 + 1)
+                denominator = term_tf + _K1 * (1 - _B + _B * dl / max(self._avg_dl, 1))
+                score += idf * (numerator / denominator)
+
+            if score > 0:
+                scores.append((idx, score))
+
+        # Sort by score descending
+        scores.sort(key=lambda x: -x[1])
+
+        results: list[SearchResult] = []
+        for idx, score in scores[:top_k]:
+            doc = self._documents[idx]
+            results.append(SearchResult(
+                doc_path=doc["path"],
+                chunk_text=doc["text"],
+                score=round(score, 4),
+                metadata={"chunk_idx": doc["chunk_idx"], "source": doc["path"]},
+            ))
+
+        return results
+
+    def save(self) -> None:
+        """Persist index to JSON file."""
+        data = {
+            "version": "1.0",
+            "total_docs": self._total_docs,
+            "avg_dl": self._avg_dl,
+            "idf": self._idf,
+            "documents": [
+                {"path": d["path"], "chunk_idx": d["chunk_idx"],
+                 "text": d["text"], "tf": d["tf"], "dl": d["dl"]}
+                for d in self._documents
+            ],
+        }
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        self._path.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
+        logger.info("[bm25] Saved index: %s (%d docs)", self._path, self._total_docs)
+
+    @classmethod
+    def load(cls, index_path: Path) -> BM25Index:
+        """Load existing index from file."""
+        index = cls(index_path)
+        if not index_path.exists():
+            return index
+
+        try:
+            data = json.loads(index_path.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError):
+            return index
+
+        index._total_docs = data.get("total_docs", 0)
+        index._avg_dl = data.get("avg_dl", 0.0)
+        index._idf = data.get("idf", {})
+        index._documents = [
+            {**d, "terms": list(d.get("tf", {}).keys())}
+            for d in data.get("documents", [])
+        ]
+
+        return index
+
+
+# ---------------------------------------------------------------------------
+# Tokenization
+# ---------------------------------------------------------------------------
+
+
+def _tokenize(text: str) -> list[str]:
+    """Tokenize text using character n-grams for Chinese + word splitting for English.
+
+    Falls back to simple splitting if jieba is unavailable.
+    """
+    if not text:
+        return []
+
+    terms: list[str] = []
+
+    # Try jieba first
+    try:
+        import jieba
+        words = jieba.lcut_for_search(text)
+        terms = [w.strip().lower() for w in words if len(w.strip()) > 1]
+        return terms
+    except ImportError:
+        pass
+
+    # Fallback: simple tokenization
+    # Chinese: bigrams
+    chinese_chars = re.findall(r'[\u4e00-\u9fff]+', text)
+    for segment in chinese_chars:
+        for i in range(len(segment) - 1):
+            terms.append(segment[i:i + 2])
+        if len(segment) >= 3:
+            for i in range(len(segment) - 2):
+                terms.append(segment[i:i + 3])
+
+    # English/code: word splitting
+    english_words = re.findall(r'[a-zA-Z_]\w{2,}', text)
+    terms.extend(w.lower() for w in english_words)
+
+    return terms
+
+
+def _compute_tf(terms: list[str]) -> dict[str, int]:
+    """Compute term frequency."""
+    tf: dict[str, int] = {}
+    for term in terms:
+        tf[term] = tf.get(term, 0) + 1
+    return tf
+
+
+def _split_into_chunks(content: str, target_size: int = 800) -> list[str]:
+    """Split document into chunks of approximately target_size characters.
+
+    Splits on paragraph boundaries (double newline) or heading boundaries.
+    """
+    if len(content) <= target_size:
+        return [content]
+
+    chunks: list[str] = []
+    current: list[str] = []
+    current_len = 0
+
+    paragraphs = re.split(r'\n\n+', content)
+
+    for para in paragraphs:
+        para_len = len(para)
+
+        if current_len + para_len > target_size and current:
+            chunks.append("\n\n".join(current))
+            current = []
+            current_len = 0
+
+        current.append(para)
+        current_len += para_len
+
+    if current:
+        chunks.append("\n\n".join(current))
+
+    return chunks

core/rag/chunker.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+"""Semantic chunker — splits by headings, paragraphs, and code structure.
+
+Migrated from engine/rag/chunker.py. Removes print() calls (no CLI I/O in core/).
+Supports two modes:
+  - "semantic" (default): splits markdown by headings, then by sentence boundaries
+  - "ast": splits source code at function/class boundaries
+"""
+
+import logging
+import re
+from dataclasses import dataclass
+
+from core.rag.loader import Document
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Chunk:
+    """A text chunk with associated metadata."""
+    text: str
+    metadata: dict
+
+
+def chunk_documents(
+    docs: list[Document],
+    size: int = 512,
+    overlap: int = 100,
+    min_length: int = 50,
+    mode: str = "semantic",
+) -> list[Chunk]:
+    """Chunk a list of documents.
+
+    Args:
+        docs: List of documents to chunk.
+        size: Maximum chunk size in characters.
+        overlap: Overlap between chunks in characters.
+        min_length: Minimum chunk length filter threshold in characters.
+        mode: Chunking mode — "semantic" (by headings) or "ast" (by code structure).
+    """
+    chunks: list[Chunk] = []
+    for doc in docs:
+        if mode == "ast" and _is_code_file(doc):
+            chunks.extend(_chunk_ast(doc, size, min_length))
+        else:
+            chunks.extend(_chunk_one(doc, size, overlap, min_length))
+    logger.info("Generated %d chunks (mode=%s)", len(chunks), mode)
+    return chunks
+
+
+def _is_code_file(doc: Document) -> bool:
+    """Check if document is a source code file (not markdown)."""
+    source = doc.metadata.get("source", "")
+    code_extensions = (".java", ".py", ".go", ".kt", ".ts", ".js", ".rs", ".cpp", ".c")
+    return any(source.endswith(ext) for ext in code_extensions)
+
+
+def _chunk_ast(doc: Document, size: int, min_length: int) -> list[Chunk]:
+    """AST-aware chunking: split at function/class boundaries."""
+    lines = doc.content.splitlines(keepends=True)
+    if not lines:
+        return []
+
+    boundaries = _find_code_boundaries(lines)
+    if not boundaries:
+        return _chunk_by_lines(doc, lines, size, min_length)
+
+    chunks: list[Chunk] = []
+    for i, start in enumerate(boundaries):
+        end = boundaries[i + 1] if i + 1 < len(boundaries) else len(lines)
+        block = "".join(lines[start:end]).strip()
+
+        if len(block) < min_length:
+            continue
+
+        if len(block) <= size:
+            chunks.append(Chunk(
+                text=block,
+                metadata={**doc.metadata, "section": _extract_symbol_name(lines[start])},
+            ))
+        else:
+            sub_chunks = _split_large_block(block, size, min_length)
+            symbol = _extract_symbol_name(lines[start])
+            for j, sub in enumerate(sub_chunks):
+                chunks.append(Chunk(
+                    text=sub,
+                    metadata={**doc.metadata, "section": f"{symbol} (part {j+1})"},
+                ))
+
+    if not chunks:
+        return _chunk_by_lines(doc, lines, size, min_length)
+    return chunks
+
+
+def _find_code_boundaries(lines: list[str]) -> list[int]:
+    """Find line indices where class/function definitions start."""
+    class_pattern = re.compile(
+        r'^\s*(?:public|protected|private|static|abstract|final|\s)*'
+        r'(?:class|interface|enum|record)\s+\w+'
+    )
+    method_pattern = re.compile(
+        r'^\s*(?:public|protected|private|static|abstract|final|synchronized|native|\s)*'
+        r'(?:[\w<>\[\],\s]+)\s+\w+\s*\('
+    )
+    py_pattern = re.compile(r'^(?:class|def|async\s+def)\s+\w+')
+
+    boundaries: list[int] = []
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+        if not stripped or stripped.startswith("//") or stripped.startswith("#"):
+            continue
+        if (class_pattern.match(stripped)
+                or (method_pattern.match(stripped) and not stripped.startswith("return"))
+                or py_pattern.match(stripped)):
+            boundaries.append(i)
+    return boundaries
+
+
+def _extract_symbol_name(line: str) -> str:
+    """Extract symbol name from a declaration line."""
+    m = re.search(r'(?:class|interface|enum|record|def)\s+(\w+)', line)
+    if m:
+        return m.group(1)
+    m = re.search(r'(\w+)\s*\(', line)
+    if m:
+        return m.group(1)
+    return line.strip()[:30]
+
+
+def _split_large_block(block: str, size: int, min_length: int) -> list[str]:
+    """Split an oversized code block at logical boundaries."""
+    lines = block.splitlines(keepends=True)
+    chunks: list[str] = []
+    current: list[str] = []
+    current_len = 0
+
+    for line in lines:
+        if current_len + len(line) > size and current_len >= min_length:
+            chunks.append("".join(current))
+            overlap_lines = current[-3:] if len(current) > 3 else []
+            current = overlap_lines.copy()
+            current_len = sum(len(ln) for ln in current)
+        current.append(line)
+        current_len += len(line)
+
+    if current:
+        chunks.append("".join(current))
+    return chunks
+
+
+def _chunk_by_lines(doc: Document, lines: list[str], size: int, min_length: int) -> list[Chunk]:
+    """Fallback: chunk code file by character size with overlap."""
+    chunks: list[Chunk] = []
+    current: list[str] = []
+    current_len = 0
+
+    for line in lines:
+        if current_len + len(line) > size and current_len >= min_length:
+            chunks.append(Chunk(
+                text="".join(current),
+                metadata={**doc.metadata, "section": ""},
+            ))
+            current = current[-3:]
+            current_len = sum(len(ln) for ln in current)
+        current.append(line)
+        current_len += len(line)
+
+    if current and current_len >= min_length:
+        chunks.append(Chunk(
+            text="".join(current),
+            metadata={**doc.metadata, "section": ""},
+        ))
+    return chunks
+
+
+def _chunk_one(doc: Document, size: int, overlap: int, min_length: int) -> list[Chunk]:
+    """Split by markdown headings, then by sentence boundaries for large sections."""
+    sections = re.split(r'\n(?=#{1,3}\s)', doc.content)
+    chunks: list[Chunk] = []
+
+    for section in sections:
+        section = section.strip()
+        if not section:
+            continue
+
+        lines = section.split("\n")
+        section_title = lines[0].lstrip("#").strip() if lines[0].startswith("#") else ""
+        content_text = "\n".join(lines[1:]).strip() if len(lines) > 1 else ""
+
+        if len(content_text) < min_length:
+            if section_title and content_text:
+                chunks.append(Chunk(
+                    text=section,
+                    metadata={**doc.metadata, "section": section_title},
+                ))
+            continue
+
+        if len(section) <= size:
+            chunks.append(Chunk(
+                text=section,
+                metadata={**doc.metadata, "section": section_title},
+            ))
+        else:
+            title_prefix = f"## {section_title}\n" if section_title else ""
+            available_size = size - len(title_prefix)
+            for sub in _sliding_window_sentence(content_text, available_size, overlap):
+                chunk_text = title_prefix + sub if title_prefix else sub
+                chunks.append(Chunk(
+                    text=chunk_text,
+                    metadata={**doc.metadata, "section": section_title},
+                ))
+    return chunks
+
+
+_SENTENCE_END = re.compile(r'(?<=[。！？；\n.!?;])\s*')
+
+
+def _split_sentences(text: str) -> list[str]:
+    """Split by sentence/paragraph boundaries."""
+    parts = _SENTENCE_END.split(text)
+    sentences: list[str] = []
+    buf = ""
+    for p in parts:
+        buf += p
+        if len(buf) >= 20:
+            sentences.append(buf)
+            buf = ""
+    if buf.strip():
+        if sentences:
+            sentences[-1] += buf
+        else:
+            sentences.append(buf)
+    return sentences
+
+
+def _sliding_window_sentence(text: str, size: int, overlap: int) -> list[str]:
+    """Sliding window split by sentence boundaries."""
+    sentences = _split_sentences(text)
+    if not sentences:
+        return [text] if text.strip() else []
+
+    result: list[str] = []
+    i = 0
+    while i < len(sentences):
+        chunk_parts: list[str] = []
+        chunk_len = 0
+        j = i
+        while j < len(sentences):
+            s = sentences[j]
+            if chunk_len + len(s) > size and chunk_parts:
+                break
+            chunk_parts.append(s)
+            chunk_len += len(s)
+            j += 1
+
+        result.append("".join(chunk_parts))
+
+        if j >= len(sentences):
+            break
+
+        overlap_start = j
+        overlap_len = 0
+        for k in range(j - 1, i - 1, -1):
+            if overlap_len + len(sentences[k]) > overlap:
+                break
+            overlap_len += len(sentences[k])
+            overlap_start = k
+
+        i = overlap_start if overlap_start > i else j
+
+    return result

core/rag/embedder.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+"""Embedding wrapper — supports Ollama / OpenAI / DashScope / ChromaDB built-in.
+
+Migrated from engine/rag/embedder.py. Removes global config dependency;
+accepts Config object explicitly. Embedding API calls are allowed in core/
+(they are not generative LLM calls).
+
+Rate-limit retry is built into each backend.
+"""
+
+import logging
+import threading
+import time
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from core.config import Config
+
+logger = logging.getLogger(__name__)
+
+MAX_RETRIES = 5  # Max retries for 429 rate limiting
+
+
+def get_embeddings(texts: list[str], config: "Config") -> list[list[float]]:
+    """Select embedding backend based on configuration.
+
+    Args:
+        texts: List of text strings to embed.
+        config: Config object providing embed_backend, embed_model, etc.
+
+    Returns:
+        List of embedding vectors (one per input text).
+    """
+    if not texts:
+        return []
+    backend = config.embed_backend
+    if backend == "chromadb-default":
+        return _chromadb_default_embed(texts)
+    elif backend == "dashscope":
+        return _dashscope_embed(texts, config)
+    elif backend == "openai":
+        return _openai_embed(texts, config)
+    else:  # ollama
+        return _ollama_embed(texts, config)
+
+
+_default_ef = None
+_default_ef_lock = threading.Lock()
+
+
+def _chromadb_default_embed(texts: list[str]) -> list[list[float]]:
+    """ChromaDB built-in embedding (all-MiniLM-L6-v2, runs locally)."""
+    global _default_ef
+    if _default_ef is None:
+        with _default_ef_lock:
+            if _default_ef is None:
+                from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
+                _default_ef = DefaultEmbeddingFunction()
+    return _default_ef(texts)
+
+
+def _ollama_embed(texts: list[str], config: "Config") -> list[list[float]]:
+    """Ollama local embedding with retry on rate limit."""
+    import requests
+
+    for attempt in range(1, MAX_RETRIES + 1):
+        resp = requests.post(
+            f"{config.embed_base_url}/api/embed",
+            json={"model": config.embed_model, "input": texts},
+            timeout=60,
+        )
+        if resp.status_code == 429:
+            if attempt >= MAX_RETRIES:
+                raise RuntimeError(
+                    f"Ollama embedding rate limited after {MAX_RETRIES} retries"
+                )
+            retry_after = int(resp.headers.get("Retry-After", 5))
+            time.sleep(retry_after)
+            continue
+        resp.raise_for_status()
+        return resp.json()["embeddings"]
+    raise RuntimeError(f"Ollama embedding failed after {MAX_RETRIES} retries")
+
+
+def _openai_embed(texts: list[str], config: "Config") -> list[list[float]]:
+    """OpenAI-compatible API (also works with vLLM, LocalAI, etc.).
+
+    Automatically batches requests (max 10 per call) to respect API limits.
+    """
+    import requests
+
+    headers = {"Content-Type": "application/json"}
+    api_key = config.embed_api_key
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    batch_size = 10
+    results: list[list[float]] = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i:i + batch_size]
+        for attempt in range(1, MAX_RETRIES + 1):
+            resp = requests.post(
+                f"{config.embed_base_url}/v1/embeddings",
+                headers=headers,
+                json={"model": config.embed_model, "input": batch},
+                timeout=60,
+            )
+            if resp.status_code == 429:
+                if attempt >= MAX_RETRIES:
+                    raise RuntimeError(
+                        f"OpenAI embedding rate limited after {MAX_RETRIES} retries"
+                    )
+                retry_after = int(resp.headers.get("Retry-After", 5))
+                time.sleep(retry_after)
+                continue
+            resp.raise_for_status()
+            data = resp.json()
+            results.extend([item["embedding"] for item in data["data"]])
+            break
+    return results
+
+
+def _dashscope_embed(texts: list[str], config: "Config") -> list[list[float]]:
+    """DashScope text-embedding with retry on rate limit."""
+    import dashscope
+
+    api_key = config.embed_api_key
+    results: list[list[float]] = []
+    batch_size = 25  # DashScope max 25 per request
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i:i + batch_size]
+        for attempt in range(1, MAX_RETRIES + 1):
+            resp = dashscope.TextEmbedding.call(
+                model=config.embed_model,
+                input=batch,
+                api_key=api_key,
+            )
+            if resp.status_code == 429:
+                if attempt >= MAX_RETRIES:
+                    raise RuntimeError(
+                        f"DashScope embedding rate limited after {MAX_RETRIES} retries"
+                    )
+                time.sleep(attempt * 2)
+                continue
+            if resp.status_code != 200:
+                raise RuntimeError(f"DashScope embedding failed: {resp.message}")
+            for item in resp.output["embeddings"]:
+                results.append(item["embedding"])
+            break
+    return results