mikoshi 0.1.9__py3-none-any.whl

mikoshi/config.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+DEFAULT_MAX_BYTES = 1_000_000
+DEFAULT_CHUNK_LINES = 120
+DEFAULT_CHUNK_OVERLAP = 20
+DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+DEFAULT_OPENAI_EMBED_MODEL = "text-embedding-3-small"
+
+
+@dataclass(frozen=True)
+class EmbeddingConfig:
+    provider: str
+    model: str
+    openai_api_key: str | None
+    openai_base_url: str | None
+
+
+@dataclass(frozen=True)
+class Config:
+    max_bytes: int
+    chunk_lines: int
+    chunk_overlap: int
+    embeddings: EmbeddingConfig
+    index_root: Path
+    quiet_external_libs: bool
+
+
+class ConfigError(RuntimeError):
+    pass
+
+
+def _env_bool(name: str, default: bool) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    value = raw.strip().lower()
+    if value in {"1", "true", "yes", "on"}:
+        return True
+    if value in {"0", "false", "no", "off"}:
+        return False
+    return default
+
+
+_EXTERNAL_LIBS_CONFIGURED = False
+
+
+def configure_external_libs(quiet: bool) -> None:
+    global _EXTERNAL_LIBS_CONFIGURED
+    if not quiet or _EXTERNAL_LIBS_CONFIGURED:
+        return
+
+    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+    os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    os.environ["PYTHONWARNINGS"] = "ignore"
+    os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+    os.environ.setdefault("MKL_NUM_THREADS", "1")
+    os.environ.setdefault("VECLIB_MAXIMUM_THREADS", "1")
+    os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
+
+    import logging
+
+    logging.getLogger("transformers").setLevel(logging.ERROR)
+    logging.getLogger("sentence_transformers").setLevel(logging.ERROR)
+    logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
+
+    try:
+        from transformers.utils import logging as transformers_logging
+
+        transformers_logging.set_verbosity_error()
+        if hasattr(transformers_logging, "disable_progress_bar"):
+            transformers_logging.disable_progress_bar()
+    except Exception:
+        pass
+
+    _EXTERNAL_LIBS_CONFIGURED = True
+
+
+def load_config() -> Config:
+    provider = os.getenv("MIKOSHI_EMBEDDINGS_PROVIDER", "local").strip().lower()
+    if provider not in {"local", "openai"}:
+        raise ConfigError(
+            "MIKOSHI_EMBEDDINGS_PROVIDER must be 'local' or 'openai'."
+        )
+
+    model = os.getenv("MIKOSHI_EMBEDDINGS_MODEL", DEFAULT_EMBEDDING_MODEL)
+    openai_model = os.getenv("MIKOSHI_OPENAI_EMBED_MODEL", DEFAULT_OPENAI_EMBED_MODEL)
+
+    openai_api_key = os.getenv("MIKOSHI_OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY")
+    openai_base_url = os.getenv("MIKOSHI_OPENAI_BASE_URL") or os.getenv("OPENAI_BASE_URL")
+
+    if provider == "openai":
+        if not openai_api_key:
+            raise ConfigError(
+                "OpenAI embeddings selected but no API key found. "
+                "Set MIKOSHI_OPENAI_API_KEY or OPENAI_API_KEY."
+            )
+        model = openai_model
+
+    index_root = Path(os.getenv("MIKOSHI_INDEX_ROOT", "~/.mikoshi")).expanduser()
+    max_bytes = int(os.getenv("MIKOSHI_MAX_BYTES", str(DEFAULT_MAX_BYTES)))
+    chunk_lines = int(os.getenv("MIKOSHI_CHUNK_LINES", str(DEFAULT_CHUNK_LINES)))
+    chunk_overlap = int(os.getenv("MIKOSHI_CHUNK_OVERLAP", str(DEFAULT_CHUNK_OVERLAP)))
+    quiet_external_libs = _env_bool("MIKOSHI_QUIET_EXTERNAL_LIBS", True)
+
+    embeddings = EmbeddingConfig(
+        provider=provider,
+        model=model,
+        openai_api_key=openai_api_key,
+        openai_base_url=openai_base_url,
+    )
+
+    return Config(
+        max_bytes=max_bytes,
+        chunk_lines=chunk_lines,
+        chunk_overlap=chunk_overlap,
+        embeddings=embeddings,
+        index_root=index_root,
+        quiet_external_libs=quiet_external_libs,
+    )

mikoshi/entitlements.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+DEFAULT_PLAN = "free"
+DEFAULT_FEATURES = ("local_index", "mcp")
+
+
+class UpgradeRequired(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class Entitlements:
+    plan: str
+    features: frozenset[str]
+
+    def has(self, feature: str) -> bool:
+        return feature in self.features
+
+    def require(self, feature: str) -> None:
+        if not self.has(feature):
+            raise UpgradeRequired(f"Available with Mikoshi Pro: {feature}")
+
+
+def from_plan_features(plan: str | None, features: list[str] | None) -> Entitlements:
+    normalized_plan = (plan or DEFAULT_PLAN).strip().lower()
+    normalized_features = frozenset(
+        {item.strip() for item in (features or list(DEFAULT_FEATURES)) if item}
+    )
+    return Entitlements(plan=normalized_plan, features=normalized_features)

mikoshi/hashing.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import hashlib
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def sha256_text(text: str) -> str:
+    return sha256_bytes(text.encode("utf-8"))

mikoshi/ignore.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Iterable
+
+from pathspec import PathSpec
+
+COMMON_IGNORE_PATTERNS = [
+    "**/.git/",
+    "**/.hg/",
+    "**/.svn/",
+    "**/.DS_Store",
+    "**/__pycache__/",
+    "**/*.pyc",
+    "**/.venv/",
+    "**/venv/",
+    "**/node_modules/",
+    "**/dist/",
+    "**/build/",
+    "**/.idea/",
+    "**/.vscode/",
+    "**/.mypy_cache/",
+    "**/.pytest_cache/",
+    "**/coverage/",
+    "**/.coverage",
+    "**/*.log",
+]
+
+COMMON_IGNORE_DIRS = {
+    ".git",
+    ".hg",
+    ".svn",
+    "node_modules",
+    "dist",
+    "build",
+    ".venv",
+    "venv",
+    "__pycache__",
+    ".idea",
+    ".vscode",
+    ".mypy_cache",
+    ".pytest_cache",
+}
+
+
+def _convert_gitignore_pattern(
+    pattern: str,
+    base_dir: Path,
+    repo_root: Path,
+) -> str | None:
+    stripped = pattern.strip()
+    if not stripped or stripped.startswith("#"):
+        return None
+
+    negated = stripped.startswith("!")
+    if negated:
+        stripped = stripped[1:]
+
+    if not stripped:
+        return None
+
+    rel_base = base_dir.relative_to(repo_root).as_posix()
+    prefix = "" if rel_base == "." else f"{rel_base}/"
+
+    anchored = stripped.startswith("/")
+    if anchored:
+        body = stripped.lstrip("/")
+        converted = f"{prefix}{body}"
+    else:
+        if "/" in stripped:
+            converted = f"{prefix}{stripped}"
+        else:
+            converted = f"{prefix}**/{stripped}"
+
+    return f"!{converted}" if negated else converted
+
+
+def _load_gitignore_files(repo_root: Path) -> list[Path]:
+    gitignores: list[Path] = []
+    for root, dirs, files in os.walk(repo_root):
+        dirs[:] = [d for d in dirs if d not in COMMON_IGNORE_DIRS]
+        if ".gitignore" in files:
+            gitignores.append(Path(root) / ".gitignore")
+    return sorted(gitignores, key=lambda p: len(p.relative_to(repo_root).parts))
+
+
+def _load_patterns(repo_root: Path) -> list[str]:
+    patterns: list[str] = list(COMMON_IGNORE_PATTERNS)
+    for gitignore in _load_gitignore_files(repo_root):
+        try:
+            content = gitignore.read_text(encoding="utf-8")
+        except OSError:
+            continue
+        for line in content.splitlines():
+            converted = _convert_gitignore_pattern(line, gitignore.parent, repo_root)
+            if converted:
+                patterns.append(converted)
+    return patterns
+
+
+def build_ignore_spec(repo_root: Path) -> PathSpec:
+    patterns = _load_patterns(repo_root)
+    return PathSpec.from_lines("gitwildmatch", patterns)
+
+
+class IgnoreMatcher:
+    def __init__(self, repo_root: Path) -> None:
+        self.repo_root = repo_root
+        self.spec = build_ignore_spec(repo_root)
+
+    def is_ignored(self, path: Path) -> bool:
+        try:
+            relpath = path.relative_to(self.repo_root).as_posix()
+        except ValueError:
+            relpath = path.as_posix()
+        if path.is_dir() and not relpath.endswith("/"):
+            relpath = f"{relpath}/"
+        if self.spec.match_file(relpath):
+            return True
+        if not path.is_dir():
+            parent = path.parent
+            while parent != self.repo_root and parent != parent.parent:
+                parent_rel = parent.relative_to(self.repo_root).as_posix()
+                if not parent_rel.endswith("/"):
+                    parent_rel = f"{parent_rel}/"
+                if self.spec.match_file(parent_rel):
+                    return True
+                parent = parent.parent
+        return False
+
+    def is_ignored_rel(self, relpath: str) -> bool:
+        return self.spec.match_file(relpath)
+
+
+def iter_ignored(paths: Iterable[Path], matcher: IgnoreMatcher) -> Iterable[Path]:
+    for path in paths:
+        if matcher.is_ignored(path):
+            yield path

mikoshi/indexing/__init__.py

@@ -0,0 +1,9 @@
+from .indexer import index_repo
+from .index_store import IndexStore
+from .file_scanner import scan_repo
+
+__all__ = [
+    "index_repo",
+    "IndexStore",
+    "scan_repo",
+]

mikoshi/indexing/file_scanner.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from mikoshi.ignore import IgnoreMatcher
+
+
+class ScanError(RuntimeError):
+    pass
+
+
+def _is_binary(path: Path) -> bool:
+    try:
+        with path.open("rb") as handle:
+            chunk = handle.read(4096)
+    except OSError:
+        return True
+    if b"\x00" in chunk:
+        return True
+    try:
+        chunk.decode("utf-8")
+    except UnicodeDecodeError:
+        return True
+    return False
+
+
+def scan_repo(
+    repo_root: Path,
+    matcher: IgnoreMatcher,
+    max_bytes: int,
+) -> list[Path]:
+    if not repo_root.exists() or not repo_root.is_dir():
+        raise ScanError(f"Repository path does not exist: {repo_root}")
+
+    files: list[Path] = []
+    for root, dirs, filenames in os.walk(repo_root):
+        root_path = Path(root)
+        dirs[:] = [
+            d
+            for d in dirs
+            if not matcher.is_ignored(root_path / d)
+        ]
+        for name in filenames:
+            file_path = root_path / name
+            if matcher.is_ignored(file_path):
+                continue
+            if file_path.is_symlink():
+                continue
+            try:
+                size = file_path.stat().st_size
+            except OSError:
+                continue
+            if size > max_bytes:
+                continue
+            if _is_binary(file_path):
+                continue
+            files.append(file_path)
+
+    return sorted(files)

mikoshi/indexing/index_store.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+import faiss
+
+from mikoshi.hashing import sha256_text
+from mikoshi.utils.types import Chunk, IndexMeta
+
+
+class IndexStore:
+    def __init__(self, repo_root: Path, index_root: Path) -> None:
+        self.repo_root = repo_root
+        self.repo_id = repo_id_for_path(repo_root)
+        self.store_dir = index_root / self.repo_id
+        self.meta_path = self.store_dir / "meta.json"
+        self.chunks_path = self.store_dir / "chunks.jsonl"
+        self.embeddings_path = self.store_dir / "embeddings.npy"
+        self.faiss_path = self.store_dir / "index.faiss"
+
+    def exists(self) -> bool:
+        return self.meta_path.exists() and self.chunks_path.exists()
+
+    def ensure_dir(self) -> None:
+        self.store_dir.mkdir(parents=True, exist_ok=True)
+
+    def clear(self) -> None:
+        if self.store_dir.exists():
+            shutil.rmtree(self.store_dir)
+
+    def load_meta(self) -> IndexMeta | None:
+        if not self.meta_path.exists():
+            return None
+        data = json.loads(self.meta_path.read_text(encoding="utf-8"))
+        return IndexMeta.from_dict(data)
+
+    def save_meta(self, meta: IndexMeta) -> None:
+        self.ensure_dir()
+        self.meta_path.write_text(
+            json.dumps(meta.to_dict(), indent=2, sort_keys=True),
+            encoding="utf-8",
+        )
+
+    def load_chunks(self) -> list[Chunk]:
+        if not self.chunks_path.exists():
+            return []
+        chunks: list[Chunk] = []
+        with self.chunks_path.open("r", encoding="utf-8") as handle:
+            for line in handle:
+                line = line.strip()
+                if not line:
+                    continue
+                chunks.append(Chunk.from_dict(json.loads(line)))
+        return chunks
+
+    def save_chunks(self, chunks: Iterable[Chunk]) -> None:
+        self.ensure_dir()
+        with self.chunks_path.open("w", encoding="utf-8") as handle:
+            for chunk in chunks:
+                handle.write(json.dumps(chunk.to_dict(), ensure_ascii=False))
+                handle.write("\n")
+
+    def load_embeddings(self) -> np.ndarray | None:
+        if not self.embeddings_path.exists():
+            return None
+        return np.load(self.embeddings_path)
+
+    def save_embeddings(self, embeddings: np.ndarray) -> None:
+        self.ensure_dir()
+        np.save(self.embeddings_path, embeddings)
+
+    def load_faiss(self) -> faiss.Index | None:
+        if not self.faiss_path.exists():
+            return None
+        return faiss.read_index(str(self.faiss_path))
+
+    def save_faiss(self, index: faiss.Index) -> None:
+        self.ensure_dir()
+        faiss.write_index(index, str(self.faiss_path))
+
+
+def repo_id_for_path(repo_root: Path) -> str:
+    return sha256_text(str(repo_root.resolve()))[:16]

mikoshi/indexing/indexer.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Iterable
+
+import faiss
+import numpy as np
+
+from mikoshi.chunking import chunk_text
+from mikoshi.config import Config, load_config
+from mikoshi.hashing import sha256_bytes, sha256_text
+from mikoshi.ignore import IgnoreMatcher
+from mikoshi.indexing.file_scanner import scan_repo
+from mikoshi.indexing.index_store import IndexStore
+from mikoshi.retrieval.semantic import get_embeddings_provider, normalize_embeddings
+from mikoshi.utils.timer import Timer
+from mikoshi.utils.types import Chunk, IndexMeta, IndexResult
+
+
+class IndexerError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class IndexedFile:
+    path: Path
+    relpath: str
+    file_hash: str
+    text: str
+
+
+def _read_file(path: Path) -> bytes:
+    return path.read_bytes()
+
+
+def _prepare_file(path: Path, repo_root: Path) -> IndexedFile:
+    data = _read_file(path)
+    file_hash = sha256_bytes(data)
+    text = data.decode("utf-8", errors="ignore")
+    relpath = path.relative_to(repo_root).as_posix()
+    return IndexedFile(path=path, relpath=relpath, file_hash=file_hash, text=text)
+
+
+def _chunks_from_text(
+    relpath: str,
+    file_hash: str,
+    text: str,
+    max_lines: int,
+    overlap: int,
+) -> list[Chunk]:
+    spans = chunk_text(text, max_lines, overlap)
+    chunks: list[Chunk] = []
+    for span in spans:
+        chunk_id = sha256_text(
+            f"{relpath}:{span.start_line}:{span.end_line}:{span.text}"
+        )
+        chunks.append(
+            Chunk(
+                id=chunk_id,
+                relpath=relpath,
+                start_line=span.start_line,
+                end_line=span.end_line,
+                text=span.text,
+                file_hash=file_hash,
+                vector_idx=None,
+            )
+        )
+    return chunks
+
+
+def _with_vector_idx(chunks: Iterable[Chunk]) -> list[Chunk]:
+    return [
+        Chunk(
+            id=chunk.id,
+            relpath=chunk.relpath,
+            start_line=chunk.start_line,
+            end_line=chunk.end_line,
+            text=chunk.text,
+            file_hash=chunk.file_hash,
+            vector_idx=idx,
+        )
+        for idx, chunk in enumerate(chunks)
+    ]
+
+
+def _reuse_chunks(
+    prev_chunks: list[Chunk],
+    prev_embeddings: np.ndarray | None,
+) -> tuple[dict[str, list[Chunk]], dict[str, np.ndarray]]:
+    chunks_by_file: dict[str, list[Chunk]] = {}
+    embedding_by_id: dict[str, np.ndarray] = {}
+
+    if not prev_chunks:
+        return chunks_by_file, embedding_by_id
+
+    if prev_embeddings is not None:
+        for chunk in prev_chunks:
+            if chunk.vector_idx is None:
+                continue
+            if 0 <= chunk.vector_idx < len(prev_embeddings):
+                embedding_by_id[chunk.id] = prev_embeddings[chunk.vector_idx]
+
+    for chunk in prev_chunks:
+        chunks_by_file.setdefault(chunk.relpath, []).append(chunk)
+
+    return chunks_by_file, embedding_by_id
+
+
+def index_repo(repo_path: str, config: Config | None = None) -> IndexResult:
+    config = config or load_config()
+    repo_root = Path(repo_path).expanduser().resolve()
+
+    matcher = IgnoreMatcher(repo_root)
+    store = IndexStore(repo_root, config.index_root)
+
+    prev_meta = store.load_meta()
+    prev_chunks = store.load_chunks()
+    prev_embeddings = store.load_embeddings()
+    prev_files = prev_meta.files if prev_meta else {}
+
+    if prev_meta:
+        if (
+            prev_meta.chunk_lines != config.chunk_lines
+            or prev_meta.chunk_overlap != config.chunk_overlap
+            or prev_meta.max_bytes != config.max_bytes
+            or prev_meta.embedding_provider != config.embeddings.provider
+            or prev_meta.model != config.embeddings.model
+        ):
+            prev_meta = None
+            prev_chunks = []
+            prev_embeddings = None
+            prev_files = {}
+
+    prev_chunks_by_file, prev_embedding_by_id = _reuse_chunks(
+        prev_chunks,
+        prev_embeddings,
+    )
+
+    with Timer() as timer:
+        files = scan_repo(repo_root, matcher, config.max_bytes)
+        indexed_files: list[IndexedFile] = []
+        for path in files:
+            try:
+                indexed_files.append(_prepare_file(path, repo_root))
+            except OSError:
+                continue
+
+        file_hashes = {item.relpath: item.file_hash for item in indexed_files}
+
+        new_chunks: list[Chunk] = []
+        embedding_slots: list[np.ndarray | None] = []
+        pending_texts: list[str] = []
+        pending_indices: list[int] = []
+
+        for item in indexed_files:
+            relpath = item.relpath
+            unchanged = prev_files.get(relpath) == item.file_hash
+            if unchanged and relpath in prev_chunks_by_file and prev_embedding_by_id:
+                existing_chunks = prev_chunks_by_file[relpath]
+                if all(chunk.id in prev_embedding_by_id for chunk in existing_chunks):
+                    for chunk in existing_chunks:
+                        new_chunks.append(chunk)
+                        embedding_slots.append(prev_embedding_by_id[chunk.id])
+                    continue
+
+            chunks = _chunks_from_text(
+                relpath,
+                item.file_hash,
+                item.text,
+                config.chunk_lines,
+                config.chunk_overlap,
+            )
+            for chunk in chunks:
+                pending_indices.append(len(new_chunks))
+                new_chunks.append(chunk)
+                embedding_slots.append(None)
+                pending_texts.append(chunk.text)
+
+        provider = None
+        if pending_texts:
+            provider = get_embeddings_provider(config)
+            pending_embeddings = provider.embed_texts(pending_texts)
+        else:
+            pending_embeddings = np.zeros((0, 0), dtype=np.float32)
+
+        if provider is not None and provider.dimension:
+            dimension = provider.dimension
+        elif prev_embeddings is not None and prev_embeddings.size:
+            dimension = int(prev_embeddings.shape[1])
+        elif prev_meta and prev_meta.embedding_dim:
+            dimension = prev_meta.embedding_dim
+        else:
+            dimension = int(pending_embeddings.shape[1]) if pending_embeddings.size else 1
+
+        for slot_idx, embedding in zip(pending_indices, pending_embeddings):
+            embedding_slots[slot_idx] = embedding
+
+        if embedding_slots:
+            all_embeddings = np.stack(embedding_slots).astype(np.float32)
+        else:
+            all_embeddings = np.zeros((0, dimension), dtype=np.float32)
+        all_embeddings = normalize_embeddings(all_embeddings)
+
+        final_chunks = _with_vector_idx(new_chunks)
+
+        index = faiss.IndexFlatIP(dimension)
+        if all_embeddings.size:
+            index.add(all_embeddings)
+
+        now = datetime.now(timezone.utc).isoformat()
+        meta = IndexMeta(
+            repo_id=store.repo_id,
+            repo_path=str(repo_root),
+            created_at=prev_meta.created_at if prev_meta else now,
+            updated_at=now,
+            embedding_provider=config.embeddings.provider,
+            model=config.embeddings.model,
+            embedding_dim=dimension,
+            chunk_lines=config.chunk_lines,
+            chunk_overlap=config.chunk_overlap,
+            max_bytes=config.max_bytes,
+            files=file_hashes,
+            chunks=len(final_chunks),
+        )
+
+        store.save_chunks(final_chunks)
+        store.save_embeddings(all_embeddings)
+        store.save_faiss(index)
+        store.save_meta(meta)
+
+    return IndexResult(
+        repo_id=store.repo_id,
+        chunks_indexed=len(final_chunks),
+        took_ms=timer.ms,
+    )