@scemoon/cdh 1.0.8 → 1.0.9-beta.2

package/onecode/codebase/__init__.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Optional
+
+from onecode.codebase.chunker import CodeChunk
+from onecode.config import CodebaseConfig
+from onecode.codebase.indexer import CodebaseIndexer
+from onecode.codebase.retriever import CodebaseRetriever
+from onecode.codebase.storage import CodebaseStorage
+from onecode.codebase.tools import CodebaseSearchTool
+
+logger = logging.getLogger(__name__)
+
+_MAX_INJECT_TOKENS = 2000
+
+
+class CodebaseEngine:
+    def __init__(self, project_dir: Path, config: Optional[CodebaseConfig] = None):
+        self.project_dir = project_dir.resolve()
+        self.config = config or CodebaseConfig()
+        self.storage = CodebaseStorage(project_dir)
+        self.indexer = CodebaseIndexer(project_dir, self.config)
+        self.retriever = CodebaseRetriever(project_dir, self.config)
+        self._indexed = False
+
+    async def ensure_indexed(self, force: bool = False) -> None:
+        if self._indexed and not force:
+            return
+        result = await self.indexer.index(force=force)
+        self._indexed = True
+        self.retriever.mark_dirty()
+        return result
+
+    async def retrieve(self, query: str, top_k: Optional[int] = None) -> list[CodeChunk]:
+        await self.ensure_indexed()
+        return await self.retriever.retrieve(query, top_k=top_k)
+
+    def format_context(self, chunks: list[CodeChunk], max_tokens: int = _MAX_INJECT_TOKENS) -> str:
+        if not chunks:
+            return ""
+        parts: list[str] = []
+        token_budget = max_tokens
+        for c in chunks:
+            header = f"<chunk file=\"{c.file_path}\" lines=\"{c.start_line}-{c.end_line}\">"
+            footer = "</chunk>"
+            chunk_text = f"{header}\n{c.content}\n{footer}"
+            estimated_tokens = len(chunk_text) // 4 + 1
+            if estimated_tokens > token_budget:
+                if not parts:
+                    trimmed = c.content[:token_budget * 4]
+                    chunk_text = f"{header}\n{trimmed}\n{footer}"
+                    parts.append(chunk_text)
+                break
+            parts.append(chunk_text)
+            token_budget -= estimated_tokens
+        return "<codebase_context>\n" + "\n".join(parts) + "\n</codebase_context>"
+
+
+__all__ = [
+    "CodebaseEngine",
+    "CodebaseConfig",
+    "CodeChunk",
+    "CodebaseIndexer",
+    "CodebaseRetriever",
+    "CodebaseStorage",
+    "CodebaseSearchTool",
+]

package/onecode/codebase/chunker.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+
+_LANG_MAP: dict[str, str] = {
+    ".py": "python",
+    ".js": "javascript",
+    ".ts": "typescript",
+    ".tsx": "typescriptreact",
+    ".jsx": "javascriptreact",
+    ".go": "go",
+    ".rs": "rust",
+    ".java": "java",
+    ".rb": "ruby",
+    ".php": "php",
+    ".c": "c",
+    ".cpp": "cpp",
+    ".h": "c",
+    ".hpp": "cpp",
+    ".swift": "swift",
+    ".kt": "kotlin",
+    ".scala": "scala",
+    ".sql": "sql",
+    ".yaml": "yaml",
+    ".yml": "yaml",
+    ".json": "json",
+    ".toml": "toml",
+    ".md": "markdown",
+    ".css": "css",
+    ".scss": "scss",
+    ".html": "html",
+    ".svelte": "svelte",
+    ".vue": "vue",
+    ".sh": "shell",
+    ".bash": "shell",
+    ".zsh": "shell",
+    ".fish": "shell",
+    ".tf": "terraform",
+    ".tfvars": "terraform",
+}
+
+
+@dataclass
+class CodeChunk:
+    file_path: str
+    start_line: int
+    end_line: int
+    content: str
+    language: str = ""
+
+    @property
+    def identifier(self) -> str:
+        return f"{self.file_path}:{self.start_line}-{self.end_line}"
+
+
+def detect_language(file_path: str) -> str:
+    ext = os.path.splitext(file_path)[1].lower()
+    return _LANG_MAP.get(ext, "")
+
+
+def chunk_file(file_path: Path, strategy: str = "line",
+               chunk_lines: int = 50, overlap: int = 10) -> list[CodeChunk]:
+    return _chunk_by_lines(file_path, chunk_lines, overlap)
+
+
+def _chunk_by_lines(file_path: Path, chunk_lines: int, overlap: int) -> list[CodeChunk]:
+    try:
+        text = file_path.read_text(encoding="utf-8")
+    except Exception:
+        return []
+
+    lines = text.splitlines(keepends=True)
+    if not lines:
+        return []
+
+    language = detect_language(str(file_path))
+    step = max(chunk_lines - overlap, 1)
+    chunks: list[CodeChunk] = []
+
+    for start in range(0, len(lines), step):
+        end = min(start + chunk_lines, len(lines))
+        content = "".join(lines[start:end]).rstrip()
+        if not content.strip():
+            continue
+        chunks.append(CodeChunk(
+            file_path=str(file_path),
+            start_line=start + 1,
+            end_line=end,
+            content=content,
+            language=language,
+        ))
+        if end == len(lines):
+            break
+
+    return chunks
+
+
+def chunk_text(text: str, file_path: str, chunk_lines: int = 50, overlap: int = 10) -> list[CodeChunk]:
+    lines = text.splitlines(keepends=True)
+    if not lines:
+        return []
+
+    language = detect_language(file_path)
+    step = max(chunk_lines - overlap, 1)
+    chunks: list[CodeChunk] = []
+
+    for start in range(0, len(lines), step):
+        end = min(start + chunk_lines, len(lines))
+        content = "".join(lines[start:end]).rstrip()
+        if not content.strip():
+            continue
+        chunks.append(CodeChunk(
+            file_path=file_path,
+            start_line=start + 1,
+            end_line=end,
+            content=content,
+            language=language,
+        ))
+        if end == len(lines):
+            break
+
+    return chunks

package/onecode/codebase/indexer.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+from pathspec import PathSpec
+from pathspec.patterns import GitWildMatchPattern
+
+from onecode.codebase.chunker import chunk_file, CodeChunk
+from onecode.config import CodebaseConfig
+from onecode.codebase.storage import CodebaseStorage
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class IndexResult:
+    total_files: int = 0
+    indexed_files: int = 0
+    skipped_files: int = 0
+    failed_files: int = 0
+    total_chunks: int = 0
+    duration_ms: float = 0.0
+    errors: list[str] = field(default_factory=list)
+
+
+class CodebaseIndexer:
+    def __init__(self, project_dir: Path, config: CodebaseConfig):
+        self.project_dir = project_dir.resolve()
+        self.config = config
+        self.storage = CodebaseStorage(project_dir)
+        self._spec = self._build_spec()
+
+    def _build_spec(self) -> Optional[PathSpec]:
+        patterns = self.config.exclude_patterns
+        if not patterns:
+            return None
+        return PathSpec.from_lines(GitWildMatchPattern, patterns)
+
+    def walk_files(self) -> list[Path]:
+        include_exts = {ext.lower() for ext in self.config.include_extensions}
+        files: list[Path] = []
+        for root, dirs, filenames in os.walk(self.project_dir):
+            root_path = Path(root)
+            rel_root = root_path.relative_to(self.project_dir)
+            dirs[:] = [
+                d for d in dirs
+                if not self._is_excluded(str(rel_root / d))
+            ]
+            for fname in filenames:
+                fpath = root_path / fname
+                rel = fpath.relative_to(self.project_dir)
+                if self._is_excluded(str(rel)):
+                    continue
+                if fpath.suffix.lower() in include_exts:
+                    files.append(fpath)
+        return files
+
+    def _is_excluded(self, rel_path: str) -> bool:
+        if self._spec is None:
+            return False
+        return self._spec.match_file(rel_path)
+
+    def needs_reindex(self, file_path: Path) -> bool:
+        if not file_path.exists():
+            return True
+        stored_mtime = self.storage.get_file_mtime(str(file_path))
+        actual_mtime = file_path.stat().st_mtime
+        return actual_mtime > stored_mtime
+
+    async def index(self, force: bool = False) -> IndexResult:
+        result = IndexResult()
+        start = time.monotonic()
+
+        if force:
+            self.storage.clear_all()
+
+        files = self.walk_files()
+        result.total_files = len(files)
+
+        _YIELD_EVERY = 16
+        _BATCH_SIZE = 32
+
+        pending_chunks: list[tuple[list[CodeChunk], float]] = []
+
+        for i, fpath in enumerate(files):
+            if i and (i % _YIELD_EVERY) == 0:
+                await asyncio.sleep(0)
+
+            if not force and not self.needs_reindex(fpath):
+                result.skipped_files += 1
+                continue
+
+            try:
+                mtime = fpath.stat().st_mtime
+                chunks = await asyncio.to_thread(
+                    chunk_file,
+                    fpath,
+                    strategy=self.config.chunk_strategy,
+                    chunk_lines=self.config.chunk_lines,
+                    overlap=self.config.chunk_overlap,
+                )
+                if chunks:
+                    pending_chunks.append((chunks, mtime))
+                    result.indexed_files += 1
+                    result.total_chunks += len(chunks)
+
+                if len(pending_chunks) >= _BATCH_SIZE:
+                    self.storage.save_chunks_batch(pending_chunks)
+                    pending_chunks.clear()
+            except Exception as e:
+                logger.warning("Failed to index %s: %s", fpath, e)
+                result.failed_files += 1
+                result.errors.append(f"{fpath}: {e}")
+
+        if pending_chunks:
+            self.storage.save_chunks_batch(pending_chunks)
+
+        result.duration_ms = (time.monotonic() - start) * 1000
+        logger.info(
+            "Indexed %d files (%d skipped, %d failed) → %d chunks in %.0fms",
+            result.indexed_files, result.skipped_files, result.failed_files,
+            result.total_chunks, result.duration_ms,
+        )
+        return result
+
+    async def index_file(self, file_path: Path) -> None:
+        if not file_path.exists():
+            self.storage.remove_file(str(file_path))
+            return
+        include_exts = {ext.lower() for ext in self.config.include_extensions}
+        if file_path.suffix.lower() not in include_exts:
+            self.storage.remove_file(str(file_path))
+            return
+        try:
+            mtime = file_path.stat().st_mtime
+            chunks = await asyncio.to_thread(
+                chunk_file,
+                file_path,
+                strategy=self.config.chunk_strategy,
+                chunk_lines=self.config.chunk_lines,
+                overlap=self.config.chunk_overlap,
+            )
+            if chunks:
+                self.storage.save_chunks(chunks, mtime=mtime)
+            else:
+                self.storage.remove_file(str(file_path))
+        except Exception as e:
+            logger.warning("Failed to index %s: %s", file_path, e)

package/onecode/codebase/retriever.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import math
+from typing import Optional
+
+import httpx
+
+from onecode.codebase.chunker import CodeChunk
+from onecode.config import CodebaseConfig
+from onecode.codebase.storage import CodebaseStorage
+from onecode.memory.recall import BM25
+
+logger = logging.getLogger(__name__)
+
+
+def _cosine_similarity(a: list[float], b: list[float]) -> float:
+    dot = sum(x * y for x, y in zip(a, b))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(y * y for y in b))
+    if norm_a == 0 or norm_b == 0:
+        return 0.0
+    return dot / (norm_a * norm_b)
+
+
+class CodebaseRetriever:
+    def __init__(self, project_dir, config: CodebaseConfig):
+        self.storage = CodebaseStorage(project_dir)
+        self.config = config
+        self._bm25: Optional[BM25] = None
+        self._chunks: list[CodeChunk] = []
+        self._dirty = True
+        self._chunks_by_file: dict[str, list[CodeChunk]] = {}
+        self._file_chunk_counts: dict[str, int] = {}
+
+    def _ensure_indexed(self) -> None:
+        if not self._dirty:
+            return
+
+        if not self._chunks_by_file:
+            self._chunks = self.storage.get_all_chunks()
+            self._chunks_by_file = {}
+            for chunk in self._chunks:
+                if chunk.file_path not in self._chunks_by_file:
+                    self._chunks_by_file[chunk.file_path] = []
+                self._chunks_by_file[chunk.file_path].append(chunk)
+            for fp, chunks in self._chunks_by_file.items():
+                self._file_chunk_counts[fp] = len(chunks)
+        else:
+            self._chunks = []
+            for chunks in self._chunks_by_file.values():
+                self._chunks.extend(chunks)
+
+        if not self._chunks:
+            self._bm25 = None
+            self._dirty = False
+            logger.debug("BM25 index empty")
+            return
+
+        texts = [c.content for c in self._chunks]
+        self._bm25 = BM25()
+        self._bm25.index(texts)
+
+        self._dirty = False
+        logger.debug("BM25 index rebuilt with %d chunks from %d files", len(self._chunks), len(self._chunks_by_file))
+
+    def mark_dirty(self) -> None:
+        self._dirty = True
+
+    def update_file(self, file_path: str, chunks: list[CodeChunk]) -> None:
+        self._chunks_by_file[file_path] = chunks
+        self._file_chunk_counts[file_path] = len(chunks)
+        self._dirty = True
+
+    def remove_file(self, file_path: str) -> None:
+        if file_path in self._chunks_by_file:
+            del self._chunks_by_file[file_path]
+            del self._file_chunk_counts[file_path]
+            self._dirty = True
+
+    async def retrieve(self, query: str, top_k: Optional[int] = None) -> list[CodeChunk]:
+        k = top_k or self.config.top_k
+        if not query.strip():
+            return []
+        retriever_type = self.config.retriever
+        if retriever_type == "bm25":
+            return self._retrieve_bm25(query, k)
+        elif retriever_type == "embedding":
+            return await self._retrieve_embedding(query, k) or self._retrieve_bm25(query, k)
+        elif retriever_type == "hybrid":
+            return await self._retrieve_hybrid(query, k)
+        return self._retrieve_bm25(query, k)
+
+    def _retrieve_bm25(self, query: str, top_k: int) -> list[CodeChunk]:
+        self._ensure_indexed()
+        if not self._bm25 or not self._chunks:
+            return []
+        results = self._bm25.search(query, top_k=top_k)
+        scored: list[tuple[CodeChunk, float]] = []
+        for idx, score in results:
+            if 0 <= idx < len(self._chunks):
+                scored.append((self._chunks[idx], score))
+        scored.sort(key=lambda x: x[1], reverse=True)
+        return [c for c, _ in scored[:top_k]]
+
+    async def _retrieve_embedding(self, query: str, top_k: int) -> Optional[list[CodeChunk]]:
+        self._ensure_indexed()
+        if not self._chunks:
+            return None
+
+        query_emb = await _get_embedding(query, self.config)
+        if query_emb is None:
+            return None
+
+        embs = await _get_batch_embeddings([c.content[:1000] for c in self._chunks], self.config)
+        if not embs or len(embs) != len(self._chunks):
+            return None
+
+        scored = []
+        for i, emb in enumerate(embs):
+            if emb is not None:
+                score = _cosine_similarity(query_emb, emb)
+                scored.append((self._chunks[i], score))
+
+        scored.sort(key=lambda x: x[1], reverse=True)
+        return [c for c, _ in scored[:top_k]]
+
+    async def _retrieve_hybrid(self, query: str, top_k: int) -> list[CodeChunk]:
+        bm25_chunks = self._retrieve_bm25(query, top_k * 2)
+        emb_chunks = await self._retrieve_embedding(query, top_k * 2) or []
+        seen: set[str] = set()
+        merged: list[CodeChunk] = []
+        for chunk in bm25_chunks + emb_chunks:
+            key = chunk.identifier
+            if key not in seen:
+                seen.add(key)
+                merged.append(chunk)
+        return merged[:top_k]
+
+
+EMBEDDING_CACHE: dict[str, list[float]] = {}
+BATCH_CACHE: dict[int, Optional[list[float]]] = {}
+
+
+async def _get_embedding(text: str, config: CodebaseConfig) -> Optional[list[float]]:
+    if text in EMBEDDING_CACHE:
+        return EMBEDDING_CACHE[text]
+    try:
+        from onecode.config import load_config
+        cfg = load_config()
+        provider_name = config.embedding_provider or cfg.default_provider
+        pcfg = cfg.providers.get(provider_name)
+        if not pcfg:
+            return None
+        if provider_name == "openai":
+            emb = await _openai_embed(text, pcfg.api_key or "", pcfg.endpoint or "")
+        elif provider_name == "ollama":
+            emb = await _ollama_embed(text, pcfg.endpoint or "http://localhost:11434")
+        else:
+            emb = await _openai_embed(text, pcfg.api_key or "", pcfg.endpoint or "")
+        if emb is not None:
+            EMBEDDING_CACHE[text] = emb
+        return emb
+    except Exception as e:
+        logger.warning("Embedding failed: %s", e)
+        return None
+
+
+async def _get_batch_embeddings(texts: list[str], config: CodebaseConfig) -> list[Optional[list[float]]]:
+    global BATCH_CACHE
+    key = hash(tuple(texts))
+    if key in BATCH_CACHE:
+        return BATCH_CACHE[key]
+
+    from onecode.config import load_config
+    cfg = load_config()
+    provider_name = config.embedding_provider or cfg.default_provider
+    pcfg = cfg.providers.get(provider_name)
+
+    MAX_EMBED_CHUNKS = 500
+    if len(texts) > MAX_EMBED_CHUNKS:
+        import logging
+        logging.getLogger(__name__).warning(
+            "Embedding requested on %d chunks (max=%d); truncating to avoid timeout",
+            len(texts), MAX_EMBED_CHUNKS,
+        )
+        texts = texts[:MAX_EMBED_CHUNKS]
+
+    _MAX_CONCURRENCY = 16
+    semaphore = asyncio.Semaphore(_MAX_CONCURRENCY)
+
+    async def embed_with_sem(text: str) -> Optional[list[float]]:
+        async with semaphore:
+            return await _get_embedding(text, config)
+
+    if provider_name == "ollama" and pcfg:
+        embs = await _ollama_batch_embed(texts, pcfg.endpoint or "http://localhost:11434")
+    else:
+        embs = await asyncio.gather(*[embed_with_sem(t) for t in texts])
+
+    BATCH_CACHE[key] = list(embs)
+    return list(embs)
+
+
+async def _openai_embed(text: str, api_key: str, endpoint: str) -> Optional[list[float]]:
+    url = f"{endpoint.rstrip('/')}/embeddings"
+    async with httpx.AsyncClient(timeout=30) as client:
+        resp = await client.post(
+            url,
+            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
+            json={"input": text, "model": "text-embedding-3-small"},
+        )
+        data = resp.json()
+        emb = data.get("data", [{}])[0].get("embedding")
+        return emb if emb else None
+
+
+async def _ollama_embed(text: str, endpoint: str) -> Optional[list[float]]:
+    url = f"{endpoint.rstrip('/')}/api/embeddings"
+    async with httpx.AsyncClient(timeout=30) as client:
+        resp = await client.post(
+            url,
+            json={"model": "nomic-embed-text", "prompt": text},
+        )
+        data = resp.json()
+        emb = data.get("embedding")
+        return emb if emb else None
+
+
+async def _ollama_batch_embed(texts: list[str], endpoint: str) -> list[Optional[list[float]]]:
+    _MAX_CONCURRENCY = 16
+    semaphore = asyncio.Semaphore(_MAX_CONCURRENCY)
+
+    async def embed_with_sem(text: str) -> Optional[list[float]]:
+        async with semaphore:
+            return await _ollama_embed(text, endpoint)
+
+    embs = await asyncio.gather(*[embed_with_sem(t) for t in texts])
+    return list(embs)