codegraph-network 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
codegraph/__init__.py ADDED
@@ -0,0 +1 @@
1
+ __version__ = "0.1.0"
codegraph/bench.py ADDED
@@ -0,0 +1,93 @@
1
+ """`codegraph bench`: does the graph answer real tasks cheaper than reading files, and does semantic
2
+ search find symbols that lexical search misses?
3
+
4
+ Tasks file (YAML, or JSON which is a subset of it):
5
+
6
+ - task: "where do we check a user's credentials on login"
7
+ expect: [authenticate_user, auth.py:Session.open]
8
+
9
+ `expect` entries are anything `context()` accepts (name, qualname or file:name). Per task we measure
10
+ recall share of expected symbols inside the top-k results, lexical-only vs hybrid
11
+ tokens what an agent reads with the graph (the hybrid search output plus context() of every
12
+ expected symbol it found) vs reading every file that contains an expected symbol
13
+ """
14
+ from __future__ import annotations
15
+
16
+ import json
17
+ from dataclasses import dataclass, field
18
+ from pathlib import Path
19
+
20
+ from .indexer import index_repo
21
+ from .queries import Graph, toks
22
+
23
+
24
+ @dataclass
25
+ class TaskResult:
26
+ task: str
27
+ expected: int
28
+ unresolved: list[str] = field(default_factory=list)
29
+ lexical_found: int = 0
30
+ hybrid_found: int = 0
31
+ graph_tokens: int = 0
32
+ file_tokens: int = 0
33
+
34
+
35
+ def load_tasks(path: str | Path) -> list[dict]:
36
+ text = Path(path).read_text()
37
+ try:
38
+ import yaml
39
+ data = yaml.safe_load(text)
40
+ except ImportError:
41
+ try:
42
+ data = json.loads(text)
43
+ except ValueError as e:
44
+ raise SystemExit("bench: reading YAML needs PyYAML (pip install 'codegraph[bench]'); "
45
+ f"as JSON the file failed to parse: {e}") from None
46
+ if not isinstance(data, list) or not all(isinstance(t, dict) and "task" in t for t in data):
47
+ raise SystemExit("bench: the tasks file must be a list of {task: ..., expect: [symbols]}")
48
+ return data
49
+
50
+
51
+ def run_task(g: Graph, spec: dict, k: int) -> TaskResult:
52
+ expect = [str(e) for e in spec.get("expect") or []]
53
+ res = TaskResult(str(spec["task"]), len(expect))
54
+ wanted: dict[str, set[int]] = {}
55
+ for e in expect:
56
+ rows = g._find(e)
57
+ if rows:
58
+ wanted[e] = {r["id"] for r in rows}
59
+ else:
60
+ res.unresolved.append(e)
61
+ lex_ids = {r["id"] for r in g._search_rows(res.task, k)}
62
+ hyb = g._hybrid_rows(res.task, k)
63
+ hyb_ids = {r["id"] for r, _ in hyb}
64
+ res.lexical_found = sum(1 for ids in wanted.values() if ids & lex_ids)
65
+ found = [e for e, ids in wanted.items() if ids & hyb_ids]
66
+ res.hybrid_found = len(found)
67
+ res.graph_tokens = toks(g.search(res.task, k)) + sum(toks(g.context(e)) for e in found)
68
+ files = {r["file"] for ids in wanted.values() for r in g._rows_by_id(ids)}
69
+ res.file_tokens = g._file_tokens(files)
70
+ return res
71
+
72
+
73
+ def run(root: str | Path, tasks_path: str | Path, k: int = 10) -> str:
74
+ index_repo(root)
75
+ g = Graph(root)
76
+ try:
77
+ semantic = g._semantic() is not None
78
+ results = [run_task(g, t, k) for t in load_tasks(tasks_path)]
79
+ finally:
80
+ g.close()
81
+ out = [f"codegraph bench top-{k} semantic index: {'yes' if semantic else 'NO (hybrid == lexical)'}", ""]
82
+ for r in results:
83
+ out.append(f"- {r.task}")
84
+ out.append(f" recall lexical {r.lexical_found}/{r.expected} hybrid {r.hybrid_found}/{r.expected}"
85
+ f" tokens graph ~{r.graph_tokens} vs files ~{r.file_tokens}")
86
+ if r.unresolved:
87
+ out.append(f" not in index: {', '.join(r.unresolved)}")
88
+ exp = sum(r.expected - len(r.unresolved) for r in results) or 1
89
+ lex, hyb = sum(r.lexical_found for r in results), sum(r.hybrid_found for r in results)
90
+ gt, ft = sum(r.graph_tokens for r in results), sum(r.file_tokens for r in results)
91
+ out += ["", f"recall lexical {lex}/{exp} ({lex / exp:.0%}) hybrid {hyb}/{exp} ({hyb / exp:.0%})",
92
+ f"tokens graph ~{gt} vs reading files ~{ft}" + (f" ({ft / gt:.1f}x fewer)" if gt else "")]
93
+ return "\n".join(out)
codegraph/cli.py ADDED
@@ -0,0 +1,128 @@
1
+ """Command line front end. Argument parsing is derived from the shared tool registry."""
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ import sys
7
+
8
+ from .indexer import index_repo
9
+ from .install import PLATFORMS, hook_install, hook_uninstall, install, uninstall
10
+ from .queries import Graph
11
+ from .tools import CLI_ALIASES, CLI_FLAGS, TOOLS
12
+
13
+ # 0 answered, 3 the name matches several symbols, 4 it matches none: scripts can branch without parsing text
14
+ EXIT_CODES = {"ok": 0, "ambiguous": 3, "missing": 4}
15
+
16
+ # the reindex tool is presented under a friendlier name
17
+ CLI_NAME = {"reindex": "index"}
18
+
19
+
20
+ def _cli_name(tool_name: str) -> str:
21
+ return CLI_NAME.get(tool_name, tool_name)
22
+
23
+
24
+ # every spelling a tool answers to on the command line
25
+ CLI_SPELLINGS = {t.name: (_cli_name(t.name), *CLI_ALIASES.get(t.name, ())) for t in TOOLS}
26
+
27
+
28
+ def _cli_name(tool_name: str) -> str:
29
+ return CLI_NAME.get(tool_name, tool_name)
30
+
31
+
32
+ def _tool_for(spelling: str):
33
+ for tool in TOOLS:
34
+ if spelling in CLI_SPELLINGS[tool.name]:
35
+ return tool
36
+ raise KeyError(spelling)
37
+
38
+
39
+ def build_parser() -> argparse.ArgumentParser:
40
+ ap = argparse.ArgumentParser(
41
+ prog="codegraph",
42
+ description="Graph-based code understanding for LLMs",
43
+ epilog="run 'codegraph <tool> --help' for the arguments of one tool")
44
+ ap.add_argument("-r", "--root", default=".", help="repository root (default: cwd)")
45
+ sub = ap.add_subparsers(dest="cmd", required=True)
46
+ sub.add_parser("serve", help="run the MCP server over stdio")
47
+ bench = sub.add_parser("bench", help="measure tokens and recall (lexical vs hybrid) on a tasks file")
48
+ bench.add_argument("tasks", help="YAML/JSON list of {task: ..., expect: [symbols]}")
49
+ bench.add_argument("-k", type=int, default=10, help="results per search (default 10)")
50
+
51
+ inst = sub.add_parser("install", help="wire codegraph into coding assistants (skill, MCP config, guidance)")
52
+ inst.add_argument("platform", nargs="?", default="all", choices=(*PLATFORMS, "all"),
53
+ help="which assistant to configure (default: all)")
54
+ inst.add_argument("--no-hooks", action="store_true", help="skip the Claude Code Glob|Grep hint hook")
55
+ inst.add_argument("--uninstall", action="store_true", help="remove what install wrote")
56
+ hk = sub.add_parser("hook", help="keep the index fresh from git: post-commit/checkout/merge hooks")
57
+ hk.add_argument("action", choices=("install", "uninstall"))
58
+
59
+ for tool in TOOLS:
60
+ spellings = CLI_SPELLINGS[tool.name]
61
+ p = sub.add_parser(spellings[0], aliases=spellings[1:],
62
+ help=tool.summary, description=tool.summary)
63
+ for param in tool.params:
64
+ if param.default is None:
65
+ p.add_argument("value", help=param.help or param.name)
66
+ continue
67
+ flag, short = CLI_FLAGS.get(param.name, (param.name, None))
68
+ opts = [f"--{flag}"] + ([short] if short else [])
69
+ if param.type == "bool":
70
+ p.add_argument(*opts, action="store_true", help=param.help)
71
+ elif param.type == "int":
72
+ p.add_argument(*opts, type=int, default=param.default, help=param.help)
73
+ else:
74
+ p.add_argument(*opts, default=param.default, help=param.help)
75
+ return ap
76
+
77
+
78
+ def _kwargs(tool, args) -> dict:
79
+ out = {}
80
+ for param in tool.params:
81
+ if param.default is None:
82
+ out[param.name] = args.value
83
+ continue
84
+ flag = CLI_FLAGS.get(param.name, (param.name, None))[0]
85
+ v = getattr(args, flag.replace("-", "_"))
86
+ out[param.name] = int(v) if param.type == "int" and v is not None else v
87
+ return out
88
+
89
+
90
+ def main(argv=None):
91
+ args = build_parser().parse_args(argv)
92
+
93
+ if args.cmd == "serve":
94
+ from .server import main as serve
95
+ return serve(args.root)
96
+
97
+ if args.cmd == "bench":
98
+ from .bench import run as bench
99
+ print(bench(args.root, args.tasks, args.k))
100
+ return None
101
+
102
+ if args.cmd in ("install", "hook"):
103
+ if args.cmd == "hook":
104
+ lines = (hook_install if args.action == "install" else hook_uninstall)(args.root)
105
+ elif args.uninstall:
106
+ lines = uninstall(args.root)
107
+ else:
108
+ lines = install(args.root, args.platform, hooks=not args.no_hooks)
109
+ print("\n".join(lines))
110
+ return None
111
+
112
+ if args.cmd == "index":
113
+ print(json.dumps(index_repo(args.root, args.force), indent=2, sort_keys=True))
114
+ return None
115
+
116
+ index_repo(args.root) # cheap when nothing changed
117
+ tool = _tool_for(args.cmd)
118
+ graph = Graph(args.root)
119
+ graph.status = "ok"
120
+ try:
121
+ print(getattr(graph, tool.name)(**_kwargs(tool, args)))
122
+ return EXIT_CODES[graph.status]
123
+ finally:
124
+ graph.close()
125
+
126
+
127
+ if __name__ == "__main__":
128
+ sys.exit(main())
codegraph/community.py ADDED
@@ -0,0 +1,77 @@
1
+ """Community detection over the structural graph: groups of symbols that belong together.
2
+
3
+ Weighted label propagation, stdlib only. Leiden/Louvain give slightly better modularity but need a
4
+ dependency (igraph/networkx); for "which cluster is this symbol in" label propagation is enough, runs
5
+ in near-linear time and is deterministic (fixed visiting order, ties broken by smallest label).
6
+ """
7
+ from __future__ import annotations
8
+
9
+ from collections import defaultdict
10
+
11
+ CONF_WEIGHT = {"extracted": 1.0, "inferred": 0.6, "ambiguous": 0.25}
12
+ _FRESH = 1 << 40 # label space for nodes with no previous community
13
+ MIN_SIZE = 3 # smaller groups are folded into the neighbouring community they connect to most
14
+ MAX_SWEEPS = 12
15
+
16
+
17
+ def label_propagation(nodes: list[int], edges: list[tuple[int, int, float]],
18
+ hub_degree: int = 200, init: dict[int, int] | None = None) -> dict[int, int]:
19
+ """Return {node: community id}. Nodes with more than `hub_degree` neighbours (utility hubs such as
20
+ a logger) vote for nobody and are not voted for, so they cannot glue unrelated areas together."""
21
+ adj: dict[int, dict[int, float]] = defaultdict(lambda: defaultdict(float))
22
+ for a, b, w in edges:
23
+ if a != b:
24
+ adj[a][b] += w
25
+ adj[b][a] += w
26
+ hubs = {n for n, nb in adj.items() if len(nb) > hub_degree}
27
+ order = sorted(nodes, key=lambda n: (-len(adj.get(n, ())), n)) # dense areas settle first
28
+ if init: # warm start: keep last run's grouping, new nodes start alone (offset avoids clashes)
29
+ label = {n: init[n] if init.get(n, -1) >= 0 else _FRESH + n for n in nodes}
30
+ else:
31
+ label = {n: n for n in nodes}
32
+ for _ in range(MAX_SWEEPS):
33
+ moved = 0
34
+ for n in order:
35
+ if n in hubs or n not in adj:
36
+ continue
37
+ votes: dict[int, float] = defaultdict(float)
38
+ for m, w in adj[n].items():
39
+ if m not in hubs:
40
+ votes[label[m]] += w
41
+ if not votes:
42
+ continue
43
+ best = min(votes, key=lambda lb: (-votes[lb], lb))
44
+ if best != label[n] and votes[best] > votes.get(label[n], 0.0):
45
+ label[n] = best
46
+ moved += 1
47
+ if not moved:
48
+ break
49
+ # fold tiny groups into their strongest neighbour, then renumber 0..k-1 by size
50
+ size: dict[int, int] = defaultdict(int)
51
+ for lb in label.values():
52
+ size[lb] += 1
53
+ for n in order:
54
+ if size[label[n]] < MIN_SIZE and n in adj:
55
+ votes = defaultdict(float)
56
+ for m, w in adj[n].items():
57
+ if size[label[m]] >= MIN_SIZE:
58
+ votes[label[m]] += w
59
+ if votes:
60
+ size[label[n]] -= 1
61
+ label[n] = min(votes, key=lambda lb: (-votes[lb], lb))
62
+ size[label[n]] += 1
63
+ ranked = sorted({lb for lb in label.values()}, key=lambda lb: (-size[lb], lb))
64
+ renum = {lb: i for i, lb in enumerate(ranked)}
65
+ return {n: renum[lb] for n, lb in label.items()}
66
+
67
+
68
+ def detect(db) -> dict:
69
+ """Assign symbols.community from the edge table. Singletons (no edges) get -1."""
70
+ ids = [r[0] for r in db.execute("SELECT id FROM symbols")]
71
+ edges = [(s, d, CONF_WEIGHT.get(c, 0.6)) for s, d, c in db.execute("SELECT src, dst, conf FROM edges")]
72
+ prev = {r[0]: r[1] for r in db.execute("SELECT id, community FROM symbols")}
73
+ label = label_propagation(ids, edges, init=prev)
74
+ linked = {x for a, b, _ in edges for x in (a, b)}
75
+ db.executemany("UPDATE symbols SET community=? WHERE id=?",
76
+ [(label[i] if i in linked else -1, i) for i in ids])
77
+ return {"communities": len({label[i] for i in linked}), "linked_symbols": len(linked)}
codegraph/docs.py ADDED
@@ -0,0 +1,141 @@
1
+ """Non-code inputs: Markdown, SQL schemas and (with the `pdf` extra) PDFs.
2
+
3
+ Static like the code parser: no LLM, nothing leaves the machine. Each document becomes symbols the
4
+ rest of the graph already understands, so search, outline, PageRank and communities work on them:
5
+
6
+ - Markdown heading -> `section`, nested by heading level; `code spans` in its body become
7
+ `mentions` edges to the code symbols they name
8
+ - SQL CREATE ... -> `table` / `view` / `function` / `procedure`; FROM / JOIN / REFERENCES
9
+ inside the statement become `mentions` edges to other SQL objects
10
+ - PDF page -> `section` carrying the page text (line numbers are page numbers)
11
+ """
12
+ from __future__ import annotations
13
+
14
+ import importlib.util
15
+ import io
16
+ import re
17
+
18
+ from .parser import Parsed, Sym
19
+
20
+ DOC_EXT = {".md": "markdown", ".markdown": "markdown", ".sql": "sql"}
21
+ if importlib.util.find_spec("pypdf"):
22
+ DOC_EXT[".pdf"] = "pdf"
23
+ DOC_LANGS = frozenset(DOC_EXT.values())
24
+
25
+ HEADING = re.compile(r"^(#{1,6})\s+(.+?)\s*#*\s*$")
26
+ FENCE = re.compile(r"^\s*(```|~~~)")
27
+ CODE_SPAN = re.compile(r"`([^`\n]+)`")
28
+ DOTTED = re.compile(r"[A-Za-z_]\w*(?:(?:\.|::|->)[A-Za-z_]\w*)*")
29
+ MD_LINK = re.compile(r"\]\(([^)#\s]+)(?:#[^)]*)?\)")
30
+
31
+ SQL_CREATE = re.compile(
32
+ r"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:(?:GLOBAL\s+|LOCAL\s+)?(?:TEMP|TEMPORARY)\s+|MATERIALIZED\s+)?"
33
+ r"(TABLE|VIEW|FUNCTION|PROCEDURE)\s+(?:IF\s+NOT\s+EXISTS\s+)?([\w.\"`\[\]]+)", re.I)
34
+ SQL_REF = re.compile(r"\b(?:FROM|JOIN|REFERENCES|INTO|UPDATE)\s+([\w.\"`\[\]]+)", re.I)
35
+ SQL_COMMENT = re.compile(r"--[^\n]*|/\*.*?\*/", re.S)
36
+ SQL_KEYWORDS = frozenset("select set values lateral unnest generate_series".split())
37
+ MIN_MENTION = 3 # shorter names are noise (`i`, `db`) and would link to anything
38
+ PDF_TEXT = 2000 # characters of a page kept as its searchable doc
39
+
40
+
41
+ def _first_line(text: str, n: int = 200) -> str:
42
+ for line in text.splitlines():
43
+ line = line.strip()
44
+ if line:
45
+ return line[:n]
46
+ return ""
47
+
48
+
49
+ def _mentions(text: str) -> list[str]:
50
+ """Identifiers named in `code spans`. `Store.resolve()` yields Store and resolve."""
51
+ out: dict[str, None] = {}
52
+ for span in CODE_SPAN.findall(text):
53
+ for dotted in DOTTED.findall(span):
54
+ for part in re.split(r"\.|::|->", dotted):
55
+ if len(part) >= MIN_MENTION:
56
+ out[part] = None
57
+ return list(out)
58
+
59
+
60
+ def parse_markdown(src: bytes, stem: str) -> Parsed:
61
+ lines = src.decode("utf-8", errors="replace").splitlines()
62
+ heads: list[tuple[int, int, str]] = [] # (line, level, title)
63
+ in_fence = False
64
+ for n, line in enumerate(lines, 1):
65
+ if FENCE.match(line):
66
+ in_fence = not in_fence
67
+ elif not in_fence and (m := HEADING.match(line)):
68
+ heads.append((n, len(m.group(1)), m.group(2)))
69
+ p = Parsed()
70
+ stack: list[tuple[int, Sym]] = [] # (level, symbol) of the open ancestors
71
+ for k, (n, level, title) in enumerate(heads):
72
+ while stack and stack[-1][0] >= level:
73
+ stack.pop()
74
+ # a section runs until the next heading of the same or a higher level
75
+ end = next((m - 1 for m, lv, _ in heads[k + 1:] if lv <= level), len(lines))
76
+ parent = stack[-1][1] if stack else None
77
+ body = "\n".join(lines[n:end])
78
+ sym = Sym(id=k, name=title, qualname=f"{parent.qualname} > {title}" if parent else title,
79
+ kind="section", start=n, end=max(end, n), sig="#" * level + " " + title,
80
+ doc=_first_line(body), parent=parent.id if parent else None)
81
+ p.syms.append(sym)
82
+ stack.append((level, sym))
83
+ p.refs.extend((k, name, "mentions") for name in _mentions(body))
84
+ # relative links to other files: resolved against the file set like an import
85
+ p.imports.extend(dict.fromkeys(MD_LINK.findall("\n".join(lines))))
86
+ return p
87
+
88
+
89
+ def _strip_sql_comments(text: str) -> str:
90
+ """Blank comments but keep newlines, so line numbers still match the file."""
91
+ return SQL_COMMENT.sub(lambda m: re.sub(r"[^\n]", " ", m.group(0)), text)
92
+
93
+
94
+ def _sql_name(raw: str) -> str:
95
+ return re.sub(r"[\"`\[\]]", "", raw).rsplit(".", 1)[-1]
96
+
97
+
98
+ def parse_sql(src: bytes, stem: str) -> Parsed:
99
+ text = _strip_sql_comments(src.decode("utf-8", errors="replace"))
100
+ total = text.count("\n") + 1
101
+ found = list(SQL_CREATE.finditer(text))
102
+ p = Parsed()
103
+ for k, m in enumerate(found):
104
+ name = _sql_name(m.group(2))
105
+ start = text.count("\n", 0, m.start()) + 1
106
+ stop = found[k + 1].start() if k + 1 < len(found) else len(text)
107
+ # ponytail: a statement runs to the next CREATE; `;` inside $$ bodies makes `;` unusable.
108
+ # Trailing ALTER / INSERT statements attach to the object above them.
109
+ body = text[m.end():stop]
110
+ end = max(start, start + text.count("\n", m.start(), stop) - (1 if stop < len(text) else 0))
111
+ kind = m.group(1).lower()
112
+ p.syms.append(Sym(id=k, name=name, qualname=name, kind=kind, start=start, end=min(end, total),
113
+ sig=" ".join(m.group(0).split()), doc="", parent=None))
114
+ seen = {name}
115
+ for r in SQL_REF.finditer(body):
116
+ ref = _sql_name(r.group(1))
117
+ if ref.lower() not in SQL_KEYWORDS and ref not in seen and len(ref) >= 2:
118
+ seen.add(ref)
119
+ p.refs.append((k, ref, "mentions"))
120
+ return p
121
+
122
+
123
+ def parse_pdf(src: bytes, stem: str) -> Parsed:
124
+ from pypdf import PdfReader # optional extra, see DOC_EXT
125
+
126
+ p = Parsed()
127
+ for n, page in enumerate(PdfReader(io.BytesIO(src)).pages, 1):
128
+ text = " ".join((page.extract_text() or "").split())
129
+ if not text:
130
+ continue
131
+ title = f"{stem} p{n}: {text[:60]}"
132
+ p.syms.append(Sym(id=len(p.syms), name=title, qualname=title, kind="section", start=n, end=n,
133
+ sig=f"page {n}", doc=text[:PDF_TEXT], parent=None))
134
+ return p
135
+
136
+
137
+ PARSERS = {"markdown": parse_markdown, "sql": parse_sql, "pdf": parse_pdf}
138
+
139
+
140
+ def parse_doc(src: bytes, lang: str, stem: str) -> Parsed:
141
+ return PARSERS[lang](src, stem)
codegraph/embed.py ADDED
@@ -0,0 +1,221 @@
1
+ """Optional semantic layer: local embeddings of symbols, a content-hash vector cache and "similar" links.
2
+
3
+ Everything here is optional. Without numpy and an embedding provider (`pip install codegraph[embed]`),
4
+ `load_embedder()` returns None and the rest of codegraph behaves exactly as it does without this module.
5
+
6
+ No network at query time: the default provider runs a small local ONNX model through fastembed (the
7
+ model files are downloaded once, on first indexing). Providers sit behind the tiny `Embedder` protocol,
8
+ so another backend is a class with three members, selected with CODEGRAPH_EMBED=module:factory.
9
+ """
10
+ from __future__ import annotations
11
+
12
+ import hashlib
13
+ import importlib
14
+ import importlib.util
15
+ import os
16
+ from typing import Protocol
17
+
18
+ from .links import is_test_path
19
+ from .store import TARGET_KINDS_CALL, Store, split_words
20
+
21
+ DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"
22
+ TEXT_VERSION = "1" # bump when symbol_text() changes: every cached vector is then recomputed
23
+ EMBED_KINDS = tuple(sorted(TARGET_KINDS_CALL)) # variables / fields / modules are noise, not concepts
24
+ BATCH = 128
25
+ SIM_K = 5 # neighbours per symbol
26
+ SIM_THRESHOLD = 0.85 # cosine; tuned on this repo's own code (bge-small): below it, mostly generic noise
27
+ SIM_CHUNK = 256 # rows of the similarity matrix computed at once: bounds memory at chunk x N floats
28
+
29
+
30
+ class Embedder(Protocol):
31
+ """What codegraph needs from an embedding backend. Vectors must be L2-normalised float32."""
32
+
33
+ name: str # identifies model + settings; part of every cache key, so changing it re-embeds
34
+
35
+ def embed_passages(self, texts: list[str]): ... # -> ndarray (len(texts), dim)
36
+
37
+ def embed_query(self, text: str): ... # -> ndarray (dim,)
38
+
39
+
40
+ class FastEmbedder:
41
+ """fastembed / ONNX. The model is loaded on first use, so a no-change reindex never pays for it."""
42
+
43
+ def __init__(self, model: str = DEFAULT_MODEL):
44
+ self.name = f"fastembed:{model}"
45
+ self._model_name = model
46
+ self._model = None
47
+
48
+ def _load(self):
49
+ if self._model is None:
50
+ from fastembed import TextEmbedding
51
+ self._model = TextEmbedding(self._model_name)
52
+ return self._model
53
+
54
+ def embed_passages(self, texts):
55
+ import numpy as np
56
+ return np.asarray(list(self._load().passage_embed(texts)), dtype=np.float32)
57
+
58
+ def embed_query(self, text):
59
+ import numpy as np
60
+ return np.asarray(next(iter(self._load().query_embed([text]))), dtype=np.float32)
61
+
62
+
63
+ def load_embedder() -> Embedder | None:
64
+ """The configured provider, or None when semantic search is unavailable or switched off.
65
+
66
+ CODEGRAPH_EMBED: unset / "fastembed" (default), "off", or "package.module:factory".
67
+ CODEGRAPH_EMBED_MODEL picks another fastembed model.
68
+ """
69
+ spec = os.environ.get("CODEGRAPH_EMBED", "fastembed").strip()
70
+ if spec.lower() in ("off", "0", "none", "false", ""):
71
+ return None
72
+ if importlib.util.find_spec("numpy") is None:
73
+ return None
74
+ if ":" in spec:
75
+ mod, _, attr = spec.partition(":")
76
+ return getattr(importlib.import_module(mod), attr)()
77
+ if importlib.util.find_spec("fastembed") is None:
78
+ return None
79
+ return FastEmbedder(os.environ.get("CODEGRAPH_EMBED_MODEL", DEFAULT_MODEL))
80
+
81
+
82
+ def symbol_text(r) -> str:
83
+ """What gets embedded: the name as words, the qualified name, the signature and the doc line."""
84
+ parts = [r["kind"], split_words(r["name"]), r["qualname"], r["sig"] or "", r["doc"] or ""]
85
+ return ". ".join(p for p in parts if p)[:400]
86
+
87
+
88
+ def text_hash(embedder: Embedder, text: str) -> str:
89
+ return hashlib.sha1(f"{TEXT_VERSION}|{embedder.name}|{text}".encode()).hexdigest()
90
+
91
+
92
+ def _ph(n: int) -> str:
93
+ return ",".join("?" * n)
94
+
95
+
96
+ def embed_pending(store: Store, embedder: Embedder, progress=None) -> list[int]:
97
+ """Embed every symbol that has no vector yet; vectors already cached by content are reused.
98
+ Returns the ids of the symbols that were newly linked to a vector."""
99
+ import numpy as np
100
+ db = store.db
101
+ if store.get_meta("embed_model") != embedder.name: # other model: nothing cached is comparable
102
+ db.execute("DELETE FROM sym_vec")
103
+ db.execute("DELETE FROM links WHERE kind='similar'")
104
+ store.set_meta("embed_model", embedder.name)
105
+ rows = db.execute(
106
+ f"SELECT s.id, s.name, s.qualname, s.kind, s.sig, s.doc FROM symbols s "
107
+ f"LEFT JOIN sym_vec v ON v.sym_id = s.id WHERE v.sym_id IS NULL AND s.kind IN ({_ph(len(EMBED_KINDS))})",
108
+ EMBED_KINDS).fetchall()
109
+ if not rows:
110
+ return []
111
+ hashes = [text_hash(embedder, symbol_text(r)) for r in rows]
112
+ known: set[str] = set()
113
+ uniq = list(dict.fromkeys(hashes))
114
+ for i in range(0, len(uniq), 500):
115
+ chunk = uniq[i:i + 500]
116
+ known |= {r[0] for r in db.execute(f"SELECT hash FROM vec_cache WHERE hash IN ({_ph(len(chunk))})", chunk)}
117
+ todo = {h: symbol_text(r) for r, h in zip(rows, hashes, strict=True) if h not in known}
118
+ items = list(todo.items())
119
+ for i in range(0, len(items), BATCH):
120
+ batch = items[i:i + BATCH]
121
+ vecs = embedder.embed_passages([t for _, t in batch]).astype(np.float32)
122
+ db.executemany("INSERT OR REPLACE INTO vec_cache VALUES(?,?)",
123
+ [(h, vecs[j].tobytes()) for j, (h, _) in enumerate(batch)])
124
+ if progress and (i // BATCH) % 10 == 0:
125
+ progress(f"embedded {min(i + BATCH, len(items))}/{len(items)} symbols")
126
+ db.executemany("INSERT OR REPLACE INTO sym_vec VALUES(?,?)", [(r["id"], h) for r, h in zip(rows, hashes, strict=True)])
127
+ db.execute("DELETE FROM vec_cache WHERE hash NOT IN (SELECT hash FROM sym_vec)")
128
+ store.commit()
129
+ return [r["id"] for r in rows]
130
+
131
+
132
+ def load_matrix(db):
133
+ """(ids, unit-norm float32 matrix) of every embedded symbol, or (None, None) when there are none."""
134
+ import numpy as np
135
+ rows = db.execute("SELECT v.sym_id, c.vec FROM sym_vec v JOIN vec_cache c ON c.hash = v.hash "
136
+ "ORDER BY v.sym_id").fetchall()
137
+ if not rows:
138
+ return None, None
139
+ ids = np.fromiter((r[0] for r in rows), dtype=np.int64, count=len(rows))
140
+ mat = np.vstack([np.frombuffer(r[1], dtype=np.float32) for r in rows])
141
+ return ids, mat
142
+
143
+
144
+ def build_similar(store: Store, k: int = SIM_K, threshold: float = SIM_THRESHOLD,
145
+ only: set[int] | None = None) -> int:
146
+ """Top-k nearest neighbours above `threshold`, as `similar` links. Brute-force cosine, chunked.
147
+
148
+ `only` restricts the queries to those symbols (their links are replaced, everyone else's kept),
149
+ which is what an edit to one file needs; None rebuilds every link.
150
+
151
+ Skipped: pairs already tied together by a structural edge (a call or an inheritance says more than
152
+ "looks alike"), a symbol's own ancestors, same-named overloads, and test-to-test pairs (test names
153
+ all resemble each other, which tells you nothing).
154
+ """
155
+ import numpy as np
156
+ db = store.db
157
+ if only is None:
158
+ db.execute("DELETE FROM links WHERE kind='similar'")
159
+ ids, mat = load_matrix(db)
160
+ if ids is None or len(ids) < 2:
161
+ store.set_meta("similar_n", "0")
162
+ return 0
163
+ linked = {(a, b) for a, b in db.execute("SELECT src, dst FROM edges")}
164
+ linked |= {(b, a) for a, b in linked}
165
+ info = {r["id"]: (r["parent"], r["file"], r["name"]) for r in db.execute("SELECT id, parent, file, name FROM symbols")}
166
+
167
+ def ancestors(i: int) -> set[int]:
168
+ out, p = set(), info[i][0]
169
+ while p is not None and p not in out:
170
+ out.add(p)
171
+ p = info.get(p, (None,))[0]
172
+ return out
173
+
174
+ seen = {(a, b) if a < b else (b, a) for a, b in db.execute("SELECT src, dst FROM links WHERE kind='similar'")}
175
+ queries = np.arange(len(ids)) if only is None else np.flatnonzero(np.isin(ids, list(only)))
176
+ kk = min(k + 1, len(ids))
177
+ rows = []
178
+ for start in range(0, len(queries), SIM_CHUNK):
179
+ qi = queries[start:start + SIM_CHUNK]
180
+ sims = mat[qi] @ mat.T
181
+ sims[np.arange(len(qi)), qi] = -1.0 # not its own neighbour
182
+ top = np.argpartition(-sims, kk - 1, axis=1)[:, :kk]
183
+ for r, i in enumerate(qi):
184
+ a = int(ids[i])
185
+ for j in top[r]:
186
+ s = float(sims[r, j])
187
+ b = int(ids[j])
188
+ key = (a, b) if a < b else (b, a)
189
+ if s < threshold or key in seen or (a, b) in linked:
190
+ continue
191
+ if info[a][1:] == info[b][1:] or b in ancestors(a) or a in ancestors(b):
192
+ continue # overloads of one name in one file / a class and its own member
193
+ if is_test_path(info[a][1]) and is_test_path(info[b][1]):
194
+ continue
195
+ seen.add(key)
196
+ rows.append((key[0], key[1], "similar", s, ""))
197
+ db.executemany("INSERT INTO links VALUES(?,?,?,?,?)", rows)
198
+ store.set_meta("similar_n", str(len(ids)))
199
+ return len(rows)
200
+
201
+
202
+ FULL_REBUILD_FRACTION = 0.2 # past this much churn a full pass is cheaper than patching
203
+
204
+
205
+ def update(store: Store, progress=None) -> dict:
206
+ """Bring vectors and `similar` links up to date. Instant when nothing is pending."""
207
+ embedder = load_embedder()
208
+ if embedder is None:
209
+ return {"enabled": False}
210
+ new_ids = embed_pending(store, embedder, progress)
211
+ n = store.db.execute("SELECT COUNT(*) FROM sym_vec").fetchone()[0]
212
+ built = store.get_meta("similar_n")
213
+ if built is None or len(new_ids) > FULL_REBUILD_FRACTION * max(n, 1):
214
+ similar = build_similar(store)
215
+ elif new_ids:
216
+ similar = build_similar(store, only=set(new_ids))
217
+ else:
218
+ store.set_meta("similar_n", str(n)) # deletions only: their links went with the symbols
219
+ similar = -1 # -1: unchanged
220
+ store.commit()
221
+ return {"enabled": True, "model": embedder.name, "embedded": len(new_ids), "vectors": n, "similar": similar}