vexor 0.21.1__py3-none-any.whl → 0.23.0rc1__py3-none-any.whl

vexor/cache.py

@@ -5,9 +5,13 @@ from __future__ import annotations
 import hashlib
 import os
 import sqlite3
+from collections import OrderedDict
 from dataclasses import dataclass
+from contextlib import contextmanager
+from contextvars import ContextVar
 from datetime import datetime, timezone, timedelta
 from pathlib import Path
+from threading import Lock
 from typing import Iterable, Mapping, Sequence
 
 import numpy as np
@@ -16,10 +20,18 @@ from .utils import collect_files
 
 DEFAULT_CACHE_DIR = Path(os.path.expanduser("~")) / ".vexor"
 CACHE_DIR = DEFAULT_CACHE_DIR
+_CACHE_DIR_OVERRIDE: ContextVar[Path | None] = ContextVar(
+    "vexor_cache_dir_override",
+    default=None,
+)
 CACHE_VERSION = 6
 DB_FILENAME = "index.db"
 EMBED_CACHE_TTL_DAYS = 30
 EMBED_CACHE_MAX_ENTRIES = 50_000
+EMBED_MEMORY_CACHE_MAX_ENTRIES = 2_048
+
+_EMBED_MEMORY_CACHE: "OrderedDict[tuple[str, int | None, str], np.ndarray]" = OrderedDict()
+_EMBED_MEMORY_LOCK = Lock()
 
 
 @dataclass(slots=True)
@@ -77,11 +89,73 @@ def query_cache_key(query: str, model: str) -> str:
     return hashlib.sha1(base.encode("utf-8")).hexdigest()
 
 
-def embedding_cache_key(text: str) -> str:
-    """Return a stable hash for embedding cache lookups."""
+def embedding_cache_key(text: str, dimension: int | None = None) -> str:
+    """Return a stable hash for embedding cache lookups.
 
+    Args:
+        text: The text to hash
+        dimension: Optional embedding dimension (included in hash for dimension-aware caching)
+    """
     clean_text = text or ""
-    return hashlib.sha1(clean_text.encode("utf-8")).hexdigest()
+    # Include dimension in hash to prevent cross-dimension cache pollution
+    if dimension is not None:
+        base = f"{clean_text}|dim={dimension}"
+    else:
+        base = clean_text
+    return hashlib.sha1(base.encode("utf-8")).hexdigest()
+
+
+def _clear_embedding_memory_cache() -> None:
+    if EMBED_MEMORY_CACHE_MAX_ENTRIES <= 0:
+        return
+    with _EMBED_MEMORY_LOCK:
+        _EMBED_MEMORY_CACHE.clear()
+
+
+def _load_embedding_memory_cache(
+    model: str,
+    text_hashes: Sequence[str],
+    dimension: int | None = None,
+) -> dict[str, np.ndarray]:
+    if EMBED_MEMORY_CACHE_MAX_ENTRIES <= 0:
+        return {}
+    results: dict[str, np.ndarray] = {}
+    with _EMBED_MEMORY_LOCK:
+        for text_hash in text_hashes:
+            if not text_hash:
+                continue
+            # Include dimension in cache key to prevent cross-dimension pollution
+            key = (model, dimension, text_hash)
+            vector = _EMBED_MEMORY_CACHE.pop(key, None)
+            if vector is None:
+                continue
+            _EMBED_MEMORY_CACHE[key] = vector
+            results[text_hash] = vector
+    return results
+
+
+def _store_embedding_memory_cache(
+    *,
+    model: str,
+    embeddings: Mapping[str, np.ndarray],
+    dimension: int | None = None,
+) -> None:
+    if EMBED_MEMORY_CACHE_MAX_ENTRIES <= 0 or not embeddings:
+        return
+    with _EMBED_MEMORY_LOCK:
+        for text_hash, vector in embeddings.items():
+            if not text_hash:
+                continue
+            array = np.asarray(vector, dtype=np.float32)
+            if array.size == 0:
+                continue
+            # Include dimension in cache key to prevent cross-dimension pollution
+            key = (model, dimension, text_hash)
+            if key in _EMBED_MEMORY_CACHE:
+                _EMBED_MEMORY_CACHE.pop(key, None)
+            _EMBED_MEMORY_CACHE[key] = array
+        while len(_EMBED_MEMORY_CACHE) > EMBED_MEMORY_CACHE_MAX_ENTRIES:
+            _EMBED_MEMORY_CACHE.popitem(last=False)
 
 
 def _serialize_extensions(extensions: Sequence[str] | None) -> str:
@@ -115,9 +189,32 @@ def _chunk_values(values: Sequence[object], size: int) -> Iterable[Sequence[obje
         yield values[idx : idx + size]
 
 
+def _resolve_cache_dir() -> Path:
+    override = _CACHE_DIR_OVERRIDE.get()
+    return override if override is not None else CACHE_DIR
+
+
+@contextmanager
+def cache_dir_context(path: Path | str | None):
+    """Temporarily override the cache directory for the current context."""
+
+    if path is None:
+        yield
+        return
+    dir_path = Path(path).expanduser().resolve()
+    if dir_path.exists() and not dir_path.is_dir():
+        raise NotADirectoryError(f"Path is not a directory: {dir_path}")
+    token = _CACHE_DIR_OVERRIDE.set(dir_path)
+    try:
+        yield
+    finally:
+        _CACHE_DIR_OVERRIDE.reset(token)
+
+
 def ensure_cache_dir() -> Path:
-    CACHE_DIR.mkdir(parents=True, exist_ok=True)
-    return CACHE_DIR
+    cache_dir = _resolve_cache_dir()
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
 
 
 def set_cache_dir(path: Path | str | None) -> None:
@@ -134,8 +231,8 @@ def set_cache_dir(path: Path | str | None) -> None:
 def cache_db_path() -> Path:
     """Return the absolute path to the shared SQLite cache database."""
 
-    ensure_cache_dir()
-    return CACHE_DIR / DB_FILENAME
+    cache_dir = ensure_cache_dir()
+    return cache_dir / DB_FILENAME
 
 
 def cache_file(root: Path, model: str, include_hidden: bool) -> Path:  # pragma: no cover - kept for API parity
@@ -1304,25 +1401,38 @@ def load_embedding_cache(
     model: str,
     text_hashes: Sequence[str],
     conn: sqlite3.Connection | None = None,
+    *,
+    dimension: int | None = None,
 ) -> dict[str, np.ndarray]:
-    """Load cached embeddings keyed by (model, text_hash)."""
-
+    """Load cached embeddings keyed by (model, text_hash).
+
+    Args:
+        model: The embedding model name
+        text_hashes: Sequence of text hashes to look up (should be generated with
+            embedding_cache_key() using the same dimension parameter)
+        conn: Optional database connection
+        dimension: Embedding dimension (used for memory cache segmentation)
+    """
     unique_hashes = list(dict.fromkeys([value for value in text_hashes if value]))
     if not unique_hashes:
         return {}
+    results = _load_embedding_memory_cache(model, unique_hashes, dimension=dimension)
+    missing = [value for value in unique_hashes if value not in results]
+    if not missing:
+        return results
     db_path = cache_db_path()
     owns_connection = conn is None
     try:
         connection = conn or _connect(db_path, readonly=True)
     except sqlite3.OperationalError:
-        return {}
+        return results
     try:
         try:
             _ensure_schema_readonly(connection, tables=("embedding_cache",))
         except sqlite3.OperationalError:
-            return {}
-        results: dict[str, np.ndarray] = {}
-        for chunk in _chunk_values(unique_hashes, 900):
+            return results
+        disk_results: dict[str, np.ndarray] = {}
+        for chunk in _chunk_values(missing, 900):
             placeholders = ", ".join("?" for _ in chunk)
             rows = connection.execute(
                 f"""
@@ -1339,7 +1449,12 @@ def load_embedding_cache(
                 vector = np.frombuffer(blob, dtype=np.float32)
                 if vector.size == 0:
                     continue
-                results[row["text_hash"]] = vector
+                disk_results[row["text_hash"]] = vector
+        if disk_results:
+            _store_embedding_memory_cache(
+                model=model, embeddings=disk_results, dimension=dimension
+            )
+            results.update(disk_results)
         return results
     finally:
         if owns_connection:
@@ -1351,11 +1466,20 @@ def store_embedding_cache(
     model: str,
     embeddings: Mapping[str, np.ndarray],
     conn: sqlite3.Connection | None = None,
+    dimension: int | None = None,
 ) -> None:
-    """Store embedding vectors keyed by (model, text_hash)."""
-
+    """Store embedding vectors keyed by (model, text_hash).
+
+    Args:
+        model: The embedding model name
+        embeddings: Dict mapping text_hash -> vector (hashes should be generated with
+            embedding_cache_key() using the same dimension parameter)
+        conn: Optional database connection
+        dimension: Embedding dimension (used for memory cache segmentation)
+    """
     if not embeddings:
         return
+    _store_embedding_memory_cache(model=model, embeddings=embeddings, dimension=dimension)
     db_path = cache_db_path()
     owns_connection = conn is None
     connection = conn or _connect(db_path)

vexor/cli.py

@@ -31,14 +31,18 @@ from .config import (
     DEFAULT_MODEL,
     DEFAULT_PROVIDER,
     DEFAULT_RERANK,
+    DEFAULT_VOYAGE_MODEL,
+    DIMENSION_SUPPORTED_MODELS,
     SUPPORTED_EXTRACT_BACKENDS,
     SUPPORTED_PROVIDERS,
     SUPPORTED_RERANKERS,
     flashrank_cache_dir,
+    get_supported_dimensions,
     load_config,
     normalize_remote_rerank_url,
     resolve_remote_rerank_api_key,
     resolve_default_model,
+    supports_dimensions,
 )
 from .modes import available_modes, get_strategy
 from .services.cache_service import is_cache_current, load_index_metadata_safe
@@ -454,6 +458,7 @@ def search(
         rerank=rerank,
         flashrank_model=flashrank_model,
         remote_rerank=remote_rerank,
+        embedding_dimensions=config.embedding_dimensions,
     )
     if output_format == SearchOutputFormat.rich:
         if no_cache:
@@ -488,7 +493,7 @@ def search(
         else:
             typer.echo(message, err=True)
         raise typer.Exit(code=1)
-    except RuntimeError as exc:
+    except (RuntimeError, ValueError) as exc:
         if output_format == SearchOutputFormat.rich:
             console.print(_styled(str(exc), Styles.ERROR))
         else:
@@ -688,8 +693,9 @@ def index(
             local_cuda=bool(config.local_cuda),
             exclude_patterns=normalized_excludes,
             extensions=normalized_exts,
+            embedding_dimensions=config.embedding_dimensions,
         )
-    except RuntimeError as exc:
+    except (RuntimeError, ValueError) as exc:
         console.print(_styled(str(exc), Styles.ERROR))
         raise typer.Exit(code=1)
     if result.status == IndexStatus.EMPTY:
@@ -768,6 +774,16 @@ def config(
         "--clear-base-url",
         help=Messages.HELP_CLEAR_BASE_URL,
     ),
+    set_embedding_dimensions_option: int | None = typer.Option(
+        None,
+        "--set-embedding-dimensions",
+        help=Messages.HELP_SET_EMBEDDING_DIMENSIONS,
+    ),
+    clear_embedding_dimensions: bool = typer.Option(
+        False,
+        "--clear-embedding-dimensions",
+        help=Messages.HELP_CLEAR_EMBEDDING_DIMENSIONS,
+    ),
     set_auto_index_option: str | None = typer.Option(
         None,
         "--set-auto-index",
@@ -989,6 +1005,33 @@ def config(
         except ValueError as exc:
             raise typer.BadParameter(str(exc)) from exc
 
+    effective_embedding_dimensions = set_embedding_dimensions_option
+    effective_clear_embedding_dimensions = clear_embedding_dimensions
+    if effective_embedding_dimensions == 0:
+        effective_embedding_dimensions = None
+        effective_clear_embedding_dimensions = True
+
+    # Validate embedding dimensions if set
+    if effective_embedding_dimensions is not None:
+        if effective_embedding_dimensions < 0:
+            raise typer.BadParameter(
+                f"--set-embedding-dimensions must be non-negative, got {effective_embedding_dimensions}"
+            )
+        if effective_embedding_dimensions > 0:
+            # Resolve effective model from provider + model to account for provider defaults
+            effective_model = resolve_default_model(pending_provider, pending_model)
+            if not supports_dimensions(effective_model):
+                raise typer.BadParameter(
+                    f"Model '{effective_model}' does not support custom dimensions. "
+                    f"Supported model names/prefixes: {', '.join(DIMENSION_SUPPORTED_MODELS.keys())}"
+                )
+            supported = get_supported_dimensions(effective_model)
+            if supported and effective_embedding_dimensions not in supported:
+                raise typer.BadParameter(
+                    f"Dimension {effective_embedding_dimensions} is not supported for model '{effective_model}'. "
+                    f"Supported dimensions: {supported}"
+                )
+
     updates = apply_config_updates(
         api_key=set_api_key_option,
         clear_api_key=clear_api_key,
@@ -1007,6 +1050,8 @@ def config(
         remote_rerank_model=set_remote_rerank_model_option,
         remote_rerank_api_key=set_remote_rerank_api_key_option,
         clear_remote_rerank=clear_remote_rerank,
+        embedding_dimensions=effective_embedding_dimensions,
+        clear_embedding_dimensions=effective_clear_embedding_dimensions,
     )
 
     if updates.api_key_set:
@@ -1109,6 +1154,17 @@ def config(
         console.print(_styled(Messages.INFO_REMOTE_RERANK_API_KEY_SET, Styles.SUCCESS))
     if updates.remote_rerank_cleared and clear_remote_rerank:
         console.print(_styled(Messages.INFO_REMOTE_RERANK_CLEARED, Styles.SUCCESS))
+    if updates.embedding_dimensions_set and effective_embedding_dimensions is not None:
+        console.print(
+            _styled(
+                Messages.INFO_EMBEDDING_DIMENSIONS_SET.format(
+                    value=effective_embedding_dimensions
+                ),
+                Styles.SUCCESS,
+            )
+        )
+    if updates.embedding_dimensions_cleared:
+        console.print(_styled(Messages.INFO_EMBEDDING_DIMENSIONS_CLEARED, Styles.SUCCESS))
 
     if clear_flashrank:
         cache_dir = flashrank_cache_dir(create=False)
@@ -1188,6 +1244,7 @@ def config(
                     api="yes" if cfg.api_key else "no",
                     provider=provider,
                     model=resolve_default_model(provider, cfg.model),
+                    embedding_dimensions=cfg.embedding_dimensions if cfg.embedding_dimensions else "default",
                     batch=cfg.batch_size if cfg.batch_size is not None else DEFAULT_BATCH_SIZE,
                     concurrency=cfg.embed_concurrency,
                     extract_concurrency=cfg.extract_concurrency,