"""Vector embedding support for semantic code search. Supports multiple providers: 1. Local (sentence-transformers) - Private, fast, offline. 2. Google Gemini - High-quality, cloud-based. Requires explicit opt-in. 3. MiniMax (embo-01) - High-quality 1536-dim cloud embeddings. Requires MINIMAX_API_KEY. 4. OpenAI-compatible - Any endpoint speaking OpenAI /v1/embeddings (real OpenAI, Azure OpenAI, self-hosted gateways like new-api / LiteLLM / vLLM / LocalAI / Ollama). """ from __future__ import annotations import hashlib import logging import os import re import sqlite3 import struct import sys import time from abc import ABC, abstractmethod from pathlib import Path from typing import Any from urllib.parse import urlparse from . import __version__ as _crg_version from .graph import GraphNode, GraphStore, node_to_dict logger = logging.getLogger(__name__) # Sent on every cloud-provider HTTP request. Some providers (e.g. Fireworks) # sit behind Cloudflare and reject the urllib default ``Python-urllib/X.Y`` # UA with HTTP 403 / error 1010 ("browser signature banned"). A real UA gets # us through and gives upstream a way to identify CRG-driven traffic. _USER_AGENT = ( f"code-review-graph/{_crg_version} " "(+https://github.com/tirth8205/code-review-graph)" ) # --------------------------------------------------------------------------- # Provider Interface and Implementations # --------------------------------------------------------------------------- class EmbeddingProvider(ABC): @abstractmethod def embed(self, texts: list[str]) -> list[list[float]]: pass @abstractmethod def embed_query(self, text: str) -> list[float]: """Embed a search query (may use a different task type than indexing).""" pass @property @abstractmethod def dimension(self) -> int: pass @property @abstractmethod def name(self) -> str: pass LOCAL_DEFAULT_MODEL = "all-MiniLM-L6-v2" # Process-wide cache of loaded sentence-transformer models, keyed by model name. # Populated by ``prewarm_local_embeddings()`` at server startup (see ``main.main``) # and by ``LocalEmbeddingProvider._get_model`` on first lazy load. Sharing the # loaded model across ``LocalEmbeddingProvider`` instances avoids re-importing # ``sentence_transformers`` + ``torch`` from worker threads, which deadlocks # ``semantic_search_nodes_tool`` on Windows stdio MCP (#385 fixed the peer # tools via ``asyncio.to_thread``; this cache fixes the remaining case where # torch DLL / OpenMP init runs inside an executor thread). _MODEL_CACHE: dict[str, Any] = {} def prewarm_local_embeddings(model_name: str | None = None) -> None: """Eagerly load the local sentence-transformer model on the calling thread. Call this from the **main thread** before entering an asyncio event loop (e.g. before ``mcp.run()``) on Windows to prevent a deadlock where lazy- loading ``sentence_transformers`` + ``torch`` inside a FastMCP executor worker thread blocks indefinitely on DLL init / OpenMP thread-pool registration. No-op when ``sentence-transformers`` is not installed (cloud-provider setups remain unaffected) or when the configured model is already cached. Args: model_name: Optional override; falls back to the ``CRG_EMBEDDING_MODEL`` environment variable and then to ``LOCAL_DEFAULT_MODEL``. """ try: from sentence_transformers import SentenceTransformer # noqa: F401 except ImportError: return # cloud-only setup: nothing to pre-warm resolved = model_name or os.environ.get( "CRG_EMBEDDING_MODEL", LOCAL_DEFAULT_MODEL ) if resolved in _MODEL_CACHE: return try: _MODEL_CACHE[resolved] = LocalEmbeddingProvider(resolved)._get_model() except Exception as exc: # pragma: no cover — best-effort startup hook logger.warning("prewarm_local_embeddings(%s) skipped: %s", resolved, exc) class LocalEmbeddingProvider(EmbeddingProvider): def __init__(self, model_name: str | None = None) -> None: self._model_name = model_name or os.environ.get( "CRG_EMBEDDING_MODEL", LOCAL_DEFAULT_MODEL ) self._model = None # Lazy-loaded def _get_model(self): if self._model is None: # Check the process-wide cache first — populated either by a prior # provider instance or by ``prewarm_local_embeddings`` at startup. cached = _MODEL_CACHE.get(self._model_name) if cached is not None: self._model = cached return self._model try: from sentence_transformers import SentenceTransformer # Check environment variable, default to False to prevent RCE _rce_val = os.environ.get("CRG_ALLOW_REMOTE_CODE", "0") allow_remote_code = _rce_val.lower() in ("1", "true", "yes") self._model = SentenceTransformer( self._model_name, trust_remote_code=allow_remote_code, ) _MODEL_CACHE[self._model_name] = self._model except ImportError: raise ImportError( "sentence-transformers not installed. " "Run: pip install code-review-graph[embeddings]" ) return self._model def embed(self, texts: list[str]) -> list[list[float]]: model = self._get_model() vectors = model.encode(texts, show_progress_bar=False) return [v.tolist() for v in vectors] def embed_query(self, text: str) -> list[float]: return self.embed([text])[0] @property def dimension(self) -> int: model = self._get_model() return model.get_sentence_embedding_dimension() @property def name(self) -> str: return f"local:{self._model_name}" class GoogleEmbeddingProvider(EmbeddingProvider): def __init__(self, api_key: str, model: str = "gemini-embedding-001") -> None: try: from google import genai self._client = genai.Client(api_key=api_key) self.model = model self._dimension: int | None = None except ImportError: raise ImportError( "google-generativeai not installed. " "Run: pip install code-review-graph[google-embeddings]" ) def embed(self, texts: list[str]) -> list[list[float]]: batch_size = 100 results = [] for i in range(0, len(texts), batch_size): batch = texts[i:i + batch_size] response = self._call_with_retry( lambda b=batch: self._client.models.embed_content( model=self.model, contents=b, config={"task_type": "RETRIEVAL_DOCUMENT"}, ) ) results.extend([e.values for e in response.embeddings]) if self._dimension is None and results: self._dimension = len(results[0]) return results @staticmethod def _call_with_retry(fn, max_retries: int = 3): """Call fn with exponential backoff on transient API errors.""" for attempt in range(max_retries): try: return fn() except Exception as e: # Retry on rate-limit (429) or server errors (5xx) err_str = str(e) is_retryable = "429" in err_str or "500" in err_str or "503" in err_str if not is_retryable or attempt == max_retries - 1: raise wait = 2 ** attempt logger.warning("Gemini API error (attempt %d/%d), retrying in %ds: %s", attempt + 1, max_retries, wait, e) time.sleep(wait) def embed_query(self, text: str) -> list[float]: response = self._call_with_retry( lambda: self._client.models.embed_content( model=self.model, contents=[text], config={"task_type": "RETRIEVAL_QUERY"}, ) ) vec = response.embeddings[0].values if self._dimension is None: self._dimension = len(vec) return vec @property def dimension(self) -> int: if self._dimension is not None: return self._dimension # Default for gemini-embedding-001; updated dynamically after first call return 768 @property def name(self) -> str: return f"google:{self.model}" class MiniMaxEmbeddingProvider(EmbeddingProvider): """MiniMax embo-01 embedding provider (1536 dimensions). Uses the MiniMax Embeddings API (https://api.minimax.io/v1/embeddings) with the embo-01 model. Requires the MINIMAX_API_KEY environment variable. """ _ENDPOINT = "https://api.minimax.io/v1/embeddings" _MODEL = "embo-01" _DIMENSION = 1536 def __init__(self, api_key: str) -> None: self._api_key = api_key def _call_api(self, texts: list[str], task_type: str) -> list[list[float]]: import json as _json import urllib.request payload = _json.dumps({ "model": self._MODEL, "texts": texts, "type": task_type, }).encode("utf-8") req = urllib.request.Request( self._ENDPOINT, data=payload, headers={ "Content-Type": "application/json", "Authorization": f"Bearer {self._api_key}", "User-Agent": _USER_AGENT, "Accept": "application/json", }, ) max_retries = 3 for attempt in range(max_retries): try: import ssl _ssl_ctx = ssl.create_default_context() with urllib.request.urlopen(req, timeout=60, context=_ssl_ctx) as resp: # nosec B310 body = _json.loads(resp.read().decode("utf-8")) base_resp = body.get("base_resp", {}) if base_resp.get("status_code", 0) != 0: raise RuntimeError( f"MiniMax API error: {base_resp.get('status_msg', 'unknown')}" ) return body["vectors"] except Exception as e: err_str = str(e) is_retryable = "429" in err_str or "500" in err_str or "503" in err_str if not is_retryable or attempt == max_retries - 1: raise wait = 2 ** attempt logger.warning( "MiniMax API error (attempt %d/%d), retrying in %ds: %s", attempt + 1, max_retries, wait, e, ) time.sleep(wait) return [] # unreachable, but keeps mypy happy def embed(self, texts: list[str]) -> list[list[float]]: batch_size = 100 results: list[list[float]] = [] for i in range(0, len(texts), batch_size): batch = texts[i:i + batch_size] results.extend(self._call_api(batch, "db")) return results def embed_query(self, text: str) -> list[float]: return self._call_api([text], "query")[0] @property def dimension(self) -> int: return self._DIMENSION @property def name(self) -> str: return f"minimax:{self._MODEL}" class OpenAIEmbeddingProvider(EmbeddingProvider): """OpenAI-compatible embedding provider. Works with any endpoint that speaks the OpenAI ``/v1/embeddings`` schema: - Real OpenAI API (``https://api.openai.com/v1``) - Azure OpenAI - Self-hosted gateways: new-api, LiteLLM, vLLM, LocalAI, Ollama (openai mode) Provider identity in ``name`` includes both the model and the endpoint host (``openai:{model}@{host}``), so switching base URL while keeping the same model ID re-partitions the embeddings table and forces a clean re-embed. This is the only defense against silently mixing vector spaces from different backends (e.g. real OpenAI vs. an OpenAI-compatible gateway that ships different weights under the same model name). Dimension is detected from the first response and frozen; switching the ``model`` in the environment also changes ``provider.name`` and triggers re-embed via the same isolation key. """ _DEFAULT_BATCH_SIZE = 100 # Default ports by scheme; stripped from the host_key so the user can't # accidentally force a re-embed by toggling an explicit default port. _DEFAULT_PORTS = {"http": 80, "https": 443} def __init__( self, api_key: str, base_url: str, model: str, dimension: int | None = None, timeout: int = 120, batch_size: int | None = None, ) -> None: self._api_key = api_key self._base_url = base_url.rstrip("/") self._model = model self._dimension = dimension self._timeout = timeout self._batch_size = batch_size or self._DEFAULT_BATCH_SIZE self._host_key = self._make_host_key(self._base_url) @classmethod def _make_host_key(cls, base_url: str) -> str: """Normalize the identity key used in ``provider.name``. Codex review pushed this well past naive ``netloc`` because that alone has three leaks: 1. ``netloc`` preserves ``userinfo`` (``user:pass@host``) — we'd persist credentials into the DB's ``embeddings.provider`` column. Use ``hostname`` instead. 2. Default ports (``:80`` for http, ``:443`` for https) are semantically identical to omitting the port; keeping them would cause spurious re-embeds when the user just spelled the URL differently. 3. Path is part of the backend identity for path-routed gateways: ``https://gw/openai/v1`` and ``https://gw/vendor-b/v1`` front different models and must not share cached vectors. """ parsed = urlparse(base_url) hostname = (parsed.hostname or "").lower() scheme = (parsed.scheme or "").lower() port = parsed.port if port and port != cls._DEFAULT_PORTS.get(scheme): # Bracket IPv6 literals when appending a port. host_part = f"[{hostname}]:{port}" if ":" in hostname else f"{hostname}:{port}" else: host_part = hostname # Preserve path routing. Trim any trailing slash and any # ``/embeddings`` suffix that callers may have included — we append # that ourselves when building the request URL. path = (parsed.path or "").rstrip("/") if path.endswith("/embeddings"): path = path[: -len("/embeddings")].rstrip("/") # Include scheme: http and https to the same host+path front # different endpoints in practice (plaintext vs TLS, dev vs prod # gateway), and sharing cached vectors across them is the same # silent-mixing failure mode as switching base URL entirely. return f"{scheme}://{host_part}{path}" if path else f"{scheme}://{host_part}" def _call_api(self, texts: list[str]) -> list[list[float]]: import http.client import json as _json import socket import ssl import urllib.error import urllib.request body: dict[str, Any] = {"model": self._model, "input": texts} # OpenAI v3 models (text-embedding-3-*) support dimension reduction; # only forward the param when the user explicitly pinned one. if self._dimension is not None: body["dimensions"] = self._dimension payload = _json.dumps(body).encode("utf-8") req = urllib.request.Request( f"{self._base_url}/embeddings", data=payload, headers={ "Content-Type": "application/json", "Authorization": f"Bearer {self._api_key}", "User-Agent": _USER_AGENT, "Accept": "application/json", }, ) max_retries = 3 for attempt in range(max_retries): try: _ssl_ctx = ssl.create_default_context() try: with urllib.request.urlopen( # nosec B310 req, timeout=self._timeout, context=_ssl_ctx, ) as resp: raw = resp.read().decode("utf-8") except urllib.error.HTTPError as http_err: # 429 / 5xx: re-raise and let the outer retry loop handle it. # (We must not convert to RuntimeError here or retry below # can't tell it was a transient HTTP failure.) if http_err.code == 429 or 500 <= http_err.code < 600: raise # Other 4xx: surface the API error body instead of a bare # "400 Bad Request" — gateways like new-api return JSON # with the real reason (batch size limits, invalid model, # etc.) which is far more actionable. try: err_body = http_err.read().decode("utf-8", errors="replace") except Exception: err_body = "" err_msg = err_body or str(http_err) try: parsed = _json.loads(err_body) if isinstance(parsed, dict) and "error" in parsed: err_obj = parsed["error"] err_msg = ( err_obj.get("message", err_msg) if isinstance(err_obj, dict) else str(err_obj) ) except Exception: # nosec B110 # Non-JSON error body is fine: we already seeded # err_msg with the raw body above, so fall through. pass raise RuntimeError( f"OpenAI API HTTP {http_err.code}: {err_msg}" ) from http_err response = _json.loads(raw) if "error" in response: err = response["error"] msg = err.get("message", "unknown") if isinstance(err, dict) else str(err) raise RuntimeError(f"OpenAI API error: {msg}") data = response.get("data", []) if not data: raise RuntimeError("OpenAI API returned empty data") # OpenAI spec: data[i].index maps to input[i], but some # compatible gateways re-order results or drop entries on # partial failure, and others omit `index` entirely. Three # disjoint cases: # 1. All items have a valid int ``index``: must form a # permutation of 0..N-1, then sort and use. # 2. NO item carries an ``index`` field: trust server # order, only verify count matches. # 3. Anything in between (partial indices, str indices, # missing on some): refuse. Zipping server order in # that case would happily misalign the indexed items. any_has_index = any("index" in item for item in data) all_int_index = all( isinstance(item.get("index"), int) for item in data ) if all_int_index: expected = set(range(len(texts))) indices = [int(item["index"]) for item in data] if len(set(indices)) != len(indices) or set(indices) != expected: raise RuntimeError( "OpenAI API returned malformed indices " f"(got {indices}, expected permutation of " f"0..{len(texts) - 1}) — refusing to misalign vectors." ) data = sorted(data, key=lambda item: int(item["index"])) elif not any_has_index: if len(data) != len(texts): raise RuntimeError( f"OpenAI API returned {len(data)} embeddings for " f"{len(texts)} inputs with no index field — " "refusing to misalign vectors." ) else: # Mixed: some items have index, others don't (or carry # non-int index). Server order would silently misplace # the indexed items, so we refuse. raise RuntimeError( "OpenAI API returned mixed indexed/unindexed data — " "refusing to misalign vectors." ) vectors = [item["embedding"] for item in data] if vectors and self._dimension is None: self._dimension = len(vectors[0]) return vectors except Exception as e: # Retryable = HTTP 429/5xx, network/timeout/TLS issues. # Non-retryable = HTTP 4xx (other), malformed responses, # misaligned data length — those are caller-side bugs that # will keep failing on retry. is_retryable = False if isinstance(e, urllib.error.HTTPError): is_retryable = e.code == 429 or 500 <= e.code < 600 elif isinstance(e, ( urllib.error.URLError, socket.timeout, TimeoutError, ConnectionError, ssl.SSLError, # Reverse proxies and edge gateways surface transient # disconnects as these stdlib classes. Real incidents # have been observed on Cloudflare-fronted endpoints # and on LiteLLM when upstream providers hiccup. http.client.IncompleteRead, http.client.BadStatusLine, http.client.RemoteDisconnected, )): is_retryable = True if not is_retryable or attempt == max_retries - 1: raise wait = 2 ** attempt logger.warning( "OpenAI embeddings API error (attempt %d/%d), retrying in %ds: %s", attempt + 1, max_retries, wait, e, ) time.sleep(wait) return [] # unreachable def embed(self, texts: list[str]) -> list[list[float]]: if not texts: return [] results: list[list[float]] = [] for i in range(0, len(texts), self._batch_size): results.extend(self._call_api(texts[i:i + self._batch_size])) return results def embed_query(self, text: str) -> list[float]: return self._call_api([text])[0] @property def dimension(self) -> int: if self._dimension is not None: return self._dimension # Default for text-embedding-3-small; updated after first call. return 1536 @property def name(self) -> str: # Endpoint-aware identity: model alone is NOT enough — two backends # can serve the same model ID with different weights or dimensions, # and re-using cached embeddings across them silently corrupts # semantic ranking. Including the host partitions the embeddings # table so switching CRG_OPENAI_BASE_URL triggers a safe re-embed. return f"openai:{self._model}@{self._host_key}" CLOUD_PROVIDERS = {"google", "minimax", "openai"} def _is_localhost_url(url: str) -> bool: """Return True if url points to a localhost host (never treat as cloud egress). Uses urlparse.hostname so we compare the actual host, not a substring match that could be fooled by e.g. ``https://my-openai.127.0.0.1.nip.io``. """ try: host = (urlparse(url).hostname or "").lower() except Exception: return False # nosec B104: we're *matching* a URL hostname, not binding a listener. return host in {"127.0.0.1", "localhost", "0.0.0.0", "::1"} # nosec B104 def _warn_cloud_egress(provider_name: str) -> None: """Print a stderr warning before a cloud embedding provider is used. The warning is suppressed when ``CRG_ACCEPT_CLOUD_EMBEDDINGS=1`` is set in the environment, so scripted / CI workloads can acknowledge once and move on. Use stderr (never stdin/input) to stay compatible with the MCP stdio transport — anything we write to stdout would corrupt the JSON-RPC stream. See: #174 """ if os.environ.get("CRG_ACCEPT_CLOUD_EMBEDDINGS", "").strip() == "1": return print( f"\n⚠️ code-review-graph: about to embed code via the '{provider_name}' " "cloud provider.\n" " Your source code (function names, docstrings, file paths) will be " "sent to an external API.\n" " This is necessary for semantic search with the cloud provider you " "selected.\n" " To skip this warning in future runs, set " "CRG_ACCEPT_CLOUD_EMBEDDINGS=1 in your environment.\n" " To stay fully offline, use the default 'local' provider instead " "(no API key needed).\n", file=sys.stderr, ) _VALID_PROVIDERS = {"local", "openai", "google", "minimax"} def get_provider( provider: str | None = None, model: str | None = None, ) -> EmbeddingProvider | None: """Get an embedding provider by name. Args: provider: Provider name. One of "local", "google", "minimax", "openai", or None for local. Names are case-insensitive and surrounding whitespace is ignored; unknown names raise ValueError instead of silently falling back to the local provider. Google requires GOOGLE_API_KEY env var and explicit opt-in. MiniMax requires MINIMAX_API_KEY env var and explicit opt-in. OpenAI requires CRG_OPENAI_API_KEY + CRG_OPENAI_BASE_URL + CRG_OPENAI_MODEL env vars (or the ``model`` arg). The egress warning is skipped when the base URL points to localhost. Cloud providers emit a one-time stderr warning before use unless ``CRG_ACCEPT_CLOUD_EMBEDDINGS=1`` is set. See: #174 model: Model name/path to use. For local provider this is any sentence-transformers compatible model. Falls back to CRG_EMBEDDING_MODEL env var, then to all-MiniLM-L6-v2. For Google provider this is a Gemini model ID. For OpenAI provider this overrides CRG_OPENAI_MODEL. Raises: ValueError: If the provider name is not one of the known providers, or if required environment variables are missing. """ name = provider.strip().lower() if provider else "" if name and name not in _VALID_PROVIDERS: raise ValueError( f"Unknown embedding provider '{name}'. " "Valid: local, openai, google, minimax" ) if name == "openai": api_key = os.environ.get("CRG_OPENAI_API_KEY") base_url = os.environ.get("CRG_OPENAI_BASE_URL") resolved_model = model or os.environ.get("CRG_OPENAI_MODEL") if not api_key or not base_url or not resolved_model: missing = [ name for name, val in [ ("CRG_OPENAI_API_KEY", api_key), ("CRG_OPENAI_BASE_URL", base_url), ("CRG_OPENAI_MODEL", resolved_model), ] if not val ] raise ValueError( "Missing required environment variable(s) for the OpenAI " f"embedding provider: {', '.join(missing)}." ) dim_env = os.environ.get("CRG_OPENAI_DIMENSION") dimension = int(dim_env) if dim_env else None batch_env = os.environ.get("CRG_OPENAI_BATCH_SIZE") batch_size = int(batch_env) if batch_env else None if not _is_localhost_url(base_url): _warn_cloud_egress("openai") return OpenAIEmbeddingProvider( api_key=api_key, base_url=base_url, model=resolved_model, dimension=dimension, batch_size=batch_size, ) if name == "minimax": api_key = os.environ.get("MINIMAX_API_KEY") if not api_key: raise ValueError( "MINIMAX_API_KEY environment variable is required for " "the MiniMax embedding provider." ) _warn_cloud_egress("minimax") return MiniMaxEmbeddingProvider(api_key=api_key) if name == "google": api_key = os.environ.get("GOOGLE_API_KEY") if not api_key: raise ValueError( "GOOGLE_API_KEY environment variable is required for " "the Google embedding provider." ) _warn_cloud_egress("google") try: return GoogleEmbeddingProvider( api_key=api_key, **({"model": model} if model else {}), ) except ImportError: return None # Default: local if not _check_available(): return None try: return LocalEmbeddingProvider(model_name=model) except ImportError: return None def _check_available() -> bool: """Check whether local embedding support is available.""" try: import sentence_transformers # noqa: F401 return True except ImportError: return False # --------------------------------------------------------------------------- # SQLite vector storage # --------------------------------------------------------------------------- _EMBEDDINGS_SCHEMA = """ CREATE TABLE IF NOT EXISTS embeddings ( qualified_name TEXT PRIMARY KEY, vector BLOB NOT NULL, text_hash TEXT NOT NULL, provider TEXT NOT NULL DEFAULT 'unknown' ); """ def _encode_vector(vec: list[float]) -> bytes: """Encode a float vector as a compact binary blob.""" return struct.pack(f"{len(vec)}f", *vec) def _decode_vector(blob: bytes) -> list[float]: """Decode a binary blob back to a float vector.""" n = len(blob) // 4 # 4 bytes per float32 return list(struct.unpack(f"{n}f", blob)) def _cosine_similarity(a: list[float], b: list[float]) -> float: """Compute cosine similarity between two vectors.""" if len(a) != len(b): return 0.0 dot = sum(x * y for x, y in zip(a, b)) norm_a = sum(x * x for x in a) ** 0.5 norm_b = sum(x * x for x in b) ** 0.5 if norm_a == 0 or norm_b == 0: return 0.0 return dot / (norm_a * norm_b) _IDENTIFIER_SPLIT_RE = re.compile(r"([a-z])([A-Z])|[_./\-]+") def _split_identifier(name: str) -> str: """Split snake_case / camelCase / PascalCase / dotted into space-separated words. Examples: get_route_handler -> "get route handler" APIRoute -> "API Route" dispatch_request -> "dispatch request" full_dispatch_request -> "full dispatch request" """ if not name: return "" # Insert space between lowercase->uppercase transitions, then collapse # snake_case / dotted / hyphenated separators. spaced = re.sub(r"([a-z])([A-Z])", r"\1 \2", name) spaced = re.sub(r"[_./\-]+", " ", spaced) return " ".join(spaced.split()) def _node_to_text(node: GraphNode) -> str: """Convert a node to a searchable text representation. Designed so natural-language queries land on the right node, not just on the enclosing class. We include the dotted ``Parent.name`` form, the identifier split into words, an explicit ``"in "`` phrase, the enclosing module directory, and the language. Tested by the ``multi_hop_retrieval`` benchmark — see ``docs/REPRODUCING.md``. """ parts: list[str] = [] # 1. Dotted form first — strongest lexical signal for "method in class" if node.parent_name and node.kind != "File": parts.append(f"{node.parent_name}.{node.name}") # 2. Bare name (always present) parts.append(node.name) # 3. Split-words form of the name (only if it differs from the bare name) name_split = _split_identifier(node.name) if name_split and name_split.lower() != node.name.lower(): parts.append(name_split) # 4. Kind ("function", "class", "test", ...) if node.kind != "File": parts.append(node.kind.lower()) # 5. Parent context with the split form too if node.parent_name: parts.append(f"in {node.parent_name}") parent_split = _split_identifier(node.parent_name) if parent_split and parent_split.lower() != node.parent_name.lower(): parts.append(parent_split) # 6. Signature bits if node.params: parts.append(node.params) if node.return_type: parts.append(f"returns {node.return_type}") # 7. Module / directory context from the file path — gives queries a # term like "routing" or "client" to anchor against. if node.file_path: parent_dir = Path(node.file_path).parent.name if parent_dir and parent_dir not in (".", "src", "lib"): parts.append(parent_dir) # 8. Language if node.language: parts.append(node.language) return " ".join(parts) class EmbeddingStore: """Manages vector embeddings for graph nodes in SQLite.""" def __init__( self, db_path: str | Path, provider: str | None = None, model: str | None = None, ) -> None: self.provider = get_provider(provider, model=model) self.available = self.provider is not None self.db_path = Path(db_path) self._conn = sqlite3.connect( str(self.db_path), timeout=30, check_same_thread=False, isolation_level=None, ) self._conn.row_factory = sqlite3.Row self._conn.executescript(_EMBEDDINGS_SCHEMA) # Migration for existing DBs missing the provider column try: self._conn.execute("SELECT provider FROM embeddings LIMIT 1") except sqlite3.OperationalError: self._conn.execute( "ALTER TABLE embeddings ADD COLUMN provider " "TEXT NOT NULL DEFAULT 'unknown'" ) self._conn.commit() def __enter__(self) -> "EmbeddingStore": return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: # type: ignore[no-untyped-def] self.close() def close(self) -> None: self._conn.close() def embed_nodes(self, nodes: list[GraphNode], batch_size: int = 64) -> int: """Compute and store embeddings for a list of nodes.""" if not self.provider: return 0 # Filter to nodes that need embedding to_embed: list[tuple[GraphNode, str, str]] = [] provider_name = self.provider.name for node in nodes: if node.kind == "File": continue text = _node_to_text(node) text_hash = hashlib.sha256(text.encode()).hexdigest() existing = self._conn.execute( "SELECT text_hash, provider FROM embeddings WHERE qualified_name = ?", (node.qualified_name,), ).fetchone() # Re-embed if text changed OR provider changed if (existing and existing["text_hash"] == text_hash and existing["provider"] == provider_name): continue to_embed.append((node, text, text_hash)) if not to_embed: return 0 # Encode in batches texts = [t for _, t, _ in to_embed] vectors = self.provider.embed(texts) for (node, _text, text_hash), vec in zip(to_embed, vectors): blob = _encode_vector(vec) self._conn.execute( """INSERT OR REPLACE INTO embeddings (qualified_name, vector, text_hash, provider) VALUES (?, ?, ?, ?)""", (node.qualified_name, blob, text_hash, provider_name), ) self._conn.commit() return len(to_embed) def search(self, query: str, limit: int = 20) -> list[tuple[str, float]]: """Search for nodes by semantic similarity.""" if not self.provider: return [] provider_name = self.provider.name query_vec = self.provider.embed_query(query) # Process in chunks, only matching current provider scored: list[tuple[str, float]] = [] cursor = self._conn.execute( "SELECT qualified_name, vector FROM embeddings WHERE provider = ?", (provider_name,), ) chunk_size = 500 while True: rows = cursor.fetchmany(chunk_size) if not rows: break for row in rows: vec = _decode_vector(row["vector"]) sim = _cosine_similarity(query_vec, vec) scored.append((row["qualified_name"], sim)) scored.sort(key=lambda x: x[1], reverse=True) return scored[:limit] def remove_node(self, qualified_name: str) -> None: self._conn.execute( "DELETE FROM embeddings WHERE qualified_name = ?", (qualified_name,) ) self._conn.commit() def count(self) -> int: return self._conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0] def embed_all_nodes(graph_store: GraphStore, embedding_store: EmbeddingStore) -> int: """Embed all non-file nodes in the graph.""" if not embedding_store.available: return 0 all_files = graph_store.get_all_files() all_nodes: list[GraphNode] = [] for f in all_files: all_nodes.extend(graph_store.get_nodes_by_file(f)) return embedding_store.embed_nodes(all_nodes) def semantic_search( query: str, graph_store: GraphStore, embedding_store: EmbeddingStore, limit: int = 20, ) -> list[dict[str, Any]]: """Search nodes using vector similarity, falling back to keyword search.""" if embedding_store.available and embedding_store.count() > 0: results = embedding_store.search(query, limit=limit) output = [] for qn, score in results: node = graph_store.get_node(qn) if node: d = node_to_dict(node) d["similarity_score"] = round(score, 4) output.append(d) return output # Fallback to keyword search nodes = graph_store.search_nodes(query, limit=limit) return [node_to_dict(n) for n in nodes]