codebase-retrieval-context-engine 2.0.0__py3-none-any.whl

corbell/core/embeddings/extractor.py

@@ -0,0 +1,201 @@
+"""Code chunk extractor for embedding indexing.
+
+Extracts code chunks (functions, classes, methods) from source files.
+Extracts function/class/method chunks using Python ast; generic line-split for others.
+"""
+
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import pathspec
+
+from corbell.core.constants import EXTENSION_LANG as _SUPPORTED, SKIP_DIRS as _SKIP_DIRS
+from corbell.core.gitignore import load_gitignore
+
+
+@dataclass
+class EmbeddingRecord:
+    """A chunk of code ready to be embedded."""
+
+    id: str
+    service_id: str
+    repo: str
+    file_path: str  # relative path within repo
+    start_line: int
+    end_line: int
+    content: str
+    language: str
+    chunk_type: str  # function | class | method | block
+    symbol: Optional[str] = None  # function/class name if known
+    embedding: Optional[List[float]] = None
+
+
+class CodeChunkExtractor:
+    """Extract meaningful code chunks from source files in a repo.
+
+    Produces :class:`EmbeddingRecord` instances that can be stored in any
+    embedding backend.
+    """
+
+    def __init__(self, chunk_size: int = 50, overlap: int = 10):
+        """Initialize the extractor.
+
+        Args:
+            chunk_size: Lines per generic block chunk.
+            overlap: Overlap between consecutive generic chunks.
+        """
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    def extract_from_repo(
+        self,
+        repo_path: Path | str,
+        service_id: str,
+        max_file_bytes: int = 1024 * 1024,
+        gitignore_spec: Optional[pathspec.PathSpec] = None,
+    ) -> List[EmbeddingRecord]:
+        """Walk a repo and extract all code chunks.
+
+        Args:
+            repo_path: Root directory of the repository.
+            service_id: ID of the owning service.
+            max_file_bytes: Skip files larger than this.
+            gitignore_spec: Pre-loaded PathSpec for gitignore filtering.
+                If None, it is loaded from the repo automatically.
+
+        Returns:
+            List of :class:`EmbeddingRecord` ready for embedding.
+        """
+        repo_path = Path(repo_path)
+        if gitignore_spec is None:
+            gitignore_spec = load_gitignore(repo_path)
+        records: List[EmbeddingRecord] = []
+
+        for fp in repo_path.rglob("*"):
+            if not fp.is_file():
+                continue
+            if self._should_skip(fp, max_file_bytes):
+                continue
+            lang = _SUPPORTED.get(fp.suffix)
+            if not lang:
+                continue
+            rel = str(fp.relative_to(repo_path))
+            if gitignore_spec.match_file(rel.replace("\\", "/")):
+                continue
+            chunks = self._extract_file(fp, rel, lang, service_id, str(repo_path))
+            records.extend(chunks)
+
+        return records
+
+    # ------------------------------------------------------------------ #
+    # Internal helpers                                                     #
+    # ------------------------------------------------------------------ #
+
+    def _should_skip(self, fp: Path, max_bytes: int) -> bool:
+        if any(part in _SKIP_DIRS for part in fp.parts):
+            return True
+        try:
+            if fp.stat().st_size > max_bytes:
+                return True
+        except OSError:
+            return True
+        return False
+
+    def _extract_file(
+        self, fp: Path, rel: str, lang: str, service_id: str, repo: str
+    ) -> List[EmbeddingRecord]:
+        try:
+            content = fp.read_text(encoding="utf-8", errors="ignore")
+        except Exception:
+            return []
+
+        if lang == "python":
+            return self._extract_python(content, rel, service_id, repo)
+        return self._extract_generic(content, rel, lang, service_id, repo)
+
+    def _extract_python(
+        self, content: str, rel: str, service_id: str, repo: str
+    ) -> List[EmbeddingRecord]:
+        """Use Python's ast to extract function/class definitions."""
+        records: List[EmbeddingRecord] = []
+        lines = content.splitlines()
+
+        try:
+            tree = ast.parse(content)
+        except SyntaxError:
+            return self._extract_generic(content, rel, "python", service_id, repo)
+
+        class _Visitor(ast.NodeVisitor):
+            def __init__(self_v):
+                self_v.class_stack: List[str] = []
+
+            def _emit(self_v, node, name: str, chunk_type: str):
+                line_start = node.lineno
+                line_end = getattr(node, "end_lineno", node.lineno)
+                chunk_content = "\n".join(lines[line_start - 1 : line_end])
+                symbol = ".".join(self_v.class_stack + [name])
+                rec = EmbeddingRecord(
+                    id=f"{service_id}::{rel}::{symbol}",
+                    service_id=service_id,
+                    repo=repo,
+                    file_path=rel,
+                    start_line=line_start,
+                    end_line=line_end,
+                    content=chunk_content,
+                    language="python",
+                    chunk_type=chunk_type,
+                    symbol=symbol,
+                )
+                records.append(rec)
+
+            def visit_ClassDef(self_v, node):
+                self_v._emit(node, node.name, "class")
+                self_v.class_stack.append(node.name)
+                self_v.generic_visit(node)
+                self_v.class_stack.pop()
+
+            def visit_FunctionDef(self_v, node):
+                chunk_type = "method" if self_v.class_stack else "function"
+                self_v._emit(node, node.name, chunk_type)
+
+            visit_AsyncFunctionDef = visit_FunctionDef
+
+        _Visitor().visit(tree)
+        # Fall back to generic if nothing found
+        if not records:
+            return self._extract_generic(content, rel, "python", service_id, repo)
+        return records
+
+    def _extract_generic(
+        self, content: str, rel: str, lang: str, service_id: str, repo: str
+    ) -> List[EmbeddingRecord]:
+        """Split file into overlapping line-based blocks."""
+        lines = content.splitlines()
+        records: List[EmbeddingRecord] = []
+        step = max(1, self.chunk_size - self.overlap)
+
+        for i in range(0, len(lines), step):
+            end = min(i + self.chunk_size, len(lines))
+            chunk_lines = lines[i:end]
+            if not any(l.strip() for l in chunk_lines):
+                continue
+            chunk_content = "\n".join(chunk_lines)
+            records.append(
+                EmbeddingRecord(
+                    id=f"{service_id}::{rel}::block_{i}",
+                    service_id=service_id,
+                    repo=repo,
+                    file_path=rel,
+                    start_line=i + 1,
+                    end_line=end,
+                    content=chunk_content,
+                    language=lang,
+                    chunk_type="block",
+                )
+            )
+
+        return records

corbell/core/embeddings/factory.py

@@ -0,0 +1,48 @@
+"""Factory for creating EmbeddingStore instances by backend name.
+
+To add a new backend:
+    1. Create a class that implements :class:`~corbell.core.embeddings.base.EmbeddingStore`.
+    2. Add an ``elif backend == "<name>":`` branch below.
+    3. Users opt in via ``storage.embeddings.backend: <name>`` in workspace.yaml.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from corbell.core.embeddings.base import EmbeddingStore
+
+_SUPPORTED_BACKENDS = ("sqlite",)
+
+
+def get_embedding_store(backend: str, db_path: Path) -> EmbeddingStore:
+    """Return an :class:`EmbeddingStore` for the requested backend.
+
+    Args:
+        backend: Backend identifier string (e.g. ``"sqlite"``).
+        db_path: Path to the storage file / directory.
+
+    Returns:
+        A concrete :class:`EmbeddingStore` instance.
+
+    Raises:
+        ValueError: If ``backend`` is not a recognised backend name.
+    """
+    backend = backend.lower().strip()
+
+    if backend == "sqlite":
+        from corbell.core.embeddings.sqlite_store import SQLiteEmbeddingStore
+        return SQLiteEmbeddingStore(db_path)
+
+    # ------------------------------------------------------------------ #
+    # Future backends — add branches here, e.g.:                          #
+    #   elif backend == "kuzu":                                            #
+    #       from corbell.core.embeddings.kuzu_store import KuzuEmbeddingStore
+    #       return KuzuEmbeddingStore(db_path)                            #
+    # ------------------------------------------------------------------ #
+
+    raise ValueError(
+        f"Unknown embedding backend: {backend!r}. "
+        f"Supported backends: {', '.join(_SUPPORTED_BACKENDS)}. "
+        f"Set 'storage.embeddings.backend' in workspace.yaml."
+    )

corbell/core/embeddings/model.py

@@ -0,0 +1,401 @@
+"""Embedding model interface + SentenceTransformers implementation."""
+
+from __future__ import annotations
+
+import logging
+import os
+import random
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingModel(ABC):
+    """Abstract embedding model interface."""
+
+    @abstractmethod
+    def encode(self, texts: List[str]) -> List[List[float]]:
+        """Encode a list of texts into embedding vectors.
+
+        Args:
+            texts: List of text strings to encode.
+
+        Returns:
+            List of float vectors (one per input text).
+        """
+        ...
+
+    @property
+    @abstractmethod
+    def dimension(self) -> int:
+        """Return the embedding dimension."""
+        ...
+
+
+class SentenceTransformerModel(EmbeddingModel):
+    """Wraps ``sentence-transformers`` with lazy loading.
+
+    Uses ``all-MiniLM-L6-v2`` by default (384-dim, fast, no API key).
+    """
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        self.model_name = model_name
+        self._model = None  # lazy-loaded
+
+    def _get_model(self):
+        if self._model is None:
+            from sentence_transformers import SentenceTransformer
+            self._model = SentenceTransformer(f"sentence-transformers/{self.model_name}")
+        return self._model
+
+    def encode(self, texts: List[str]) -> List[List[float]]:
+        model = self._get_model()
+        vecs = model.encode(texts, show_progress_bar=False)
+        return [v.tolist() for v in vecs]
+
+    @property
+    def dimension(self) -> int:
+        return self._get_model().get_sentence_embedding_dimension()
+
+
+def _is_voyage_rate_limit_error(e: Exception) -> bool:
+    """Return True when a Voyage API error is a 429 rate limit."""
+    status = getattr(e, "status_code", None)
+    if status == 429:
+        return True
+    # Some Voyage SDK versions use a different attribute
+    code = getattr(e, "code", None)
+    return code == 429
+
+
+def _is_google_key_error(e: Exception) -> bool:
+    """Return True when a Google API error is caused by the key, not the request."""
+    code = getattr(e, "code", None)
+    if code in (401, 403, 429):
+        return True
+    if code == 400:
+        msg = (getattr(e, "message", None) or str(e)).lower()
+        return "api key" in msg
+    return False
+
+
+def _is_rate_limit_error(e: Exception) -> bool:
+    """Return True when a Google API error is a 429 RESOURCE_EXHAUSTED rate limit."""
+    return getattr(e, "code", None) == 429
+
+
+def _parse_gemini_version(model_name: str) -> int:
+    """Parse the version number from a gemini-embedding model name.
+
+    Examples:
+        ``gemini-embedding-001`` → 1
+        ``gemini-embedding-2`` → 2
+
+    Returns 0 if the version cannot be parsed.
+    """
+    prefix = "gemini-embedding-"
+    if not model_name.startswith(prefix):
+        return 0
+    suffix = model_name[len(prefix):]
+    try:
+        return int(suffix)
+    except ValueError:
+        return 0
+
+
+class GoogleEmbeddingModel(EmbeddingModel):
+    """Google AI (Gemini) embedding model via the google-genai SDK.
+
+    Uses ``gemini-embedding-001`` by default (768-dim, text-only).
+    Requires ``pip install corbell[google]`` and ``GOOGLE_API_KEY``.
+
+    Supports a comma-separated list of API keys for round-robin distribution
+    and automatic failover when a key is invalid or quota-exhausted.
+
+    Supports ``task_type`` to improve retrieval quality:
+    - ``RETRIEVAL_DOCUMENT`` for indexing (default)
+    - ``RETRIEVAL_QUERY`` for query-time encoding
+
+    For ``gemini-embedding-2`` and later, inline text prefixes are used
+    instead of relying solely on ``task_type`` for better retrieval quality.
+    Use ``prepare_query`` and ``prepare_document`` to format texts before
+    passing them to ``encode``.
+    """
+
+    def __init__(self, model_name: str = "gemini-embedding-001", api_key: Optional[str] = None):
+        self.model_name = model_name
+        raw = api_key or os.environ.get("GOOGLE_API_KEY") or ""
+        self._api_keys: List[str] = [k.strip() for k in raw.split(",") if k.strip()]
+        if not self._api_keys:
+            raise ValueError(
+                "GOOGLE_API_KEY is not set. "
+                "Set it in your environment or workspace.yaml:\n"
+                "  export GOOGLE_API_KEY=AIza...\n"
+                "Or use a local embedding model (e.g. all-MiniLM-L6-v2) in storage.model."
+            )
+        self._key_index: int = random.randrange(len(self._api_keys))
+        # kept for backwards-compat with tests that read _api_key directly
+        self._api_key: str = self._api_keys[0]
+
+    @property
+    def uses_prefix_format(self) -> bool:
+        """Return True when the model requires inline text prefixes for best quality.
+
+        Activated for ``gemini-embedding-2`` and all later versions (version >= 2).
+        """
+        return _parse_gemini_version(self.model_name) >= 2
+
+    def prepare_query(self, query: str) -> str:
+        """Format a query string with a task prefix for retrieval.
+
+        Only applies the prefix when ``uses_prefix_format`` is True.
+        """
+        if self.uses_prefix_format:
+            return f"task: code retrieval | query: {query}"
+        return query
+
+    def prepare_document(self, content: str, title: Optional[str] = None) -> str:
+        """Format a document chunk with a title prefix for indexing.
+
+        Only applies the prefix when ``uses_prefix_format`` is True.
+
+        Args:
+            content: Raw chunk text.
+            title: Descriptive title, typically ``"{file_path}:{symbol}"``
+                   or ``"{file_path}:L{start}-{end}"``. Defaults to ``"none"``.
+        """
+        if self.uses_prefix_format:
+            resolved_title = title or "none"
+            return f"title: {resolved_title} | text: {content}"
+        return content
+
+    _BATCH_SIZE = 100
+    _BASE_DELAY = 2.0
+    _MAX_BACKOFF = 60.0
+
+    def encode(self, texts: List[str], task_type: str = "RETRIEVAL_DOCUMENT") -> List[List[float]]:
+        """Encode a list of texts into embedding vectors.
+
+        Batches requests (100 texts/batch) and retries on rate limit (429)
+        with exponential backoff.
+
+        Args:
+            texts: List of text strings to encode.
+            task_type: Task type hint for the embedding model.
+                Use ``RETRIEVAL_DOCUMENT`` when indexing, ``RETRIEVAL_QUERY`` at query time.
+
+        Returns:
+            List of float vectors (one per input text).
+        """
+        try:
+            from google import genai
+            from google.genai import types
+        except ImportError:
+            raise ImportError("pip install corbell[google]")
+
+        all_embeddings: List[List[float]] = []
+        for batch_start in range(0, len(texts), self._BATCH_SIZE):
+            batch = texts[batch_start:batch_start + self._BATCH_SIZE]
+            contents = [types.Content(parts=[types.Part(text=t)]) for t in batch]
+            batch_result = self._embed_batch_with_retry(
+                contents, task_type, genai, types
+            )
+            all_embeddings.extend(batch_result)
+
+        return all_embeddings
+
+    def _embed_batch_with_retry(
+        self, contents, task_type: str, genai, types
+    ) -> List[List[float]]:
+        """Embed a single batch, rotating keys and retrying on rate limit.
+
+        - If all keys fail with 429 (rate limit), waits with capped exponential
+          backoff and retries indefinitely until the quota is restored.
+        - If any key fails with a non-key error, raises immediately.
+        - If all keys fail with auth errors (401/403/400+apikey), raises immediately.
+        """
+        import time
+
+        start = self._key_index
+        rate_limit_attempt = 0
+
+        while True:
+            errors: List[str] = []
+            all_rate_limited = True
+
+            for i in range(len(self._api_keys)):
+                idx = (start + i) % len(self._api_keys)
+                key = self._api_keys[idx]
+                try:
+                    client = genai.Client(api_key=key)
+                    result = client.models.embed_content(
+                        model=self.model_name,
+                        contents=contents,
+                        config=types.EmbedContentConfig(
+                            task_type=task_type,
+                            output_dimensionality=self.dimension,
+                        ),
+                    )
+                    self._key_index = (idx + 1) % len(self._api_keys)
+                    return [emb.values for emb in result.embeddings]
+                except Exception as e:
+                    if _is_google_key_error(e):
+                        if not _is_rate_limit_error(e):
+                            # Auth failure (401/403/400+apikey) — not a transient error
+                            all_rate_limited = False
+                        errors.append(f"key[{idx}]: {e}")
+                        continue
+                    raise
+
+            if not all_rate_limited:
+                raise RuntimeError(
+                    f"All {len(self._api_keys)} Google API key(s) failed with auth errors:\n"
+                    + "\n".join(errors)
+                )
+
+            # All keys are rate-limited (429) — wait and retry indefinitely
+            delay = min(self._BASE_DELAY * (2 ** rate_limit_attempt), self._MAX_BACKOFF)
+            rate_limit_attempt += 1
+            logger.warning(
+                "All %d Google API key(s) rate-limited (429). "
+                "Retrying in %.0fs (attempt %d)...",
+                len(self._api_keys),
+                delay,
+                rate_limit_attempt,
+            )
+            time.sleep(delay)
+
+    @property
+    def dimension(self) -> int:
+        dim_env = os.environ.get("CORBELL_EMBEDDING_DIM", "").strip()
+        if dim_env:
+            return int(dim_env)
+        return 768
+
+
+class VoyageEmbeddingModel(EmbeddingModel):
+    """Voyage AI embedding model via the voyageai SDK.
+
+    Uses ``voyage-code-3`` by default (1024-dim, optimized for code retrieval).
+    Requires ``pip install corbell[voyage]`` and ``VOYAGE_API_KEY``.
+
+    Supports a comma-separated list of API keys for round-robin distribution
+    and automatic failover when a key is quota-exhausted.
+
+    The ``input_type`` parameter improves retrieval quality:
+    - ``"document"`` for indexing (default)
+    - ``"query"`` for query-time encoding
+
+    Use ``prepare_query`` and ``prepare_document`` which return the text unchanged
+    (Voyage handles task differentiation via ``input_type``).
+    """
+
+    _BATCH_SIZE = 1000
+    _BASE_DELAY = 2.0
+    _MAX_BACKOFF = 60.0
+
+    def __init__(self, model_name: str = "voyage-code-3", api_key: Optional[str] = None):
+        self.model_name = model_name
+        raw = api_key or os.environ.get("VOYAGE_API_KEY") or ""
+        self._api_keys: List[str] = [k.strip() for k in raw.split(",") if k.strip()]
+        if not self._api_keys:
+            raise ValueError(
+                "VOYAGE_API_KEY is not set. "
+                "Set it in your environment or workspace.yaml:\n"
+                "  export VOYAGE_API_KEY=pa-...\n"
+                "Or use a local embedding model (e.g. all-MiniLM-L6-v2) in storage.model."
+            )
+        self._key_index: int = random.randrange(len(self._api_keys))
+        # kept for backwards-compat with tests that read _api_key directly
+        self._api_key: str = self._api_keys[0]
+
+    def prepare_query(self, query: str) -> str:
+        """Return query unchanged; Voyage handles task differentiation via input_type."""
+        return query
+
+    def prepare_document(self, content: str, title: Optional[str] = None) -> str:
+        """Return content unchanged; Voyage handles task differentiation via input_type."""
+        return content
+
+    def encode(self, texts: List[str], input_type: str = "document") -> List[List[float]]:
+        """Encode a list of texts into embedding vectors.
+
+        Batches requests (1000 texts/batch) and retries on rate limit (429)
+        with exponential backoff.
+
+        Args:
+            texts: List of text strings to encode.
+            input_type: Voyage input type hint. Use ``"document"`` when indexing,
+                ``"query"`` at query time.
+
+        Returns:
+            List of float vectors (one per input text).
+        """
+        try:
+            import voyageai
+        except ImportError:
+            raise ImportError("pip install corbell[voyage]")
+
+        all_embeddings: List[List[float]] = []
+        for batch_start in range(0, len(texts), self._BATCH_SIZE):
+            batch = texts[batch_start:batch_start + self._BATCH_SIZE]
+            batch_result = self._embed_batch_with_retry(batch, input_type, voyageai)
+            all_embeddings.extend(batch_result)
+
+        return all_embeddings
+
+    def _embed_batch_with_retry(
+        self, batch: List[str], input_type: str, voyageai
+    ) -> List[List[float]]:
+        """Embed a single batch, rotating keys and retrying on rate limit.
+
+        - If all keys fail with 429 (rate limit), waits with capped exponential
+          backoff and retries indefinitely until the quota is restored.
+        - If any key fails with a non-rate-limit error, raises immediately.
+        """
+        import time
+
+        start = self._key_index
+        rate_limit_attempt = 0
+
+        while True:
+            errors: List[str] = []
+
+            for i in range(len(self._api_keys)):
+                idx = (start + i) % len(self._api_keys)
+                key = self._api_keys[idx]
+                try:
+                    vo = voyageai.Client(api_key=key)
+                    result = vo.embed(
+                        batch,
+                        model=self.model_name,
+                        input_type=input_type,
+                        output_dimension=self.dimension,
+                    )
+                    self._key_index = (idx + 1) % len(self._api_keys)
+                    return result.embeddings
+                except Exception as e:
+                    if _is_voyage_rate_limit_error(e):
+                        errors.append(f"key[{idx}]: {e}")
+                        continue
+                    raise
+
+            # All keys are rate-limited (429) — wait and retry indefinitely
+            delay = min(self._BASE_DELAY * (2 ** rate_limit_attempt), self._MAX_BACKOFF)
+            rate_limit_attempt += 1
+            logger.warning(
+                "All %d Voyage API key(s) rate-limited (429). "
+                "Retrying in %.0fs (attempt %d)...",
+                len(self._api_keys),
+                delay,
+                rate_limit_attempt,
+            )
+            time.sleep(delay)
+
+    @property
+    def dimension(self) -> int:
+        dim_env = os.environ.get("CORBELL_EMBEDDING_DIM", "").strip()
+        if dim_env:
+            return int(dim_env)
+        return 1024