codegraph-cli 2.0.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

codegraph_cli/__init__.py

@@ -1,4 +1,4 @@
 """CodeGraph CLI package."""
 
 __all__ = ["__version__"]
-__version__ = "2.0.0"
+__version__ = "2.0.1"

codegraph_cli/agents.py

@@ -6,7 +6,7 @@ from collections import deque
 from pathlib import Path
 from typing import Dict, List, Set
 
-from .embeddings import HashEmbeddingModel
+from .embeddings import HashEmbeddingModel, TransformerEmbedder
 from .llm import LocalLLM
 from .models import ImpactReport
 from .parser import PythonGraphParser

codegraph_cli/cli.py

@@ -10,6 +10,7 @@ import typer
 from . import __version__, config
 from .cli_chat import chat_app
 from .cli_setup import setup as setup_wizard, set_llm, unset_llm, show_llm
+from .cli_setup import set_embedding, unset_embedding, show_embedding
 from .cli_v2 import v2_app
 from .graph_export import export_dot, export_html
 from .orchestrator import MCPOrchestrator
@@ -35,6 +36,11 @@ app.command("set-llm")(set_llm)
 app.command("unset-llm")(unset_llm)
 app.command("show-llm")(show_llm)
 
+# Register embedding management commands
+app.command("set-embedding")(set_embedding)
+app.command("unset-embedding")(unset_embedding)
+app.command("show-embedding")(show_embedding)
+
 
 def version_callback(value: bool):
     """Print version and exit."""

codegraph_cli/cli_chat.py

@@ -281,7 +281,7 @@ def start_chat(
     new_session: bool = typer.Option(False, "--new", "-n", help="Force start a new session"),
 ):
     """Start interactive chat session."""
-    from .embeddings import HashEmbeddingModel
+    from .embeddings import get_embedder
     from .project_context import ProjectContext
     
     pm = ProjectManager()
@@ -294,7 +294,7 @@ def start_chat(
     
     # Initialize components
     context = ProjectContext(project, pm)
-    embedding_model = HashEmbeddingModel()
+    embedding_model = get_embedder()
     llm = LocalLLM(model=llm_model, provider=llm_provider, api_key=llm_api_key, endpoint=llm_endpoint)
     rag_retriever = RAGRetriever(context.store, embedding_model)
     

codegraph_cli/cli_setup.py

@@ -8,6 +8,7 @@ from typing import Optional
 import typer
 
 from . import config_manager
+from .embeddings import EMBEDDING_MODELS
 
 app = typer.Typer(help="Setup wizard for LLM provider configuration")
 
@@ -287,6 +288,12 @@ def setup():
         print_error("Failed to save configuration!")
         raise typer.Exit(code=1)
 
+    # Offer embedding setup
+    typer.echo("")
+    setup_emb = typer.confirm("Configure embedding model for semantic search?", default=True)
+    if setup_emb:
+        _interactive_embedding_setup()
+
 
 def set_llm(
     provider: str = typer.Argument(..., help="LLM provider: ollama, groq, openai, anthropic, gemini, openrouter"),
@@ -466,5 +473,156 @@ def show_llm():
     typer.echo("")
 
 
+# ===================================================================
+# Embedding model commands
+# ===================================================================
+
+def _interactive_embedding_setup():
+    """Interactive embedding model picker (called from setup wizard)."""
+    typer.echo("")
+    typer.echo(typer.style("╭──────────────────────────────────────────────╮", fg=typer.colors.CYAN))
+    typer.echo(typer.style("│", fg=typer.colors.CYAN) + typer.style("   Embedding Model Setup                       ", bold=True) + typer.style("│", fg=typer.colors.CYAN))
+    typer.echo(typer.style("╰──────────────────────────────────────────────╯", fg=typer.colors.CYAN))
+    typer.echo("")
+    typer.echo("Choose an embedding model for semantic code search:")
+    typer.echo("Larger models give better results but need more disk/RAM.\n")
+
+    # List models with numbers
+    model_keys = list(EMBEDDING_MODELS.keys())
+    for i, key in enumerate(model_keys, 1):
+        spec = EMBEDDING_MODELS[key]
+        name_col = f"{key}".ljust(12)
+        size_col = f"({spec['size']})".ljust(14)
+        desc = spec["description"]
+        typer.echo(f"  {i}) {name_col} {size_col} {desc}")
+
+    typer.echo("")
+
+    while True:
+        choice = typer.prompt(f"Enter choice [1-{len(model_keys)}]", type=str)
+        try:
+            idx = int(choice)
+            if 1 <= idx <= len(model_keys):
+                selected = model_keys[idx - 1]
+                break
+        except ValueError:
+            # Accept model key directly
+            if choice.strip() in model_keys:
+                selected = choice.strip()
+                break
+        print_error(f"Invalid choice. Enter 1-{len(model_keys)} or a model key.")
+
+    spec = EMBEDDING_MODELS[selected]
+
+    if selected != "hash":
+        typer.echo(f"\n  Model:    {typer.style(spec['name'], fg=typer.colors.CYAN)}")
+        typer.echo(f"  Download: {typer.style(spec['size'], fg=typer.colors.YELLOW)}")
+        typer.echo(f"  Dim:      {spec['dim']}")
+        print_info("Requires: pip install codegraph-cli[embeddings]")
+    else:
+        typer.echo(f"\n  Model: {typer.style('Hash Embedding (zero-dependency)', fg=typer.colors.CYAN)}")
+        print_info("No download needed, but no semantic understanding.")
+
+    success = config_manager.save_embedding_config(selected)
+    if success:
+        print_success(f"Embedding model set to: {selected}")
+        if selected != "hash":
+            print_info(f"Model will be downloaded on first use (~{spec['size']}).")
+            print_info("Re-index your project after changing embeddings: cg index <path>")
+    else:
+        print_error("Failed to save embedding config!")
+
+
+def set_embedding(
+    model: str = typer.Argument(
+        ...,
+        help="Embedding model key: qodo-1.5b, jina-code, bge-base, minilm, hash",
+    ),
+):
+    """Set the embedding model for semantic code search.
+
+    Available models (smallest to largest):
+
+        hash        0 bytes    No download, keyword-level only
+        minilm      ~80 MB     Tiny, fast, decent quality
+        bge-base    ~440 MB    Solid general-purpose
+        jina-code   ~550 MB    Code-aware, good quality
+        qodo-1.5b   ~6.2 GB   Best quality, code-optimized
+
+    Examples:
+        cg set-embedding minilm
+        cg set-embedding jina-code
+        cg set-embedding hash
+    """
+    model = model.lower().strip()
+
+    if model not in EMBEDDING_MODELS:
+        print_error(
+            f"Unknown model '{model}'. "
+            f"Choose from: {', '.join(EMBEDDING_MODELS.keys())}"
+        )
+        raise typer.Exit(code=1)
+
+    spec = EMBEDDING_MODELS[model]
+    success = config_manager.save_embedding_config(model)
+
+    if success:
+        print_success(f"Embedding model set to: {model}")
+        typer.echo(f"  Name: {typer.style(spec['name'], fg=typer.colors.CYAN)}")
+        typer.echo(f"  Dim:  {spec['dim']}")
+        if model != "hash":
+            typer.echo(f"  Size: {spec['size']} (downloaded on first use)")
+            print_info("Re-index your project after changing: cg index <path>")
+    else:
+        print_error("Failed to save configuration!")
+        raise typer.Exit(code=1)
+
+
+def unset_embedding():
+    """Reset embedding model to default (hash — no download)."""
+    success = config_manager.clear_embedding_config()
+    if success:
+        print_success("Embedding model reset to default (hash).")
+        print_info("No neural model will be used. Re-index to apply.")
+    else:
+        print_error("Failed to reset embedding config!")
+        raise typer.Exit(code=1)
+
+
+def show_embedding():
+    """Show current embedding model configuration."""
+    typer.echo("")
+    typer.echo(typer.style("╭──────────────────────────────────────────────╮", fg=typer.colors.CYAN))
+    typer.echo(typer.style("│", fg=typer.colors.CYAN) + typer.style("   Embedding Configuration                     ", bold=True) + typer.style("│", fg=typer.colors.CYAN))
+    typer.echo(typer.style("╰──────────────────────────────────────────────╯", fg=typer.colors.CYAN))
+
+    emb_cfg = config_manager.load_embedding_config()
+    current_key = emb_cfg.get("model", "hash")
+    spec = EMBEDDING_MODELS.get(current_key)
+
+    if spec is None:
+        typer.echo(f"  Model   {typer.style(current_key, fg=typer.colors.RED)} (unknown)")
+    else:
+        typer.echo(f"  Model   {typer.style(f' {current_key} ', bg=typer.colors.CYAN, fg=typer.colors.WHITE, bold=True)}")
+        typer.echo(f"  Name    {typer.style(spec['name'], bold=True)}")
+        typer.echo(f"  Dim     {spec['dim']}")
+        typer.echo(f"  Size    {spec['size']}")
+        typer.echo(f"  Desc    {spec['description']}")
+
+    typer.echo("")
+    typer.echo(typer.style("  Available Models", bold=True))
+    typer.echo(typer.style("  ─────────────────────────────────────────", dim=True))
+    for key, s in EMBEDDING_MODELS.items():
+        marker = typer.style(" *", fg=typer.colors.GREEN) if key == current_key else "  "
+        typer.echo(f"  {marker} {key.ljust(12)} {s['size'].ljust(12)} {s['description']}")
+
+    typer.echo("")
+    typer.echo(typer.style("  Quick Commands", bold=True))
+    typer.echo(typer.style("  ─────────────────────────────────────────", dim=True))
+    typer.echo(f"  {typer.style('cg set-embedding <model>', fg=typer.colors.YELLOW)}   Switch model")
+    typer.echo(f"  {typer.style('cg unset-embedding', fg=typer.colors.YELLOW)}         Reset to hash")
+    typer.echo("")
+
+
 if __name__ == "__main__":
     app()

codegraph_cli/config.py

@@ -13,10 +13,12 @@ SUPPORTED_EXTENSIONS = {".py"}
 
 # Load configuration from TOML file (if available)
 try:
-    from .config_manager import load_config
+    from .config_manager import load_config, load_embedding_config
     _toml_config = load_config()
+    _emb_config = load_embedding_config()
 except ImportError:
     _toml_config = {}
+    _emb_config = {}
 
 # LLM Provider Configuration — loaded from ~/.codegraph/config.toml (set via `cg setup` or `cg set-llm`)
 LLM_PROVIDER = _toml_config.get("provider", "ollama")
@@ -24,6 +26,9 @@ LLM_API_KEY = _toml_config.get("api_key", "")
 LLM_MODEL = _toml_config.get("model", "qwen2.5-coder:7b")
 LLM_ENDPOINT = _toml_config.get("endpoint", "http://127.0.0.1:11434/api/generate")
 
+# Embedding model — set via `cg set-embedding` (default: "hash" = no download)
+EMBEDDING_MODEL = _emb_config.get("model", "hash")
+
 
 def ensure_base_dirs() -> None:
     """Create base directories for local storage if needed."""

codegraph_cli/config_manager.py

@@ -78,11 +78,37 @@ def load_config() -> Dict[str, Any]:
         return DEFAULT_CONFIGS["ollama"].copy()
 
 
+def load_full_config() -> Dict[str, Any]:
+    """Load the entire TOML config (all sections)."""
+    if not CONFIG_FILE.exists() or toml is None:
+        return {}
+    try:
+        with open(CONFIG_FILE, "r") as f:
+            return toml.load(f)
+    except Exception:
+        return {}
+
+
+def _save_full_config(config: Dict[str, Any]) -> bool:
+    """Write entire config dict to TOML file, preserving all sections."""
+    if toml is None:
+        return False
+    BASE_DIR.mkdir(parents=True, exist_ok=True)
+    try:
+        with open(CONFIG_FILE, "w") as f:
+            toml.dump(config, f)
+        return True
+    except Exception:
+        return False
+
+
 def save_config(provider: str, model: str, api_key: str = "", endpoint: str = "") -> bool:
     """Save LLM configuration to TOML file.
     
+    Preserves other sections (e.g. ``[embeddings]``) in the file.
+    
     Args:
-        provider: Provider name (ollama, groq, openai, anthropic)
+        provider: Provider name (ollama, groq, openai, anthropic, gemini, openrouter)
         model: Model name
         api_key: API key for cloud providers
         endpoint: Custom endpoint (for Ollama)
@@ -90,32 +116,56 @@ def save_config(provider: str, model: str, api_key: str = "", endpoint: str = ""
     Returns:
         True if saved successfully, False otherwise
     """
-    if toml is None:
-        return False
-    
-    # Ensure directory exists
-    BASE_DIR.mkdir(parents=True, exist_ok=True)
+    config = load_full_config()
     
-    # Build config
-    config = {
-        "llm": {
-            "provider": provider,
-            "model": model,
-        }
+    config["llm"] = {
+        "provider": provider,
+        "model": model,
     }
-    
     if api_key:
         config["llm"]["api_key"] = api_key
-    
     if endpoint:
         config["llm"]["endpoint"] = endpoint
     
-    try:
-        with open(CONFIG_FILE, "w") as f:
-            toml.dump(config, f)
-        return True
-    except Exception:
-        return False
+    return _save_full_config(config)
+
+
+# ------------------------------------------------------------------
+# Embedding configuration
+# ------------------------------------------------------------------
+
+def load_embedding_config() -> Dict[str, Any]:
+    """Load embedding configuration from ``[embeddings]`` section.
+    
+    Returns:
+        Dict with at least ``model`` key, or empty dict.
+    """
+    full = load_full_config()
+    return full.get("embeddings", {})
+
+
+def save_embedding_config(model_key: str) -> bool:
+    """Save embedding model choice to config TOML.
+    
+    Preserves ``[llm]`` and other sections.
+    
+    Args:
+        model_key: One of the keys from ``EMBEDDING_MODELS``
+                   (e.g. ``"minilm"``, ``"jina-code"``, ``"hash"``).
+    
+    Returns:
+        True if saved successfully.
+    """
+    config = load_full_config()
+    config["embeddings"] = {"model": model_key}
+    return _save_full_config(config)
+
+
+def clear_embedding_config() -> bool:
+    """Remove ``[embeddings]`` section from config, resetting to default."""
+    config = load_full_config()
+    config.pop("embeddings", None)
+    return _save_full_config(config)
 
 
 def get_provider_config(provider: str) -> Dict[str, Any]:

codegraph_cli/context_manager.py

@@ -56,7 +56,7 @@ class RepoMap:
           class: PythonGraphParser
           function: _resolve_call_edges
         codegraph_cli/embeddings.py
-          class: NeuralEmbedder
+          class: TransformerEmbedder
           class: HashEmbeddingModel
           function: get_embedder
           function: cosine_similarity

codegraph_cli/embeddings.py

@@ -1,23 +1,32 @@
-"""Neural embedding engine using Sentence Transformers for semantic code understanding.
-
-Local-first architecture:
-- Models are downloaded once and cached in ``~/.codegraph/models``.
-- All inference runs on-device (CPU or GPU).  No data is ever sent to
-  external APIs.
-
-Falls back to a lightweight deterministic hash-embedding when
-``sentence-transformers`` is not installed.
+"""Configurable code embedding engine with multiple model support.
+
+Supported models (configure via ``cg set-embedding``):
+
+========== ====================================== ========= ====== ======================
+Key        HuggingFace Model                      Download  Dim    Notes
+========== ====================================== ========= ====== ======================
+qodo-1.5b  Qodo/Qodo-Embed-1-1.5B                ~6.2 GB   1536   Best quality, code-optimized
+jina-code  jinaai/jina-embeddings-v2-base-code    ~550 MB    768   Good quality, code-aware
+bge-base   BAAI/bge-base-en-v1.5                  ~440 MB    768   Solid general-purpose
+minilm     sentence-transformers/all-MiniLM-L6-v2  ~80 MB    384   Tiny and fast
+hash       (none)                                     0 B    256   No ML, keyword-level only
+========== ====================================== ========= ====== ======================
+
+Architecture:
+- Models downloaded once from HuggingFace and cached in ``~/.codegraph/models``.
+- All inference runs on-device (CPU or GPU). No data leaves the machine.
+- Uses raw ``transformers`` library only — no sentence-transformers, no flash_attn.
+- Falls back to hash embeddings when ``torch``/``transformers`` are not installed.
 """
 
 from __future__ import annotations
 
 import logging
 import math
-import os
 import re
 from hashlib import blake2b
 from pathlib import Path
-from typing import Iterable, List, Optional, Union
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 from .config import BASE_DIR
 
@@ -26,44 +35,115 @@ logger = logging.getLogger(__name__)
 # Default local model cache directory
 MODEL_CACHE_DIR: Path = BASE_DIR / "models"
 
-# Preferred models in priority order
-PREFERRED_MODELS: List[str] = [
-    "all-MiniLM-L6-v2",
-    "nomic-ai/nomic-embed-text-v1.5",
-]
-
 _TOKEN_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
 
 
 # ===================================================================
-# NeuralEmbedder  (Primary – Sentence Transformers)
+# Model Registry
+# ===================================================================
+
+EMBEDDING_MODELS: Dict[str, Dict[str, Any]] = {
+    "qodo-1.5b": {
+        "name": "Qodo Embed 1.5B",
+        "hf_id": "Qodo/Qodo-Embed-1-1.5B",
+        "dim": 1536,
+        "max_tokens": 8192,
+        "size": "~6.2 GB",
+        "description": "Best quality, code-optimized (needs 8GB+ RAM)",
+        "pooling": "last_token",
+        "trust_remote_code": True,
+    },
+    "jina-code": {
+        "name": "Jina Embeddings v2 Code",
+        "hf_id": "jinaai/jina-embeddings-v2-base-code",
+        "dim": 768,
+        "max_tokens": 8192,
+        "size": "~550 MB",
+        "description": "Good quality, code-aware, lightweight",
+        "pooling": "mean",
+        "trust_remote_code": True,
+    },
+    "bge-base": {
+        "name": "BGE Base EN v1.5",
+        "hf_id": "BAAI/bge-base-en-v1.5",
+        "dim": 768,
+        "max_tokens": 512,
+        "size": "~440 MB",
+        "description": "Solid general-purpose, fast",
+        "pooling": "cls",
+        "trust_remote_code": False,
+    },
+    "minilm": {
+        "name": "MiniLM L6 v2",
+        "hf_id": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "max_tokens": 256,
+        "size": "~80 MB",
+        "description": "Tiny and fast, decent quality",
+        "pooling": "mean",
+        "trust_remote_code": False,
+    },
+    "hash": {
+        "name": "Hash Embedding",
+        "hf_id": None,
+        "dim": 256,
+        "max_tokens": None,
+        "size": "0 bytes",
+        "description": "Zero-dependency fallback, no semantics",
+        "pooling": None,
+        "trust_remote_code": False,
+    },
+}
+
+DEFAULT_MODEL = "hash"
+
+
+# ===================================================================
+# TransformerEmbedder  (handles all HuggingFace models)
 # ===================================================================
 
-class NeuralEmbedder:
-    """Semantic embedding engine powered by Sentence Transformers.
+class TransformerEmbedder:
+    """Generic HuggingFace embedding engine with configurable pooling.
 
-    The model is downloaded on first use and cached in
-    ``~/.codegraph/models`` so that subsequent runs are fully offline.
-    All computation is local – **no data leaves the machine**.
+    Supports multiple pooling strategies:
 
-    Example::
+    - **last_token** — last non-padding token (Qodo models).
+    - **mean** — mean over non-padding tokens (Jina, MiniLM).
+    - **cls** — ``[CLS]`` first token (BGE models).
 
-        embedder = NeuralEmbedder()
-        vecs = embedder.embed_documents(["def hello(): ...", "class Foo: ..."])
+    Model weights are downloaded on first use and cached in
+    ``~/.codegraph/models/`` for offline subsequent runs.
     """
 
     def __init__(
         self,
-        model_name: str = "all-MiniLM-L6-v2",
+        model_key: str,
         cache_dir: Optional[Path] = None,
         device: str = "cpu",
     ) -> None:
-        self.model_name = model_name
+        if model_key not in EMBEDDING_MODELS:
+            raise ValueError(
+                f"Unknown model: '{model_key}'. "
+                f"Available: {', '.join(EMBEDDING_MODELS.keys())}"
+            )
+
+        spec = EMBEDDING_MODELS[model_key]
+        if spec["hf_id"] is None:
+            raise ValueError(
+                f"'{model_key}' has no transformer backend. Use HashEmbeddingModel."
+            )
+
+        self.model_key = model_key
+        self.hf_id: str = spec["hf_id"]
+        self.dim: int = spec["dim"]
+        self.max_length: int = spec["max_tokens"]
+        self.pooling: str = spec["pooling"]
+        self.trust_remote_code: bool = spec["trust_remote_code"]
         self.cache_dir = cache_dir or MODEL_CACHE_DIR
         self.cache_dir.mkdir(parents=True, exist_ok=True)
         self.device = device
-        self._model: object = None  # lazy-loaded SentenceTransformer
-        self._dim: Optional[int] = None
+        self._model: Any = None
+        self._tokenizer: Any = None
 
     # ------------------------------------------------------------------
     # Lazy model loading
@@ -74,100 +154,154 @@ class NeuralEmbedder:
             return
 
         try:
-            from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]
+            import torch  # noqa: F401
+            from transformers import AutoModel, AutoTokenizer
         except ImportError:
             raise ImportError(
-                "sentence-transformers is not installed. "
-                "Install with:  pip install sentence-transformers"
+                "torch and transformers are required for neural embeddings.\n"
+                "Install with:  pip install codegraph-cli[embeddings]\n"
+                "For CPU-only (skip NVIDIA packages):\n"
+                "  pip install torch --index-url https://download.pytorch.org/whl/cpu\n"
+                "  pip install transformers"
             )
 
-        # Tell sentence-transformers where to cache
-        os.environ.setdefault(
-            "SENTENCE_TRANSFORMERS_HOME", str(self.cache_dir),
+        logger.info(
+            "Loading embedding model '%s' (%s) — first run downloads %s...",
+            self.model_key, self.hf_id, EMBEDDING_MODELS[self.model_key]["size"],
         )
 
         try:
-            self._model = SentenceTransformer(
-                self.model_name,
-                cache_folder=str(self.cache_dir),
-                device=self.device,
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.hf_id,
+                cache_dir=str(self.cache_dir),
+                trust_remote_code=self.trust_remote_code,
             )
-            self._dim = self._model.get_sentence_embedding_dimension()  # type: ignore[union-attr]
+            self._model = AutoModel.from_pretrained(
+                self.hf_id,
+                cache_dir=str(self.cache_dir),
+                trust_remote_code=self.trust_remote_code,
+            )
+            self._model.eval()
+            self._model.to(self.device)
             logger.info(
-                "Loaded model '%s' (dim=%d) on %s",
-                self.model_name, self._dim, self.device,
+                "Loaded '%s' (dim=%d, pooling=%s) on %s",
+                self.model_key, self.dim, self.pooling, self.device,
             )
         except Exception as exc:
             raise RuntimeError(
-                f"Failed to load embedding model '{self.model_name}': {exc}"
+                f"Failed to load embedding model '{self.model_key}' "
+                f"({self.hf_id}): {exc}"
             ) from exc
 
     # ------------------------------------------------------------------
-    # Public API
+    # Pooling strategies
     # ------------------------------------------------------------------
 
-    @property
-    def dim(self) -> int:
-        """Dimensionality of the embedding vectors."""
-        if self._dim is None:
-            self._load_model()
-        assert self._dim is not None
-        return self._dim
+    @staticmethod
+    def _pool_last_token(last_hidden_states: Any, attention_mask: Any) -> Any:
+        """Last non-padding token (Qodo style)."""
+        import torch
+
+        left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+        if left_padding:
+            return last_hidden_states[:, -1]
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[
+            torch.arange(batch_size, device=last_hidden_states.device),
+            sequence_lengths,
+        ]
+
+    @staticmethod
+    def _pool_mean(last_hidden_states: Any, attention_mask: Any) -> Any:
+        """Mean over non-padding tokens (Jina, MiniLM)."""
+        mask_expanded = attention_mask.unsqueeze(-1).expand(
+            last_hidden_states.size()
+        ).float()
+        sum_embeddings = (last_hidden_states * mask_expanded).sum(dim=1)
+        sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)
+        return sum_embeddings / sum_mask
+
+    @staticmethod
+    def _pool_cls(last_hidden_states: Any, attention_mask: Any) -> Any:
+        """[CLS] first token (BGE)."""
+        return last_hidden_states[:, 0]
+
+    def _pool(self, last_hidden_states: Any, attention_mask: Any) -> Any:
+        """Dispatch to the pooling strategy for this model."""
+        if self.pooling == "last_token":
+            return self._pool_last_token(last_hidden_states, attention_mask)
+        if self.pooling == "mean":
+            return self._pool_mean(last_hidden_states, attention_mask)
+        if self.pooling == "cls":
+            return self._pool_cls(last_hidden_states, attention_mask)
+        raise ValueError(f"Unknown pooling strategy: {self.pooling}")
+
+    # ------------------------------------------------------------------
+    # Encode
+    # ------------------------------------------------------------------
+
+    def _encode(self, texts: List[str]) -> List[List[float]]:
+        """Encode a batch of texts into L2-normalised embedding vectors."""
+        import torch
+        import torch.nn.functional as F
 
-    def embed_text(self, text: str) -> List[float]:
-        """Embed a single text string and return a unit-norm vector."""
         self._load_model()
-        assert self._model is not None
-        embedding = self._model.encode(  # type: ignore[union-attr]
-            text,
-            convert_to_numpy=True,
-            normalize_embeddings=True,
-            show_progress_bar=False,
+
+        batch_dict = self._tokenizer(
+            texts,
+            max_length=self.max_length,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+        )
+        batch_dict = {k: v.to(self.device) for k, v in batch_dict.items()}
+
+        with torch.no_grad():
+            outputs = self._model(**batch_dict)
+
+        embeddings = self._pool(
+            outputs.last_hidden_state, batch_dict["attention_mask"],
         )
-        return embedding.tolist()
+        embeddings = F.normalize(embeddings, p=2, dim=1)
+        return embeddings.cpu().tolist()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def embed_text(self, text: str) -> List[float]:
+        """Embed a single text string and return a unit-norm vector."""
+        return self._encode([text])[0]
 
     def embed_documents(
         self,
         texts: List[str],
-        batch_size: int = 32,
+        batch_size: int = 16,
     ) -> List[List[float]]:
-        """Embed multiple documents with batching for efficiency.
-
-        Args:
-            texts: List of text strings to embed.
-            batch_size: Number of texts per forward pass.
-
-        Returns:
-            List of embedding vectors (each normalised to unit length).
-        """
+        """Embed multiple documents with batching."""
         if not texts:
             return []
-        self._load_model()
-        assert self._model is not None
-        embeddings = self._model.encode(  # type: ignore[union-attr]
-            texts,
-            batch_size=batch_size,
-            convert_to_numpy=True,
-            normalize_embeddings=True,
-            show_progress_bar=len(texts) > 100,
-        )
-        return embeddings.tolist()
+        all_embeddings: List[List[float]] = []
+        for i in range(0, len(texts), batch_size):
+            all_embeddings.extend(self._encode(texts[i : i + batch_size]))
+        return all_embeddings
 
-    # Backward-compat alias used by legacy callers
     def embed_many(self, texts: Iterable[str]) -> List[List[float]]:
         """Alias for :meth:`embed_documents`."""
         return self.embed_documents(list(texts))
 
 
 # ===================================================================
-# HashEmbeddingModel  (Lightweight Fallback)
+# HashEmbeddingModel  (Zero-dependency fallback)
 # ===================================================================
 
 class HashEmbeddingModel:
-    """Deterministic token-hashing embedder – no ML dependencies.
+    """Deterministic token-hashing embedder — no ML dependencies.
 
-    Provides basic keyword-level similarity.  Automatically used as a
-    fallback when ``sentence-transformers`` is not available.
+    Provides basic keyword-level similarity. Used as the default when
+    ``torch``/``transformers`` are not installed or when ``hash`` is
+    selected via ``cg set-embedding hash``.
     """
 
     def __init__(self, dim: int = 256) -> None:
@@ -189,7 +323,7 @@ class HashEmbeddingModel:
         return [self.embed_text(text) for text in texts]
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
-        """Alias matching the NeuralEmbedder interface."""
+        """Alias matching the TransformerEmbedder interface."""
         return self.embed_many(texts)
 
 
@@ -198,27 +332,61 @@ class HashEmbeddingModel:
 # ===================================================================
 
 def get_embedder(
-    model_name: Optional[str] = None,
+    model_key: Optional[str] = None,
     cache_dir: Optional[Path] = None,
     device: str = "cpu",
-) -> Union[NeuralEmbedder, HashEmbeddingModel]:
-    """Return the best available embedder.
+) -> Union[TransformerEmbedder, HashEmbeddingModel]:
+    """Return the configured embedder.
+
+    Resolution order:
+
+    1. Explicit ``model_key`` argument.
+    2. ``[embeddings].model`` from ``~/.codegraph/config.toml``.
+    3. ``"hash"`` (zero-dependency fallback).
 
-    * If ``sentence-transformers`` is installed → :class:`NeuralEmbedder`.
-    * Otherwise → :class:`HashEmbeddingModel` (zero-dependency fallback).
+    If a transformer model is configured but ``torch``/``transformers``
+    are missing, falls back to hash with a warning.
     """
+    if model_key is None:
+        try:
+            from .config_manager import load_embedding_config
+            emb_cfg = load_embedding_config()
+            model_key = emb_cfg.get("model", None)
+        except Exception:
+            model_key = None
+
+    # Default to hash if nothing configured
+    if model_key is None:
+        model_key = DEFAULT_MODEL
+
+    # Hash path — no ML needed
+    if model_key == "hash":
+        return HashEmbeddingModel()
+
+    # Unknown model guard
+    if model_key not in EMBEDDING_MODELS:
+        logger.warning(
+            "Unknown embedding model '%s' — falling back to hash.", model_key,
+        )
+        return HashEmbeddingModel()
+
+    spec = EMBEDDING_MODELS[model_key]
+    if spec["hf_id"] is None:
+        return HashEmbeddingModel()
+
+    # Transformer path — check dependencies
     try:
-        import sentence_transformers  # noqa: F401
-        return NeuralEmbedder(
-            model_name=model_name or "all-MiniLM-L6-v2",
-            cache_dir=cache_dir,
-            device=device,
+        import torch  # noqa: F401
+        import transformers  # noqa: F401
+        return TransformerEmbedder(
+            model_key=model_key, cache_dir=cache_dir, device=device,
         )
     except ImportError:
         logger.warning(
-            "sentence-transformers not installed – "
-            "using hash-based embeddings (no semantic understanding). "
-            "Install with: pip install sentence-transformers"
+            "Embedding model '%s' requires torch + transformers. "
+            "Falling back to hash embeddings.  Install with: "
+            "pip install codegraph-cli[embeddings]",
+            model_key,
         )
         return HashEmbeddingModel()
 

codegraph_cli/orchestrator.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from typing import Dict, List
 
 from .agents import GraphAgent, RAGAgent, SummarizationAgent
-from .embeddings import HashEmbeddingModel
+from .embeddings import get_embedder
 from .llm import LocalLLM
 from .models import ImpactReport, SearchResult
 from .rag import RAGRetriever
@@ -25,7 +25,7 @@ class MCPOrchestrator:
         llm_endpoint: str | None = None,
     ):
         self.store = store
-        self.embedding_model = HashEmbeddingModel()
+        self.embedding_model = get_embedder()
         self.graph_agent = GraphAgent(store, self.embedding_model)
         self.rag_agent = RAGAgent(RAGRetriever(store, self.embedding_model))
         self.summarization_agent = SummarizationAgent(

codegraph_cli/rag.py

@@ -11,7 +11,7 @@ import json
 import logging
 from typing import Any, Dict, List, Optional, Union
 
-from .embeddings import HashEmbeddingModel, NeuralEmbedder, cosine_similarity
+from .embeddings import HashEmbeddingModel, TransformerEmbedder, cosine_similarity
 from .models import SearchResult
 from .storage import GraphStore
 
@@ -29,14 +29,14 @@ class RAGRetriever:
        cosine similarity in Python.
 
     The ``embedding_model`` argument accepts either a
-    :class:`~codegraph_cli.embeddings.NeuralEmbedder` or the legacy
+    :class:`~codegraph_cli.embeddings.TransformerEmbedder` or the lightweight
     :class:`~codegraph_cli.embeddings.HashEmbeddingModel`.
     """
 
     def __init__(
         self,
         store: GraphStore,
-        embedding_model: Union[NeuralEmbedder, HashEmbeddingModel, Any],
+        embedding_model: Union[TransformerEmbedder, HashEmbeddingModel, Any],
     ) -> None:
         self.store = store
         self.embedding_model = embedding_model

{codegraph_cli-2.0.0.dist-info → codegraph_cli-2.1.0.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: codegraph-cli
-Version: 2.0.0
+Version: 2.1.0
 Summary: AI-powered code intelligence CLI with multi-agent analysis, impact graphs, and conversational coding.
-Author-email: Ali Nasir <ali@codegraph.dev>
+Author-email: Ali Nasir <muhammadalinasir00786@gmail.com>
 License: MIT
 Project-URL: Homepage, https://github.com/al1-nasir/codegraph-cli
 Project-URL: Documentation, https://github.com/al1-nasir/codegraph-cli#readme
@@ -31,7 +31,6 @@ Requires-Dist: typer<1.0.0,>=0.12.0
 Requires-Dist: toml>=0.10.2
 Requires-Dist: lancedb>=0.4.0
 Requires-Dist: pyarrow>=14.0.0
-Requires-Dist: sentence-transformers>=2.2.0
 Requires-Dist: tree-sitter>=0.24.0
 Requires-Dist: tree-sitter-python>=0.23.0
 Requires-Dist: tree-sitter-javascript>=0.23.0
@@ -45,9 +44,13 @@ Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
 Requires-Dist: pytest-mock>=3.11.0; extra == "dev"
 Requires-Dist: build>=1.0.0; extra == "dev"
 Requires-Dist: twine>=5.0.0; extra == "dev"
+Provides-Extra: embeddings
+Requires-Dist: torch>=2.0.0; extra == "embeddings"
+Requires-Dist: transformers<5.0.0,>=4.48.0; extra == "embeddings"
 Provides-Extra: all
 Requires-Dist: crewai>=0.80.0; extra == "all"
-Requires-Dist: google-generativeai>=0.5.0; extra == "all"
+Requires-Dist: torch>=2.0.0; extra == "all"
+Requires-Dist: transformers<5.0.0,>=4.48.0; extra == "all"
 Dynamic: license-file
 
 # CodeGraph CLI

{codegraph_cli-2.0.0.dist-info → codegraph_cli-2.1.0.dist-info}/RECORD

@@ -1,33 +1,33 @@
-codegraph_cli/__init__.py,sha256=RbUQDQExUJrOtUpjy4I7oXTwd1VT89usD5EJmSCngcs,78
-codegraph_cli/agents.py,sha256=euxEz8LhHX3K9lpllUnAngU22oEP-ZwxeBDYDLCJgY8,6989
+codegraph_cli/__init__.py,sha256=qTFuIhMU-qKms6nhobwg3YUgDBKR0JenO_3Pq5VgHEk,78
+codegraph_cli/agents.py,sha256=i4VpklF2WLgpS7bmCPcH5lAzohxErZLP5wvssmEK38w,7010
 codegraph_cli/bug_detector.py,sha256=soT4luB5eQx6qrU5rgFCsG44rdo9jRpV0hn-b0f3LPo,16419
 codegraph_cli/chat_agent.py,sha256=dbkEY3zaPJh0ztYaVkCwkTw5zSLGArHkChC_6JWOneg,13685
 codegraph_cli/chat_session.py,sha256=GVey-hnfsa9fa6k2PY1sgy1wtrYSUHKE5cJDV2hG-tg,7038
-codegraph_cli/cli.py,sha256=V3BdUTgGt5nBXFAsJu-qhM3uZx_c9DW4plnXfSKVy6I,10949
-codegraph_cli/cli_chat.py,sha256=O6qujvTMQK3vovy5Di-__UlUuosacbcJX7pPudJPboc,14184
+codegraph_cli/cli.py,sha256=eEzH4TOgyMAFJpVhh2hU0MD2oh61s1hBomeSFx3I3qE,11199
+codegraph_cli/cli_chat.py,sha256=8vk0zrhFQ6MGUa4KomZfnlBXN-Tw-D6aWjqeQPVFxL8,14172
 codegraph_cli/cli_diagnose.py,sha256=gT4qHayC_uWRMsr1Tf92BCFJfRcXAMq8XdEImatrSkU,4260
 codegraph_cli/cli_refactor.py,sha256=_u5RvsF3-KV5C_QnErA4sowlkIAmlxSeLeWKBmSusCI,8176
-codegraph_cli/cli_setup.py,sha256=SC_w0yEHkVBkfXVWx2OaIrQCC5NLouC5ciJ6QAEug1c,17527
+codegraph_cli/cli_setup.py,sha256=f8KdcE0Tf9HQ_ewQm1R_4OZ91bOmi0kuM8eQ05Vs7is,24749
 codegraph_cli/cli_test.py,sha256=ZFPIRhbZ9YYIuSWJyPYLi9PEdHZAI9h8FkWXXRYfqcw,5561
 codegraph_cli/cli_v2.py,sha256=iuw3h5gtvsTg5SdUFXSdLx1Ttiq-oUDM7ZugMqMfETg,9465
 codegraph_cli/codegen_agent.py,sha256=F73YZIIVgE5pOvJsKBl0cv22VW3rP_SGj2viwZS-rqE,9193
-codegraph_cli/config.py,sha256=os7C8-LBGwYFz84wjEnM-48jco-UUKN-I6PY1kufiCE,1072
-codegraph_cli/config_manager.py,sha256=MzZffQC_OpN_Udcm8mhFiWlk_MX5ANW1hKNm5csvrSc,9805
-codegraph_cli/context_manager.py,sha256=Bupg3pAQiDQH1owN4LAcWcpMLOGtwIWjvdzs5Cns61g,16819
+codegraph_cli/config.py,sha256=rOq4lDvqmoly1pfEukzPeCUb76BMqK7cUbzDSFHhsC8,1291
+codegraph_cli/config_manager.py,sha256=K81Ca7jHzHlwxoJsSeRezl8V-iGGJD_IEGE7ZWo3eG0,11422
+codegraph_cli/context_manager.py,sha256=qEKjI7llcLX9y8NFTDs3aiHDm7nDF9jTbhu3tHHOk6w,16824
 codegraph_cli/crew_agents.py,sha256=PKb0skEmxBy2_Ryq67XccmPizLFLgPNs43xOqjSHcGM,6006
 codegraph_cli/crew_chat.py,sha256=tdo8Zf9lOp5-XMdvNntLxo5hKrDcEOe3165cH22-nvQ,6149
 codegraph_cli/crew_tools.py,sha256=wg39mkPDHeE2Wuy3q54xl5d243MHQZBcXCWR2rHrd3M,19339
 codegraph_cli/diff_engine.py,sha256=VGwPG_pZFVz8lGuVHZz_0nhrDocglugw6TumMmnHdTY,8968
-codegraph_cli/embeddings.py,sha256=DfbF4STJXioCU7DXMUQ0LJ3OjLS34ZxiZ5LYcTWny9c,8035
+codegraph_cli/embeddings.py,sha256=YoR6OjiIFC628EnLhNWbw2-_YWqtxSlL--tNWHGsKRk,14611
 codegraph_cli/graph_export.py,sha256=gPyRrOc4_gnW-JaHmmp2pAD60PiZIj_uYA6b0xfU5O0,4562
 codegraph_cli/llm.py,sha256=RpGjJKhUvejmtCHTb9FpGInwPtfaEkHBChBSBTwxUUo,23170
 codegraph_cli/models.py,sha256=o6Wlu8TtWEPDWgq0AhB1xJtxzVfViBMQoCW_4AS29p0,794
 codegraph_cli/models_v2.py,sha256=8zS16hT4SlIahMBwDZ7j4I8fdm3YyWv5qD0urJv1LsI,5521
-codegraph_cli/orchestrator.py,sha256=ARVUpNR2EJIP5jw6CcYe9StqSRmtBAs_2TOIQF8ORwU,1802
+codegraph_cli/orchestrator.py,sha256=AguYRsZ-xu-biM3-uZMhRf6QaoQEqSlGxC7eL3fZqXE,1790
 codegraph_cli/parser.py,sha256=vtKOwirs30O9UxJ6siHzvEWLx4-PxMn5dAfhb42QBG4,29193
 codegraph_cli/performance_analyzer.py,sha256=f9PNMZQ_8jWvzs4osPYgTW2eOsvDytIRmfWWO5DuWCs,10090
 codegraph_cli/project_context.py,sha256=9tSEDEPRmfEQfLcyWXjPa8IGFC1sZI1ysochoxrm4y0,7672
-codegraph_cli/rag.py,sha256=pePkEj8O1Xv1MGHr_KkByiLNNSqmek1nipJx4J_QTx8,7105
+codegraph_cli/rag.py,sha256=DTijL8uZjdEeShQHWFtC_EkoINRgzT3Cr_oHuHHQcfA,7125
 codegraph_cli/refactor_agent.py,sha256=ktQyhUn5YjhbXt7IVgKV7JgSZaT4AivWLLpMlZ7NLXw,16657
 codegraph_cli/security_scanner.py,sha256=rPf8PcYMBllco4PkrxfILJEqKaj1UuEKqCupVycKpo8,15681
 codegraph_cli/storage.py,sha256=XR_w6nJ_ge4r72bfxuuY8Zt8qi8CtHm4EE268EB5kBE,14340
@@ -35,9 +35,9 @@ codegraph_cli/testgen_agent.py,sha256=rqlKbLeEnjfzAZhQUXqLPwFKwRIpiHriTPxVgPCuR_
 codegraph_cli/validation_engine.py,sha256=pzoRH_b06gWfiDZ5Yiecf0SWDWs4oJ66JokggGZZbaw,9029
 codegraph_cli/vector_store.py,sha256=qbIBVDoNOha8JgZwrk7_Jdb7RMYUnBLphJfmqQdrVN4,9912
 codegraph_cli/templates/graph_interactive.html,sha256=PFpU69DbY-Vkcu5UTiqOva_LrZjN2erdz7VXPgNSt6Q,7813
-codegraph_cli-2.0.0.dist-info/licenses/LICENSE,sha256=3PiQTjpJW4DDJz8k5pk-WqX9TrVQD3fNrVNzbTEyW-A,1066
-codegraph_cli-2.0.0.dist-info/METADATA,sha256=UJY9GJXu2iwKVX7V3_Iw8HdfyKEkzBsMyn-DuopVkHg,11024
-codegraph_cli-2.0.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
-codegraph_cli-2.0.0.dist-info/entry_points.txt,sha256=_p5CutxbiWjGVTx9GPeYJ30XOblccdf7SCCNtCkPnaA,45
-codegraph_cli-2.0.0.dist-info/top_level.txt,sha256=XKmdlLsrhdgVW-pN4vzdo-ZTl-9_Rk94SXcM2YRAmHk,14
-codegraph_cli-2.0.0.dist-info/RECORD,,
+codegraph_cli-2.1.0.dist-info/licenses/LICENSE,sha256=3PiQTjpJW4DDJz8k5pk-WqX9TrVQD3fNrVNzbTEyW-A,1066
+codegraph_cli-2.1.0.dist-info/METADATA,sha256=p1GW3UN_gchcHGMyx7_mMT9zZpGkTNbHEcR2DBst6xg,11183
+codegraph_cli-2.1.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+codegraph_cli-2.1.0.dist-info/entry_points.txt,sha256=_p5CutxbiWjGVTx9GPeYJ30XOblccdf7SCCNtCkPnaA,45
+codegraph_cli-2.1.0.dist-info/top_level.txt,sha256=XKmdlLsrhdgVW-pN4vzdo-ZTl-9_Rk94SXcM2YRAmHk,14
+codegraph_cli-2.1.0.dist-info/RECORD,,