Compute and store embeddings for a list of nodes.
(self, nodes: list[GraphNode], batch_size: int = 64)
| 890 | self._conn.close() |
| 891 | |
| 892 | def embed_nodes(self, nodes: list[GraphNode], batch_size: int = 64) -> int: |
| 893 | """Compute and store embeddings for a list of nodes.""" |
| 894 | if not self.provider: |
| 895 | return 0 |
| 896 | |
| 897 | # Filter to nodes that need embedding |
| 898 | to_embed: list[tuple[GraphNode, str, str]] = [] |
| 899 | provider_name = self.provider.name |
| 900 | |
| 901 | for node in nodes: |
| 902 | if node.kind == "File": |
| 903 | continue |
| 904 | text = _node_to_text(node) |
| 905 | text_hash = hashlib.sha256(text.encode()).hexdigest() |
| 906 | |
| 907 | existing = self._conn.execute( |
| 908 | "SELECT text_hash, provider FROM embeddings WHERE qualified_name = ?", |
| 909 | (node.qualified_name,), |
| 910 | ).fetchone() |
| 911 | |
| 912 | # Re-embed if text changed OR provider changed |
| 913 | if (existing and existing["text_hash"] == text_hash |
| 914 | and existing["provider"] == provider_name): |
| 915 | continue |
| 916 | to_embed.append((node, text, text_hash)) |
| 917 | |
| 918 | if not to_embed: |
| 919 | return 0 |
| 920 | |
| 921 | # Encode in batches |
| 922 | texts = [t for _, t, _ in to_embed] |
| 923 | vectors = self.provider.embed(texts) |
| 924 | |
| 925 | for (node, _text, text_hash), vec in zip(to_embed, vectors): |
| 926 | blob = _encode_vector(vec) |
| 927 | self._conn.execute( |
| 928 | """INSERT OR REPLACE INTO embeddings (qualified_name, vector, text_hash, provider) |
| 929 | VALUES (?, ?, ?, ?)""", |
| 930 | (node.qualified_name, blob, text_hash, provider_name), |
| 931 | ) |
| 932 | |
| 933 | self._conn.commit() |
| 934 | return len(to_embed) |
| 935 | |
| 936 | def search(self, query: str, limit: int = 20) -> list[tuple[str, float]]: |
| 937 | """Search for nodes by semantic similarity.""" |