|
1
|
""" |
|
2
|
SemanticSearch — embedding-based similarity search over the navegador graph. |
|
3
|
|
|
4
|
Embeds function/class docstrings via an LLMProvider, stores the embedding |
|
5
|
vectors as JSON in a node property (``embedding``), and retrieves the top-k |
|
6
|
most similar nodes to a natural-language query using cosine similarity. |
|
7
|
|
|
8
|
Usage:: |
|
9
|
|
|
10
|
from navegador.graph import GraphStore |
|
11
|
from navegador.llm import get_provider |
|
12
|
from navegador.intelligence.search import SemanticSearch |
|
13
|
|
|
14
|
store = GraphStore.sqlite(".navegador/graph.db") |
|
15
|
provider = get_provider("openai") |
|
16
|
ss = SemanticSearch(store, provider) |
|
17
|
|
|
18
|
# Build / refresh the index (idempotent — re-embeds all nodes) |
|
19
|
ss.index() |
|
20
|
|
|
21
|
# Query |
|
22
|
results = ss.search("function that validates JWT tokens", limit=5) |
|
23
|
for r in results: |
|
24
|
print(r["name"], r["score"]) |
|
25
|
""" |
|
26
|
|
|
27
|
from __future__ import annotations |
|
28
|
|
|
29
|
import json |
|
30
|
import math |
|
31
|
from typing import TYPE_CHECKING, Any |
|
32
|
|
|
33
|
if TYPE_CHECKING: |
|
34
|
from navegador.graph.store import GraphStore |
|
35
|
from navegador.llm import LLMProvider |
|
36
|
|
|
37
|
|
|
38
|
# Cypher to fetch all embeddable nodes (those with a docstring or description) |
|
39
|
_EMBEDDABLE_NODES = """ |
|
40
|
MATCH (n) |
|
41
|
WHERE (n:Function OR n:Class OR n:Method OR n:Concept OR n:Rule OR n:Decision) |
|
42
|
AND (n.docstring IS NOT NULL OR n.description IS NOT NULL) |
|
43
|
RETURN |
|
44
|
labels(n)[0] AS type, |
|
45
|
n.name AS name, |
|
46
|
coalesce(n.file_path, '') AS file_path, |
|
47
|
coalesce(n.docstring, n.description, '') AS text |
|
48
|
LIMIT $limit |
|
49
|
""" |
|
50
|
|
|
51
|
# Cypher to fetch nodes that already have a stored embedding |
|
52
|
_NODES_WITH_EMBEDDINGS = """ |
|
53
|
MATCH (n) |
|
54
|
WHERE n.embedding IS NOT NULL |
|
55
|
RETURN |
|
56
|
labels(n)[0] AS type, |
|
57
|
n.name AS name, |
|
58
|
coalesce(n.file_path, '') AS file_path, |
|
59
|
coalesce(n.docstring, n.description, '') AS text, |
|
60
|
n.embedding AS embedding |
|
61
|
""" |
|
62
|
|
|
63
|
# Upsert the embedding property on a matched node |
|
64
|
_SET_EMBEDDING = """ |
|
65
|
MATCH (n) |
|
66
|
WHERE n.name = $name AND ($file_path = '' OR n.file_path = $file_path) |
|
67
|
SET n.embedding = $embedding |
|
68
|
""" |
|
69
|
|
|
70
|
|
|
71
|
class SemanticSearch: |
|
72
|
""" |
|
73
|
Embedding-based semantic search over the navegador graph. |
|
74
|
|
|
75
|
Embeddings are stored as a JSON string (serialised ``list[float]``) in the |
|
76
|
``embedding`` property of each node so they survive graph restarts without |
|
77
|
any external vector store. |
|
78
|
|
|
79
|
Args: |
|
80
|
store: A :class:`~navegador.graph.GraphStore` instance. |
|
81
|
provider: An :class:`~navegador.llm.LLMProvider` that implements |
|
82
|
:meth:`embed`. |
|
83
|
""" |
|
84
|
|
|
85
|
def __init__(self, store: "GraphStore", provider: "LLMProvider") -> None: |
|
86
|
self._store = store |
|
87
|
self._provider = provider |
|
88
|
|
|
89
|
# ── Index ───────────────────────────────────────────────────────────────── |
|
90
|
|
|
91
|
def index(self, limit: int = 1000) -> int: |
|
92
|
""" |
|
93
|
Embed all function/class/concept docstrings and store on the nodes. |
|
94
|
|
|
95
|
Args: |
|
96
|
limit: Maximum number of nodes to index in one pass. |
|
97
|
|
|
98
|
Returns: |
|
99
|
The number of nodes that were (re-)embedded. |
|
100
|
""" |
|
101
|
result = self._store.query(_EMBEDDABLE_NODES, {"limit": limit}) |
|
102
|
rows = result.result_set or [] |
|
103
|
indexed = 0 |
|
104
|
for row in rows: |
|
105
|
node_type, name, file_path, text = row[0], row[1], row[2], row[3] |
|
106
|
if not text: |
|
107
|
continue |
|
108
|
label = f"[{node_type}] {name}: {text}" |
|
109
|
vector = self._provider.embed(label) |
|
110
|
self._store.query( |
|
111
|
_SET_EMBEDDING, |
|
112
|
{ |
|
113
|
"name": name, |
|
114
|
"file_path": file_path, |
|
115
|
"embedding": json.dumps(vector), |
|
116
|
}, |
|
117
|
) |
|
118
|
indexed += 1 |
|
119
|
return indexed |
|
120
|
|
|
121
|
# ── Search ──────────────────────────────────────────────────────────────── |
|
122
|
|
|
123
|
def search(self, query: str, limit: int = 10) -> list[dict[str, Any]]: |
|
124
|
""" |
|
125
|
Embed *query* and return the *limit* most similar indexed nodes. |
|
126
|
|
|
127
|
Each result dict has keys: ``type``, ``name``, ``file_path``, |
|
128
|
``text``, ``score`` (cosine similarity, 0–1). |
|
129
|
|
|
130
|
Args: |
|
131
|
query: Natural-language search query. |
|
132
|
limit: Maximum number of results to return. |
|
133
|
|
|
134
|
Returns: |
|
135
|
List of result dicts sorted by descending similarity score. |
|
136
|
""" |
|
137
|
query_vec = self._provider.embed(query) |
|
138
|
|
|
139
|
result = self._store.query(_NODES_WITH_EMBEDDINGS, {}) |
|
140
|
rows = result.result_set or [] |
|
141
|
|
|
142
|
scored: list[dict[str, Any]] = [] |
|
143
|
for row in rows: |
|
144
|
node_type, name, file_path, text, emb_json = (row[0], row[1], row[2], row[3], row[4]) |
|
145
|
if not emb_json: |
|
146
|
continue |
|
147
|
try: |
|
148
|
node_vec: list[float] = json.loads(emb_json) |
|
149
|
except (json.JSONDecodeError, TypeError): |
|
150
|
continue |
|
151
|
score = self._cosine_similarity(query_vec, node_vec) |
|
152
|
scored.append( |
|
153
|
{ |
|
154
|
"type": node_type, |
|
155
|
"name": name, |
|
156
|
"file_path": file_path, |
|
157
|
"text": text, |
|
158
|
"score": score, |
|
159
|
} |
|
160
|
) |
|
161
|
|
|
162
|
scored.sort(key=lambda x: x["score"], reverse=True) |
|
163
|
return scored[:limit] |
|
164
|
|
|
165
|
# ── Internals ───────────────────────────────────────────────────────────── |
|
166
|
|
|
167
|
@staticmethod |
|
168
|
def _cosine_similarity(a: list[float], b: list[float]) -> float: |
|
169
|
""" |
|
170
|
Compute the cosine similarity between two vectors. |
|
171
|
|
|
172
|
Args: |
|
173
|
a: First embedding vector. |
|
174
|
b: Second embedding vector. |
|
175
|
|
|
176
|
Returns: |
|
177
|
Cosine similarity in the range ``[-1, 1]``. Returns ``0.0`` if |
|
178
|
either vector is the zero vector or the lengths differ. |
|
179
|
""" |
|
180
|
if len(a) != len(b): |
|
181
|
return 0.0 |
|
182
|
dot = sum(x * y for x, y in zip(a, b)) |
|
183
|
mag_a = math.sqrt(sum(x * x for x in a)) |
|
184
|
mag_b = math.sqrt(sum(x * x for x in b)) |
|
185
|
if mag_a == 0.0 or mag_b == 0.0: |
|
186
|
return 0.0 |
|
187
|
return dot / (mag_a * mag_b) |
|
188
|
|