wehub-resource-sync
1007 lines
38 KiB
Python
1007 lines
38 KiB
Python
"""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 <Parent>"`` 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]
|