| """SQLite-backed vector store for CodeRAG retrieval demos.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import re |
| import sqlite3 |
| from contextlib import closing |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Iterable |
|
|
| import numpy as np |
|
|
| from .records import FileIndex |
| from .synthetic_queries import build_synthetic_questions |
|
|
|
|
| DEFAULT_DIMENSIONS = 384 |
| SCHEMA_VERSION = "coderag-sqlite-vector-db-v1" |
| TOKEN_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|\d+(?:\.\d+)?") |
|
|
|
|
| @dataclass(frozen=True) |
| class SearchHit: |
| key: str |
| file_path: str |
| qualified_name: str |
| symbol_type: str |
| start_line: int |
| end_line: int |
| raw_code: str |
| own_sketch: str |
| propagated_text: str |
| source_scores: dict[str, float] |
| best_source: str |
| best_score: float |
| matched_text: str |
| matched_scope: str | None |
| matched_evidence: list[str] |
|
|
| def to_dict(self) -> dict: |
| return { |
| "key": self.key, |
| "file_path": self.file_path, |
| "qualified_name": self.qualified_name, |
| "symbol_type": self.symbol_type, |
| "start_line": self.start_line, |
| "end_line": self.end_line, |
| "raw_code": self.raw_code, |
| "own_sketch": self.own_sketch, |
| "propagated_text": self.propagated_text, |
| "source_scores": self.source_scores, |
| "best_source": self.best_source, |
| "best_score": self.best_score, |
| "matched_text": self.matched_text, |
| "matched_scope": self.matched_scope, |
| "matched_evidence": self.matched_evidence, |
| } |
|
|
|
|
| def _tokens(text: str) -> list[str]: |
| words = [match.group(0).lower() for match in TOKEN_RE.finditer(text)] |
| bigrams = [ |
| f"{left}_{right}" |
| for left, right in zip(words, words[1:]) |
| if left != right |
| ] |
| return [*words, *bigrams] |
|
|
|
|
| def embed_text(text: str, dimensions: int = DEFAULT_DIMENSIONS) -> np.ndarray: |
| """Create a deterministic normalized hashing vector for local search.""" |
| vector = np.zeros(dimensions, dtype=np.float32) |
| for token in _tokens(text): |
| digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest() |
| bucket = int.from_bytes(digest[:4], "little") % dimensions |
| sign = 1.0 if digest[4] & 1 else -1.0 |
| vector[bucket] += sign |
| norm = float(np.linalg.norm(vector)) |
| if norm > 0: |
| vector /= norm |
| return vector |
|
|
|
|
| def _blob(vector: np.ndarray) -> bytes: |
| return np.asarray(vector, dtype=np.float32).tobytes() |
|
|
|
|
| def _vector(blob: bytes) -> np.ndarray: |
| return np.frombuffer(blob, dtype=np.float32) |
|
|
|
|
| def _record_key(record: dict, root_path: str | None = None) -> str: |
| path = Path(record["file_path"]) |
| if root_path: |
| try: |
| file_path = path.resolve().relative_to(Path(root_path).resolve()).as_posix() |
| except ValueError: |
| file_path = path.as_posix() |
| else: |
| file_path = path.as_posix() |
| return f"{file_path}::{record['qualified_name']}" |
|
|
|
|
| class CodeRagVectorDB: |
| """Small vector database used by the CodeRAG React demo.""" |
|
|
| def __init__(self, path: str | Path, *, dimensions: int = DEFAULT_DIMENSIONS): |
| self.path = Path(path) |
| self.dimensions = dimensions |
|
|
| def connect(self) -> sqlite3.Connection: |
| self.path.parent.mkdir(parents=True, exist_ok=True) |
| connection = sqlite3.connect(self.path) |
| connection.row_factory = sqlite3.Row |
| return connection |
|
|
| def initialize(self) -> None: |
| with closing(self.connect()) as connection, connection: |
| connection.executescript( |
| """ |
| CREATE TABLE IF NOT EXISTS metadata ( |
| key TEXT PRIMARY KEY, |
| value TEXT NOT NULL |
| ); |
| |
| CREATE TABLE IF NOT EXISTS functions ( |
| key TEXT PRIMARY KEY, |
| file_path TEXT NOT NULL, |
| qualified_name TEXT NOT NULL, |
| symbol_type TEXT NOT NULL, |
| start_line INTEGER NOT NULL DEFAULT 0, |
| end_line INTEGER NOT NULL DEFAULT 0, |
| raw_code TEXT NOT NULL DEFAULT '', |
| own_sketch TEXT NOT NULL DEFAULT '', |
| propagated_text TEXT NOT NULL DEFAULT '' |
| ); |
| |
| CREATE TABLE IF NOT EXISTS vectors ( |
| id INTEGER PRIMARY KEY AUTOINCREMENT, |
| function_key TEXT NOT NULL, |
| source TEXT NOT NULL, |
| scope TEXT, |
| text TEXT NOT NULL, |
| evidence_json TEXT NOT NULL DEFAULT '[]', |
| vector BLOB NOT NULL, |
| FOREIGN KEY(function_key) REFERENCES functions(key) |
| ); |
| |
| CREATE INDEX IF NOT EXISTS idx_vectors_source |
| ON vectors(source); |
| CREATE INDEX IF NOT EXISTS idx_vectors_function |
| ON vectors(function_key); |
| """ |
| ) |
| connection.execute( |
| """ |
| INSERT OR REPLACE INTO metadata(key, value) |
| VALUES ('schema_version', ?), ('dimensions', ?) |
| """, |
| (SCHEMA_VERSION, str(self.dimensions)), |
| ) |
|
|
| def clear(self) -> None: |
| with closing(self.connect()) as connection, connection: |
| connection.execute("DELETE FROM vectors") |
| connection.execute("DELETE FROM functions") |
|
|
| def build( |
| self, |
| *, |
| repository_index_path: str | Path, |
| synthetic_questions_path: str | Path | None = None, |
| ) -> dict: |
| self.initialize() |
| self.clear() |
| repository_index = json.loads( |
| Path(repository_index_path).read_text(encoding="utf-8") |
| ) |
| root_path = repository_index.get("root_path") |
| function_count = 0 |
| vector_count = 0 |
| with closing(self.connect()) as connection, connection: |
| for record in repository_index["functions"]: |
| key = _record_key(record, root_path) |
| self._upsert_function(connection, key, record) |
| for source, field in ( |
| ("raw_code", "raw_code"), |
| ("cross_file_propagated", "propagated_text"), |
| ): |
| text = record.get(field, "") |
| if text.strip(): |
| self._insert_vector( |
| connection, |
| function_key=key, |
| source=source, |
| text=text, |
| scope=None, |
| evidence=[], |
| ) |
| vector_count += 1 |
| function_count += 1 |
|
|
| synthetic_count = 0 |
| if synthetic_questions_path is not None: |
| synthetic_path = Path(synthetic_questions_path) |
| if synthetic_path.exists(): |
| synthetic = json.loads( |
| synthetic_path.read_text(encoding="utf-8") |
| ) |
| for payload in synthetic["repositories"].values(): |
| for key, questions in payload["functions"].items(): |
| self._upsert_placeholder(connection, key) |
| for question in questions: |
| self._insert_vector( |
| connection, |
| function_key=key, |
| source="semantic_synthetic", |
| text=question["text"], |
| scope=question.get("scope"), |
| evidence=question.get("evidence", []), |
| ) |
| vector_count += 1 |
| synthetic_count += 1 |
|
|
| connection.execute( |
| """ |
| INSERT OR REPLACE INTO metadata(key, value) |
| VALUES |
| ('repository_index_path', ?), |
| ('synthetic_questions_path', ?), |
| ('function_count', ?), |
| ('vector_count', ?), |
| ('synthetic_question_count', ?) |
| """, |
| ( |
| str(repository_index_path), |
| str(synthetic_questions_path or ""), |
| str(function_count), |
| str(vector_count), |
| str(synthetic_count), |
| ), |
| ) |
| return self.stats() |
|
|
| def build_from_file_index(self, index: FileIndex) -> dict: |
| """Replace the vector database with records from one pasted-code index.""" |
| self.initialize() |
| self.clear() |
| records_by_key = { |
| f"{Path(record.file_path).as_posix()}::{record.qualified_name}": record |
| for record in index.functions |
| } |
| graph = {} |
| for key, record in records_by_key.items(): |
| graph[key] = [ |
| f"{Path(record.file_path).as_posix()}::{callee}" |
| for callee in record.direct_callees |
| if f"{Path(record.file_path).as_posix()}::{callee}" in records_by_key |
| ] |
| synthetic = build_synthetic_questions(records_by_key, graph) |
| vector_count = 0 |
| with closing(self.connect()) as connection, connection: |
| for key, record in records_by_key.items(): |
| self._upsert_function(connection, key, record.to_dict()) |
| for source, text in ( |
| ("raw_code", record.raw_code), |
| ("cross_file_propagated", record.propagated_text), |
| ): |
| if text.strip(): |
| self._insert_vector( |
| connection, |
| function_key=key, |
| source=source, |
| text=text, |
| scope=None, |
| evidence=[], |
| ) |
| vector_count += 1 |
| for question in synthetic.get(key, []): |
| self._insert_vector( |
| connection, |
| function_key=key, |
| source="semantic_synthetic", |
| text=question.text, |
| scope=question.scope, |
| evidence=list(question.evidence), |
| ) |
| vector_count += 1 |
|
|
| connection.execute( |
| """ |
| INSERT OR REPLACE INTO metadata(key, value) |
| VALUES |
| ('repository_index_path', ?), |
| ('synthetic_questions_path', ?), |
| ('function_count', ?), |
| ('vector_count', ?), |
| ('synthetic_question_count', ?), |
| ('active_index_mode', ?), |
| ('active_file_path', ?) |
| """, |
| ( |
| str(index.file_path), |
| "generated-from-pasted-code", |
| str(len(records_by_key)), |
| str(vector_count), |
| str(sum(len(items) for items in synthetic.values())), |
| "pasted_code", |
| str(index.file_path), |
| ), |
| ) |
| stats = self.stats() |
| stats["indexing_order"] = index.indexing_order |
| stats["call_edge_count"] = len(index.call_edges) |
| return stats |
|
|
| def stats(self) -> dict: |
| self.initialize() |
| with closing(self.connect()) as connection: |
| metadata = { |
| row["key"]: row["value"] |
| for row in connection.execute("SELECT key, value FROM metadata") |
| } |
| sources = { |
| row["source"]: row["count"] |
| for row in connection.execute( |
| """ |
| SELECT source, COUNT(*) AS count |
| FROM vectors |
| GROUP BY source |
| ORDER BY source |
| """ |
| ) |
| } |
| function_count = connection.execute( |
| "SELECT COUNT(*) AS count FROM functions" |
| ).fetchone()["count"] |
| vector_count = connection.execute( |
| "SELECT COUNT(*) AS count FROM vectors" |
| ).fetchone()["count"] |
| return { |
| "schema_version": metadata.get("schema_version", SCHEMA_VERSION), |
| "database_path": str(self.path), |
| "dimensions": int(metadata.get("dimensions", self.dimensions)), |
| "function_count": function_count, |
| "vector_count": vector_count, |
| "sources": sources, |
| "repository_index_path": metadata.get("repository_index_path", ""), |
| "synthetic_questions_path": metadata.get( |
| "synthetic_questions_path", |
| "", |
| ), |
| "active_index_mode": metadata.get("active_index_mode", "repository"), |
| "active_file_path": metadata.get("active_file_path", ""), |
| } |
|
|
| def search( |
| self, |
| query: str, |
| *, |
| top_k: int = 8, |
| per_source: int = 5, |
| sources: Iterable[str] | None = None, |
| ) -> dict: |
| self.initialize() |
| requested_sources = list(sources or ( |
| "raw_code", |
| "cross_file_propagated", |
| "semantic_synthetic", |
| )) |
| query_vector = embed_text(query, self.dimensions) |
| if not np.any(query_vector): |
| return { |
| "query": query, |
| "hits": [], |
| "candidate_count": 0, |
| "per_source": per_source, |
| "sources": requested_sources, |
| "stats": self.stats(), |
| } |
| with closing(self.connect()) as connection: |
| rows = connection.execute( |
| """ |
| SELECT |
| vectors.id, |
| vectors.function_key, |
| vectors.source, |
| vectors.scope, |
| vectors.text, |
| vectors.evidence_json, |
| vectors.vector, |
| functions.file_path, |
| functions.qualified_name, |
| functions.symbol_type, |
| functions.start_line, |
| functions.end_line, |
| functions.raw_code, |
| functions.own_sketch, |
| functions.propagated_text |
| FROM vectors |
| JOIN functions ON functions.key = vectors.function_key |
| WHERE vectors.source IN ({}) |
| """.format(",".join("?" for _ in requested_sources)), |
| requested_sources, |
| ).fetchall() |
|
|
| ranked_by_source: dict[str, list[tuple[float, sqlite3.Row]]] = { |
| source: [] for source in requested_sources |
| } |
| for row in rows: |
| score = float(np.dot(_vector(row["vector"]), query_vector)) |
| ranked_by_source[row["source"]].append((score, row)) |
|
|
| candidate_rows: dict[str, dict] = {} |
| for source, scored_rows in ranked_by_source.items(): |
| for score, row in sorted( |
| scored_rows, |
| key=lambda item: item[0], |
| reverse=True, |
| )[:per_source]: |
| key = row["function_key"] |
| candidate = candidate_rows.setdefault( |
| key, |
| { |
| "row": row, |
| "source_scores": {}, |
| "best_score": score, |
| "best_source": source, |
| "matched_text": row["text"], |
| "matched_scope": row["scope"], |
| "matched_evidence": json.loads(row["evidence_json"]), |
| }, |
| ) |
| candidate["source_scores"][source] = round(score, 6) |
| if score > candidate["best_score"]: |
| candidate["best_score"] = score |
| candidate["best_source"] = source |
| candidate["matched_text"] = row["text"] |
| candidate["matched_scope"] = row["scope"] |
| candidate["matched_evidence"] = json.loads( |
| row["evidence_json"] |
| ) |
|
|
| hits = [] |
| for key, candidate in candidate_rows.items(): |
| row = candidate["row"] |
| hits.append(SearchHit( |
| key=key, |
| file_path=row["file_path"], |
| qualified_name=row["qualified_name"], |
| symbol_type=row["symbol_type"], |
| start_line=row["start_line"], |
| end_line=row["end_line"], |
| raw_code=row["raw_code"], |
| own_sketch=row["own_sketch"], |
| propagated_text=row["propagated_text"], |
| source_scores=candidate["source_scores"], |
| best_source=candidate["best_source"], |
| best_score=round(candidate["best_score"], 6), |
| matched_text=candidate["matched_text"], |
| matched_scope=candidate["matched_scope"], |
| matched_evidence=candidate["matched_evidence"], |
| )) |
| hits.sort(key=lambda hit: hit.best_score, reverse=True) |
| return { |
| "query": query, |
| "hits": [hit.to_dict() for hit in hits[:top_k]], |
| "candidate_count": len(candidate_rows), |
| "per_source": per_source, |
| "sources": requested_sources, |
| "stats": self.stats(), |
| } |
|
|
| def _upsert_function( |
| self, |
| connection: sqlite3.Connection, |
| key: str, |
| record: dict, |
| ) -> None: |
| connection.execute( |
| """ |
| INSERT OR REPLACE INTO functions( |
| key, |
| file_path, |
| qualified_name, |
| symbol_type, |
| start_line, |
| end_line, |
| raw_code, |
| own_sketch, |
| propagated_text |
| ) |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) |
| """, |
| ( |
| key, |
| record.get("file_path", key.split("::", 1)[0]), |
| record.get("qualified_name", key.split("::", 1)[-1]), |
| record.get("symbol_type", "function"), |
| int(record.get("start_line", 0) or 0), |
| int(record.get("end_line", 0) or 0), |
| record.get("raw_code", ""), |
| record.get("own_sketch", ""), |
| record.get("propagated_text", ""), |
| ), |
| ) |
|
|
| def _upsert_placeholder( |
| self, |
| connection: sqlite3.Connection, |
| key: str, |
| ) -> None: |
| path, qualified = key.split("::", 1) |
| connection.execute( |
| """ |
| INSERT OR IGNORE INTO functions( |
| key, |
| file_path, |
| qualified_name, |
| symbol_type |
| ) |
| VALUES (?, ?, ?, ?) |
| """, |
| (key, path, qualified, "function"), |
| ) |
|
|
| def _insert_vector( |
| self, |
| connection: sqlite3.Connection, |
| *, |
| function_key: str, |
| source: str, |
| text: str, |
| scope: str | None, |
| evidence: list[str], |
| ) -> None: |
| connection.execute( |
| """ |
| INSERT INTO vectors( |
| function_key, |
| source, |
| scope, |
| text, |
| evidence_json, |
| vector |
| ) |
| VALUES (?, ?, ?, ?, ?, ?) |
| """, |
| ( |
| function_key, |
| source, |
| scope, |
| text, |
| json.dumps(evidence), |
| _blob(embed_text(text, self.dimensions)), |
| ), |
| ) |
|
|
